Refactor transformations and add name transformation (#56)

* Add name field transformation

* Add AST debug logging

* Add basic name transformation

* Refactor matchers to add function name transform

* Add missing function name

* Clean up transformations refactor

* Refactor branch debug logging

* Fix purple color setting name

* Add function and class name scopes

* Fixed name matching for exported/decorated nodes

* Clean up transformations refactor (again)

Co-authored-by: Pokey Rule <[email protected]>

Author: brxck

src/Types.ts

@@ -47,12 +47,15 @@ export type ScopeType =
   | "argumentOrParameter"
   | "arrowFunction"
   | "class"
+  | "className"
   | "comment"
   | "dictionary"
   | "functionCall"
+  | "functionName"
   | "ifStatement"
   | "list"
   | "listElement"
+  | "name"
   | "namedFunction"
   | "pair"
   | "pairKey"
@@ -244,13 +247,25 @@ export interface Graph {
   readonly editStyles: EditStyles;
 }
 
+export interface DecorationColorSetting {
+  dark: string;
+  light: string;
+  highContrast: string;
+}
+
 export type NodeMatcher = (
   editor: vscode.TextEditor,
   node: SyntaxNode
 ) => SelectionWithContext | null;
 
-export interface DecorationColorSetting {
-  dark: string;
-  light: string;
-  highContrast: string;
-}
+/**
+ * Returns the desired relative of the provided node.
+ * Returns null if matching node not found.
+ **/
+export type NodeFinder = (node: SyntaxNode) => SyntaxNode | null;
+
+/** Returns a selection for a given SyntaxNode */
+export type SelectionExtractor = (
+  editor: vscode.TextEditor,
+  node: SyntaxNode
+) => SelectionWithContext | null;

src/debug.ts

@@ -0,0 +1,24 @@
+import * as vscode from "vscode";
+import { SyntaxNode } from "web-tree-sitter";
+
+export function logBranchTypes(getNodeAtLocation: any) {
+  return (event: vscode.TextEditorSelectionChangeEvent) => {
+    const location = new vscode.Location(
+      vscode.window.activeTextEditor!.document.uri,
+      event.selections[0]
+    );
+
+    const ancestors: SyntaxNode[] = [];
+    let node: SyntaxNode = getNodeAtLocation(location);
+    while (node.parent != null) {
+      ancestors.unshift(node.parent);
+      node = node.parent;
+    }
+
+    ancestors.forEach((node, i) => console.debug(">".repeat(i + 1), node.type));
+    const leafText = ancestors[ancestors.length - 1].text
+      .replace(/\s+/g, " ")
+      .substring(0, 100);
+    console.debug(">".repeat(ancestors.length), `"${leafText}"`);
+  };
+}

src/extension.ts

@@ -14,6 +14,7 @@ import {
   SelectionWithEditor,
 } from "./Types";
 import makeGraph from "./makeGraph";
+import { logBranchTypes } from "./debug";
 
 export async function activate(context: vscode.ExtensionContext) {
   const fontMeasurements = new FontMeasurements(context);
@@ -206,6 +207,9 @@ export async function activate(context: vscode.ExtensionContext) {
     vscode.window.onDidChangeActiveTextEditor(addDecorationsDebounced),
     vscode.window.onDidChangeVisibleTextEditors(addDecorationsDebounced),
     vscode.window.onDidChangeTextEditorSelection(addDecorationsDebounced),
+    vscode.window.onDidChangeTextEditorSelection(
+      logBranchTypes(getNodeAtLocation)
+    ),
     vscode.workspace.onDidChangeTextDocument(handleEdit),
     {
       dispose() {

src/languages/getPojoMatchers.ts

@@ -1,42 +1,41 @@
 import { SyntaxNode } from "web-tree-sitter";
-import { TextEditor } from "vscode";
-import {
-  delimitedMatcher,
-  hasType,
-  simpleSelectionExtractor,
-  makeRange,
-  childNodeMatcher,
-  getNodeWithLeadingDelimiter,
-} from "../nodeMatchers";
 import { getKeyNode, getValueNode } from "../treeSitterUtils";
+import {
+  delimitedSelector,
+  selectWithLeadingDelimiter,
+} from "../nodeSelectors";
+import { composedMatcher, matcher, typeMatcher } from "../nodeMatchers";
+import { nodeFinder, typedNodeFinder } from "../nodeFinders";
 
+// Matchers for "plain old javascript objects", like those found in JSON
 export function getPojoMatchers(
   dictionaryTypes: string[],
   listTypes: string[],
   listElementMatcher: (node: SyntaxNode) => boolean
 ) {
   return {
-    dictionary: hasType(...dictionaryTypes),
-    pair: delimitedMatcher(
-      (node) => node.type === "pair",
-      (node) => node.type === "," || node.type === "}" || node.type === "{",
-      ", "
+    dictionary: typeMatcher(...dictionaryTypes),
+    pair: matcher(
+      nodeFinder((node) => node.type === "pair"),
+      delimitedSelector(
+        (node) => node.type === "," || node.type === "}" || node.type === "{",
+        ", "
+      )
     ),
-    pairKey(editor: TextEditor, node: SyntaxNode) {
-      if (node.type !== "pair") {
-        return null;
-      }
-
-      return simpleSelectionExtractor(getKeyNode(node)!);
-    },
-    value: childNodeMatcher(getValueNode, getNodeWithLeadingDelimiter),
-    list: hasType(...listTypes),
-    listElement: delimitedMatcher(
-      (node) =>
-        listTypes.includes(node.parent?.type ?? "") && listElementMatcher(node),
-      (node) => node.type === "," || node.type === "[" || node.type === "]",
-      ", "
+    pairKey: composedMatcher([typedNodeFinder("pair"), getKeyNode]),
+    value: matcher(getValueNode, selectWithLeadingDelimiter),
+    list: typeMatcher(...listTypes),
+    listElement: matcher(
+      nodeFinder(
+        (node) =>
+          listTypes.includes(node.parent?.type ?? "") &&
+          listElementMatcher(node)
+      ),
+      delimitedSelector(
+        (node) => node.type === "," || node.type === "[" || node.type === "]",
+        ", "
+      )
     ),
-    string: hasType("string"),
+    string: typeMatcher("string"),
   };
 }

src/languages/json.ts

@@ -25,13 +25,16 @@ const nodeMatchers: Record<ScopeType, NodeMatcher> = {
   ...getPojoMatchers(["object"], ["array"], isValue),
   ifStatement: notSupported,
   class: notSupported,
+  className: notSupported,
   statement: notSupported,
   arrowFunction: notSupported,
   functionCall: notSupported,
   argumentOrParameter: notSupported,
   namedFunction: notSupported,
+  functionName: notSupported,
   comment: notSupported,
   type: notSupported,
+  name: notSupported,
 };
 
 export default nodeMatchers;

src/languages/python.ts

@@ -1,14 +1,27 @@
 import { SyntaxNode } from "web-tree-sitter";
 import { getPojoMatchers } from "./getPojoMatchers";
 import {
-  childNodeMatcher,
-  delimitedMatcher,
-  getNodeWithLeadingDelimiter,
-  hasType,
-  possiblyWrappedNode,
+  cascadingMatcher,
+  composedMatcher,
+  matcher,
+  notSupported,
+  typeMatcher,
 } from "../nodeMatchers";
-import { NodeMatcher, ScopeType } from "../Types";
-import { getDefinitionNode } from "../treeSitterUtils";
+import { NodeFinder, NodeMatcher, ScopeType } from "../Types";
+import {
+  getDefinitionNode,
+  getLeftNode,
+  getNameNode,
+} from "../treeSitterUtils";
+import {
+  nodeFinder,
+  typedNodeFinder,
+  findPossiblyWrappedNode,
+} from "../nodeFinders";
+import {
+  delimitedSelector,
+  selectWithLeadingDelimiter,
+} from "../nodeSelectors";
 
 // TODO figure out how to properly use super types
 // Generated by the following command:
@@ -113,10 +126,10 @@ const ARGUMENT_TYPES = [
   "keyword_argument",
 ];
 
-function possiblyDecoratedDefinition(...typeNames: string[]): NodeMatcher {
-  return possiblyWrappedNode(
-    (node) => node.type === "decorated_definition",
-    (node) => typeNames.includes(node.type),
+function possiblyDecoratedDefinition(...typeNames: string[]): NodeFinder {
+  return findPossiblyWrappedNode(
+    typedNodeFinder("decorated_definition"),
+    typedNodeFinder(...typeNames),
     (node) => [getDefinitionNode(node)]
   );
 }
@@ -130,23 +143,39 @@ const nodeMatchers: Record<ScopeType, NodeMatcher> = {
     ["list", "list_comprehension"],
     (node) => LIST_ELEMENT_TYPES.includes(node.type)
   ),
-  ifStatement: hasType("if_statement"),
-  class: possiblyDecoratedDefinition("class_definition"),
-  statement: hasType(...STATEMENT_TYPES),
-  arrowFunction: hasType("lambda"),
-  functionCall: hasType("call"),
-  argumentOrParameter: delimitedMatcher(
-    (node) =>
-      (node.parent?.type === "argument_list" &&
-        ARGUMENT_TYPES.includes(node.type)) ||
-      (PARAMETER_LIST_TYPES.includes(node.parent?.type ?? "") &&
-        PARAMETER_TYPES.includes(node.type)),
-    (node) => node.type === "," || node.type === "(" || node.type === ")",
-    ", "
+  ifStatement: typeMatcher("if_statement"),
+  class: matcher(possiblyDecoratedDefinition("class_definition")),
+  statement: typeMatcher(...STATEMENT_TYPES),
+  name: cascadingMatcher(
+    matcher(getNameNode),
+    matcher((node) => (node.type === "assignment" ? getLeftNode(node) : null))
+  ),
+  functionName: composedMatcher([
+    typedNodeFinder("function_definition"),
+    getNameNode,
+  ]),
+  className: composedMatcher([
+    typedNodeFinder("class_definition"),
+    getNameNode,
+  ]),
+  arrowFunction: typeMatcher("lambda"),
+  functionCall: typeMatcher("call"),
+  argumentOrParameter: matcher(
+    nodeFinder(
+      (node) =>
+        (node.parent?.type === "argument_list" &&
+          ARGUMENT_TYPES.includes(node.type)) ||
+        (PARAMETER_LIST_TYPES.includes(node.parent?.type ?? "") &&
+          PARAMETER_TYPES.includes(node.type))
+    ),
+    delimitedSelector(
+      (node) => node.type === "," || node.type === "(" || node.type === ")",
+      ", "
+    )
   ),
-  namedFunction: possiblyDecoratedDefinition("function_definition"),
-  comment: hasType("comment"),
-  type: childNodeMatcher(getTypeNode, getNodeWithLeadingDelimiter),
+  namedFunction: matcher(possiblyDecoratedDefinition("function_definition")),
+  comment: typeMatcher("comment"),
+  type: matcher(getTypeNode, selectWithLeadingDelimiter),
 };
 
 export default nodeMatchers;