ast-search-python 0.1.0

package/build/index.d.ts

@@ -0,0 +1,14 @@
+import type { LanguageBackend, LanguageRegistry, Match } from "ast-search/plugin";
+export declare class PythonLanguageBackend implements LanguageBackend {
+    readonly langId = "python";
+    readonly name = "Python";
+    readonly extensions: Set<string>;
+    parse(source: string, _filePath: string): unknown;
+    query(ast: unknown, selector: string, source: string, filePath: string): Match[];
+    validateSelector(selector: string): void;
+}
+/**
+ * Register the Python backend with an ast-search LanguageRegistry.
+ * Called by ast-search core when --plugin ast-search-python is passed.
+ */
+export declare function register(registry: LanguageRegistry): void;

package/build/index.js

@@ -0,0 +1,53 @@
+import { createRequire } from "module";
+import { expandShorthands } from "./shorthands.js";
+import { runTreeSitterQuery, validateTreeSitterQuery } from "./query.js";
+// tree-sitter is a CommonJS module; use createRequire for ESM compat.
+const _require = createRequire(import.meta.url);
+// Lazy-initialize the parser so the native addon is only loaded when
+// a Python file is actually parsed (not at module import time).
+let _parser;
+let _language;
+let _QueryClass;
+function getRuntime() {
+    if (!_parser) {
+        const Parser = _require("tree-sitter");
+        const pythonModule = _require("tree-sitter-python");
+        _language = pythonModule.language;
+        _QueryClass = Parser.Query;
+        const p = new Parser();
+        p.setLanguage(pythonModule);
+        _parser = p;
+    }
+    return {
+        parser: _parser,
+        language: _language,
+        QueryClass: _QueryClass,
+    };
+}
+export class PythonLanguageBackend {
+    constructor() {
+        this.langId = "python";
+        this.name = "Python";
+        this.extensions = new Set([".py", ".pyw"]);
+    }
+    parse(source, _filePath) {
+        const { parser } = getRuntime();
+        return parser.parse(source);
+    }
+    query(ast, selector, source, filePath) {
+        const { language, QueryClass } = getRuntime();
+        const expanded = expandShorthands(selector);
+        return runTreeSitterQuery(ast, expanded, source, filePath, language, QueryClass);
+    }
+    validateSelector(selector) {
+        const { language, QueryClass } = getRuntime();
+        validateTreeSitterQuery(expandShorthands(selector), language, QueryClass);
+    }
+}
+/**
+ * Register the Python backend with an ast-search LanguageRegistry.
+ * Called by ast-search core when --plugin ast-search-python is passed.
+ */
+export function register(registry) {
+    registry.register(new PythonLanguageBackend());
+}

package/build/query.d.ts

@@ -0,0 +1,23 @@
+import type { Match } from "ast-search/plugin";
+interface TSNode {
+    startPosition: {
+        row: number;
+        column: number;
+    };
+    startIndex: number;
+    endIndex: number;
+    type: string;
+}
+interface TSCapture {
+    node: TSNode;
+    name: string;
+}
+interface TSQueryConstructor {
+    new (language: unknown, pattern: string): TSQuery;
+}
+interface TSQuery {
+    captures(node: TSNode): TSCapture[];
+}
+export declare function runTreeSitterQuery(ast: unknown, pattern: string, source: string, filePath: string, language: unknown, QueryClass: TSQueryConstructor): Match[];
+export declare function validateTreeSitterQuery(pattern: string, language: unknown, QueryClass: TSQueryConstructor): void;
+export {};

package/build/query.js

@@ -0,0 +1,41 @@
+function assertValidPattern(pattern) {
+    const trimmed = pattern.trim();
+    // A valid tree-sitter pattern must start with "(" (S-expression) or contain
+    // "@" (capture). Bare words crash tree-sitter with a native SIGSEGV.
+    if (!trimmed.startsWith("(") && !trimmed.includes("@")) {
+        throw new Error(`Invalid tree-sitter query "${trimmed}": must start with "(" or include a "@capture". ` +
+            `Use a shorthand (e.g. "fn", "call") or write a full S-expression (e.g. "(function_definition) @fn").`);
+    }
+}
+export function runTreeSitterQuery(ast, pattern, source, filePath, language, QueryClass) {
+    assertValidPattern(pattern);
+    const tree = ast;
+    const q = new QueryClass(language, pattern);
+    const captures = q.captures(tree.rootNode);
+    const seen = new Set();
+    const results = [];
+    for (const capture of captures) {
+        const node = capture.node;
+        if (seen.has(node))
+            continue;
+        seen.add(node);
+        const text = source.slice(node.startIndex, node.endIndex);
+        const firstLine = text.split("\n")[0].trimEnd();
+        results.push({
+            file: filePath,
+            line: node.startPosition.row + 1, // 1-indexed to match JS backend
+            col: node.startPosition.column,
+            source: firstLine,
+        });
+    }
+    return results;
+}
+export function validateTreeSitterQuery(pattern, language, QueryClass) {
+    try {
+        assertValidPattern(pattern);
+        new QueryClass(language, pattern);
+    }
+    catch (e) {
+        throw new Error(`Invalid tree-sitter query: ${e instanceof Error ? e.message : String(e)}`);
+    }
+}

package/build/shorthands.d.ts

@@ -0,0 +1,2 @@
+export declare const PYTHON_SHORTHANDS: Record<string, string>;
+export declare function expandShorthands(selector: string): string;

package/build/shorthands.js

@@ -0,0 +1,68 @@
+// Shorthands for common Python AST node types, expressed as tree-sitter
+// S-expression patterns. Each shorthand includes a `@_` capture so that
+// `captures()` returns results without the user needing to add a capture name.
+//
+// For attribute matching or complex patterns, write tree-sitter S-expressions
+// directly (and include at least one capture):
+//   (call function: (identifier) @name (#eq? @name "foo"))
+//
+// NOTE: In tree-sitter-python 0.21+, async functions are still
+// `function_definition` (no separate `async_function_definition` type).
+// To match only async functions, use the full query:
+//   (function_definition . "async" .) @fn
+export const PYTHON_SHORTHANDS = {
+    // Shared semantics with JS backend
+    fn: "(function_definition) @_",
+    call: "(call) @_",
+    class: "(class_definition) @_",
+    assign: "(assignment) @_",
+    return: "(return_statement) @_",
+    await: "(await) @_",
+    yield: "(yield) @_",
+    import: "(import_statement) @_",
+    if: "(if_statement) @_",
+    for: "(for_statement) @_",
+    while: "(while_statement) @_",
+    // Python-specific
+    from: "(import_from_statement) @_",
+    raise: "(raise_statement) @_",
+    with: "(with_statement) @_",
+    lambda: "(lambda) @_",
+    comp: "(list_comprehension) @_",
+    dictcomp: "(dictionary_comprehension) @_",
+    setcomp: "(set_comprehension) @_",
+    genexp: "(generator_expression) @_",
+    decorator: "(decorator) @_",
+    assert: "(assert_statement) @_",
+    delete: "(delete_statement) @_",
+    global: "(global_statement) @_",
+    nonlocal: "(nonlocal_statement) @_",
+    augassign: "(augmented_assignment) @_",
+    decorated: "(decorated_definition) @_",
+};
+export function expandShorthands(selector) {
+    // Replace bare shorthand words (not inside quotes and not preceded by @)
+    // with their S-expression. The negative lookbehind prevents expanding
+    // capture names like @fn into @(function_definition) @_.
+    const keys = Object.keys(PYTHON_SHORTHANDS);
+    const pattern = new RegExp(`(?<!@)\\b(${keys.join("|")})\\b`, "g");
+    const parts = [];
+    let i = 0;
+    while (i < selector.length) {
+        const ch = selector[i];
+        if (ch === '"' || ch === "'") {
+            let j = i + 1;
+            while (j < selector.length && selector[j] !== ch)
+                j++;
+            parts.push(selector.slice(i, j + 1));
+            i = j + 1;
+        }
+        else {
+            const nextQuote = selector.slice(i).search(/['"]/);
+            const segment = nextQuote === -1 ? selector.slice(i) : selector.slice(i, i + nextQuote);
+            parts.push(segment.replace(pattern, (m) => { var _a; return (_a = PYTHON_SHORTHANDS[m]) !== null && _a !== void 0 ? _a : m; }));
+            i = nextQuote === -1 ? selector.length : i + nextQuote;
+        }
+    }
+    return parts.join("");
+}

package/package.json

@@ -0,0 +1,50 @@
+{
+  "name": "ast-search-python",
+  "version": "0.1.0",
+  "description": "Python language plugin for ast-search",
+  "type": "module",
+  "main": "./build/index.js",
+  "exports": {
+    ".": {
+      "import": "./build/index.js",
+      "types": "./build/index.d.ts"
+    }
+  },
+  "scripts": {
+    "build": "tsc",
+    "test": "NODE_OPTIONS=--experimental-vm-modules jest",
+    "prepublishOnly": "tsc"
+  },
+  "author": "shiplet",
+  "license": "ISC",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/willey-shiplet/ast-search.git",
+    "directory": "packages/ast-search-python"
+  },
+  "files": [
+    "build/*.js",
+    "build/*.d.ts",
+    "README.md"
+  ],
+  "publishConfig": {
+    "access": "public",
+    "provenance": true
+  },
+  "dependencies": {
+    "tree-sitter": "^0.21.1",
+    "tree-sitter-python": "^0.21.0"
+  },
+  "devDependencies": {
+    "@jest/globals": "^29.7.0",
+    "@types/jest": "^29.5.12",
+    "@types/node": "^20.12.7",
+    "ast-search": "workspace:*",
+    "jest": "^29.7.0",
+    "ts-jest": "^29.1.2",
+    "typescript": "^5.4.5"
+  },
+  "peerDependencies": {
+    "ast-search": ">=0.3.0"
+  }
+}