From 3ae9425d75719fc3e1216d38bc0467f67dfd3000 Mon Sep 17 00:00:00 2001
From: Mariusz Jurowicz <mariusz.jurowicz@xsolve.pl>
Date: Thu, 5 Jul 2018 14:10:12 +0200
Subject: [PATCH] #5742 define line magic highlighter mode

---
 js/notebook/src/extension.js                  |  3 +-
 js/notebook/src/extension/codeMirror.ts       | 82 ++++++++++++++++---
 .../src/extension/groovyModeExtension.ts      |  7 +-
 js/notebook/src/types/global.env.d.ts         |  1 +
 4 files changed, 78 insertions(+), 15 deletions(-)

diff --git a/js/notebook/src/extension.js b/js/notebook/src/extension.js
index c2c1ddd41a..73beac7e20 100644
--- a/js/notebook/src/extension.js
+++ b/js/notebook/src/extension.js
@@ -113,7 +113,8 @@ define([
     if (inNotebook) {
       // setup things to run on loading config/notebook
 
-      bxCodeMirror.extendHighlightModes(Jupyter, CodeMirror);
+      CodeMirror.defineInitHook(bxCodeMirror.addLineMagicsOverlay);
+      bxCodeMirror.extendHighlightModes(Jupyter);
       GroovyMode.extendWithLineComment(Jupyter, CodeMirror);
 
       Jupyter.notebook.config.loaded
diff --git a/js/notebook/src/extension/codeMirror.ts b/js/notebook/src/extension/codeMirror.ts
index ba0fbfa025..2038c3cee8 100644
--- a/js/notebook/src/extension/codeMirror.ts
+++ b/js/notebook/src/extension/codeMirror.ts
@@ -14,20 +14,78 @@
  *  limitations under the License.
  */
 
-export function extendHighlightModes(Jupyter: any, CodeMirror: any) {
-  Jupyter.CodeCell.options_default.highlight_modes = _.extend(
-    Jupyter.CodeCell.options_default.highlight_modes,
-    {
-      magic_groovy: { reg: ['^%%groovy'] },
-      magic_java: { reg: ['^%%java'] },
-      magic_scala: { reg: ['^%%scala'] },
-      magic_kotlin: { reg: ['^%%kotlin'] },
-      magic_clojure: { reg: ['^%%clojure'] },
-      magic_sql: { reg: ['^%%sql'] }
+/// <reference path='../types/index.d.ts'/>
+
+export function extendHighlightModes(Jupyter: any) {
+  Jupyter.CodeCell.options_default.highlight_modes = {
+    ...Jupyter.CodeCell.options_default.highlight_modes,
+    magic_python: {reg: ['^%%python']},
+    magic_groovy: {reg: ['^%%groovy']},
+    magic_java: {reg: ['^%%java']},
+    magic_scala: {reg: ['^%%scala']},
+    magic_kotlin: {reg: ['^%%kotlin']},
+    magic_clojure: {reg: ['^%%clojure']},
+    magic_sql: {reg: ['^%%sql']},
+    magic_html: {reg: ['^%%html']}
+  };
+}
+
+const lineMagicOverlay = {
+  startState() {
+    return { firstMatched: false, inMagicLine: false };
+  },
+
+  token(stream, state) {
+    if (!state.firstMatched) {
+      state.firstMatched = true;
+
+      if (stream.match("%", false)) {
+        state.inMagicLine = true;
+      }
+    }
+
+    if (state.inMagicLine) {
+      stream.eat(() => {
+        return true;
+      });
+
+      if (stream.eol()) {
+        state.inMagicLine = false;
+      }
+
+      return "line_magic";
     }
-  );
+
+    stream.skipToEnd();
+
+    return null;
+  }
+};
+
+export function autoHighlightLineMagics(code_mirror) {
+  const current_mode = code_mirror.getOption('mode');
+  const first_line = code_mirror.getLine(0);
+  const re = /^%\w+/;
+
+  if (first_line.match(re) !== null) {
+    // Add an overlay mode to recognize the first line as "magic" instead
+    // of the mode used for the rest of the cell.
+    CodeMirror.defineMode('line_magic', (config) => {
+      return CodeMirror.overlayMode(CodeMirror.getMode(config, current_mode), lineMagicOverlay);
+    });
+
+    code_mirror.setOption('mode', 'line_magic');
+  }
+}
+
+export function addLineMagicsOverlay(code_mirror) {
+  autoHighlightLineMagics(code_mirror);
+  code_mirror.on("focus", () => autoHighlightLineMagics(code_mirror));
+  code_mirror.on("change", () => autoHighlightLineMagics(code_mirror));
+  code_mirror.on("blur", () => autoHighlightLineMagics(code_mirror));
 }
 
 export default {
-  extendHighlightModes
+  extendHighlightModes,
+  addLineMagicsOverlay
 };
diff --git a/js/notebook/src/extension/groovyModeExtension.ts b/js/notebook/src/extension/groovyModeExtension.ts
index 1ce6d02fa3..c7288331cc 100644
--- a/js/notebook/src/extension/groovyModeExtension.ts
+++ b/js/notebook/src/extension/groovyModeExtension.ts
@@ -14,10 +14,12 @@
  *  limitations under the License.
  */
 
+import {addLineMagicsOverlay} from "./codeMirror";
+
 export namespace GroovyMode {
   export const LINE_COMMENT_CHAR = '//';
 
-  export function setCodeMirrorLineComment(cell: any) {
+  export function setCodeMirrorLineComment(cell: any, CodeMirror) {
     if (cell.cell_type !== 'code') {
       return;
     }
@@ -32,11 +34,12 @@ export namespace GroovyMode {
     }
 
     cell.auto_highlight();
+    addLineMagicsOverlay(cm);
   }
 
   export function extendWithLineComment(Jupyter: any, CodeMirror: any) {
     CodeMirror.extendMode('groovy', { lineComment: LINE_COMMENT_CHAR });
 
-    Jupyter.notebook.get_cells().map(setCodeMirrorLineComment);
+    Jupyter.notebook.get_cells().map((cell) => setCodeMirrorLineComment(cell, CodeMirror));
   }
 }
diff --git a/js/notebook/src/types/global.env.d.ts b/js/notebook/src/types/global.env.d.ts
index a924f45f39..5891d28e57 100644
--- a/js/notebook/src/types/global.env.d.ts
+++ b/js/notebook/src/types/global.env.d.ts
@@ -31,6 +31,7 @@ interface MapConstructor {
 }
 
 declare var Map: MapConstructor;
+declare var CodeMirror: any;
 
 declare interface NumberConstructor {
   isNaN: (number: number) => boolean,