Backport PR #15805: Fix codemirror highlight for Python builtin (#15882)

Co-authored-by: Allan Chain <36528777+AllanChain@users.noreply.github.com>

packages/codemirror-extension/src/services.tsx

@@ -3,7 +3,7 @@
  * Distributed under the terms of the Modified BSD License.
  */
 
-import { StreamLanguage } from '@codemirror/language';
+import { LanguageSupport, StreamLanguage } from '@codemirror/language';
 import { IYText } from '@jupyter/ydoc';
 import {
   JupyterFrontEnd,
@@ -20,6 +20,7 @@ import {
   IEditorLanguageRegistry,
   IEditorThemeRegistry,
   parseMathIPython,
+  pythonBuiltin,
   ybinding
 } from '@jupyterlab/codemirror';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
@@ -62,17 +63,22 @@ export const languagePlugin: JupyterFrontEndPlugin<IEditorLanguageRegistry> = {
       name: 'ipythongfm',
       mime: 'text/x-ipythongfm',
       load: async () => {
-        const [m, tex] = await Promise.all([
+        const [m, py, tex] = await Promise.all([
           import('@codemirror/lang-markdown'),
+          import('@codemirror/lang-python'),
           import('@codemirror/legacy-modes/mode/stex')
         ]);
-        return m.markdown({
+        const mdlang = m.markdown({
           base: m.markdownLanguage,
           codeLanguages: (info: string) => languages.findBest(info) as any,
           extensions: [
             parseMathIPython(StreamLanguage.define(tex.stexMath).parser)
           ]
         });
+        return new LanguageSupport(mdlang.language, [
+          mdlang.support,
+          pythonBuiltin(py.pythonLanguage)
+        ]);
       }
     });
     return languages;

packages/codemirror/src/index.ts

@@ -15,3 +15,4 @@ export * from './mimetype';
 export * from './searchprovider';
 export * from './theme';
 export * from './token';
+export * from './pythonBuiltin';

packages/codemirror/src/language.ts

@@ -16,6 +16,7 @@ import { highlightTree } from '@lezer/highlight';
 
 import { jupyterHighlightStyle } from './theme';
 import { IEditorLanguage, IEditorLanguageRegistry } from './token';
+import { pythonBuiltin } from './pythonBuiltin';
 
 /**
  * CodeMirror language registry
@@ -473,7 +474,10 @@ export namespace EditorLanguageRegistry {
         filename: /^(BUCK|BUILD)$/,
         async load() {
           const m = await import('@codemirror/lang-python');
-          return m.python();
+          return new LanguageSupport(
+            m.pythonLanguage,
+            pythonBuiltin(m.pythonLanguage)
+          );
         }
       },
       {
@@ -485,7 +489,10 @@ export namespace EditorLanguageRegistry {
           // to activate feature such as code folding.
           // return Promise.resolve(legacy(mkPython({ singleOperators: /\?/ })));
           const m = await import('@codemirror/lang-python');
-          return m.python();
+          return new LanguageSupport(
+            m.pythonLanguage,
+            pythonBuiltin(m.pythonLanguage)
+          );
         }
       },
       {

packages/codemirror/src/pythonBuiltin.ts

@@ -0,0 +1,159 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+import { Language, syntaxTree } from '@codemirror/language';
+import { RangeSetBuilder } from '@codemirror/state';
+import {
+  Decoration,
+  DecorationSet,
+  EditorView,
+  ViewPlugin,
+  ViewUpdate
+} from '@codemirror/view';
+import { NodeProp, SyntaxNodeRef, Tree } from '@lezer/common';
+
+export class PythonBuiltin {
+  decorations: DecorationSet;
+  decoratedTo: number;
+  langPython: Language;
+  tree: Tree;
+  mark: Decoration;
+
+  constructor(view: EditorView, langPython: Language) {
+    this.langPython = langPython;
+    this.tree = syntaxTree(view.state);
+    this.mark = Decoration.mark({ class: 'cm-builtin' });
+    this.decorations = this.buildDeco(view);
+    this.decoratedTo = view.viewport.to;
+  }
+
+  update(update: ViewUpdate) {
+    let tree = syntaxTree(update.state);
+    let { viewport } = update.view,
+      decoratedToMapped = update.changes.mapPos(this.decoratedTo, 1);
+    if (
+      tree.length < viewport.to &&
+      tree.type == this.tree.type &&
+      decoratedToMapped >= viewport.to
+    ) {
+      this.decorations = this.decorations.map(update.changes);
+      this.decoratedTo = decoratedToMapped;
+    } else if (tree != this.tree || update.viewportChanged) {
+      this.tree = tree;
+      this.decorations = this.buildDeco(update.view);
+      this.decoratedTo = viewport.to;
+    }
+  }
+
+  buildDeco(view: EditorView) {
+    if (!this.tree.length) return Decoration.none;
+
+    let builder = new RangeSetBuilder<Decoration>();
+    const enter = (node: SyntaxNodeRef) => {
+      const cursor = node.node.cursor();
+      // Handle nested language, e.g. Markdown
+      const mounted = cursor.tree && cursor.tree.prop(NodeProp.mounted);
+      if (mounted && mounted.overlay) {
+        node.node
+          .enter(mounted.overlay[0].from + node.from, 1)
+          ?.cursor()
+          .iterate(enter);
+      }
+      if (
+        this.langPython.isActiveAt(view.state, node.from + 1) &&
+        node.name === 'VariableName'
+      ) {
+        const variableName = view.state.sliceDoc(node.from, node.to);
+        if (builtins.includes(variableName)) {
+          builder.add(node.from, node.to, this.mark);
+        }
+      }
+    };
+    for (let { from, to } of view.visibleRanges) {
+      this.tree.iterate({ enter, from, to });
+    }
+    return builder.finish();
+  }
+}
+
+export function pythonBuiltin(langPython: Language) {
+  return ViewPlugin.define(view => new PythonBuiltin(view, langPython), {
+    decorations: v => v.decorations
+  });
+}
+
+const builtins = [
+  'abs',
+  'aiter',
+  'all',
+  'any',
+  'anext',
+  'ascii',
+  'bin',
+  'bool',
+  'breakpoint',
+  'bytearray',
+  'bytes',
+  'callable',
+  'chr',
+  'classmethod',
+  'compile',
+  'complex',
+  'delattr',
+  'dict',
+  'dir',
+  'divmod',
+  'enumerate',
+  'eval',
+  'exec',
+  'filter',
+  'float',
+  'format',
+  'frozenset',
+  'getattr',
+  'globals',
+  'hasattr',
+  'hash',
+  'help',
+  'hex',
+  'id',
+  'input',
+  'int',
+  'isinstance',
+  'issubclass',
+  'iter',
+  'len',
+  'list',
+  'locals',
+  'map',
+  'max',
+  'memoryview',
+  'min',
+  'next',
+  'object',
+  'oct',
+  'open',
+  'ord',
+  'pow',
+  'print',
+  'property',
+  'range',
+  'repr',
+  'reversed',
+  'round',
+  'set',
+  'setattr',
+  'slice',
+  'sorted',
+  'staticmethod',
+  'str',
+  'sum',
+  'super',
+  'tuple',
+  'type',
+  'vars',
+  'zip',
+  '__import__'
+];

packages/codemirror/src/theme.ts

@@ -81,6 +81,10 @@ export const jupyterEditorTheme = EditorView.theme({
 
   '.cm-tooltip': {
     backgroundColor: 'var(--jp-layout-color1)'
+  },
+
+  '.cm-builtin': {
+    color: 'var(--jp-mirror-editor-builtin-color)'
   }
 });