feed-the-machine 1.3.1 → 1.4.0

package/ftm-map/scripts/parser.py

@@ -1,17 +1,25 @@
 """
-Tree-sitter based parser for extracting symbols and relationships from source code.
-Uses tree-sitter-language-pack for multi-language support and per-language .scm query files.
+Tree-sitter based parser for extracting definition and reference tags from source code.
+
+Uses Aider-style @name.definition.* / @name.reference.* capture convention in
+per-language .scm query files for structured tag extraction. Falls back to Pygments
+lexer for reference extraction when tree-sitter queries lack reference patterns.
 """
 import hashlib
 import os
 import sys
-from dataclasses import dataclass
+from collections import namedtuple
 from pathlib import Path
 from typing import Optional
 
 import tree_sitter as ts
 from tree_sitter_language_pack import get_language, get_parser
 
+# Tag namedtuple: the single output type for all extraction
+# kind is "def" or "ref"
+# rel_fname is relative path, fname is absolute path
+Tag = namedtuple("Tag", ["rel_fname", "fname", "line", "name", "kind"])
+
 QUERIES_DIR = os.path.join(os.path.dirname(__file__), "queries")
 
 # Map file extensions to tree-sitter language names
@@ -35,67 +43,6 @@ EXTENSION_MAP = {
     ".sh": "bash",
 }
 
-# Node types that represent definitions in generic AST walk, keyed by node type
-DEFINITION_TYPES = {
-    # TypeScript / JavaScript
-    "function_declaration": "function",
-    "method_definition": "method",
-    "class_declaration": "class",
-    "arrow_function": "function",
-    "lexical_declaration": "variable",
-    "variable_declaration": "variable",
-    "interface_declaration": "class",
-    "type_alias_declaration": "type",
-    "enum_declaration": "class",
-    # Python
-    "function_definition": "function",
-    "class_definition": "class",
-    "decorated_definition": None,  # unwrap to inner definition
-    # Imports
-    "import_statement": "import",
-    "import_from_statement": "import",
-}
-
-# Node types that carry a symbol name field
-NAME_TYPES = frozenset({
-    "identifier",
-    "property_identifier",
-    "type_identifier",
-})
-
-# Node types representing call expressions per language
-CALL_TYPES = frozenset({
-    "call_expression",  # JS/TS/Go/Rust
-    "call",             # Ruby
-})
-
-# Node types representing import statements
-IMPORT_TYPES = frozenset({
-    "import_statement",
-    "import_from_statement",
-})
-
-
-@dataclass
-class Symbol:
-    name: str
-    kind: str           # function, class, method, variable, import, type
-    file_path: str
-    start_line: int
-    end_line: int
-    signature: str = ""
-    doc_comment: str = ""
-    content_hash: str = ""
-
-
-@dataclass
-class Relationship:
-    source_name: str
-    target_name: str
-    kind: str           # calls, imports, extends, implements, uses
-    source_file: str
-    target_file: str = ""  # may be unknown for cross-file refs
-
 
 # ---------------------------------------------------------------------------
 # Public API
@@ -112,55 +59,44 @@ def compute_content_hash(content: str) -> str:
     return hashlib.sha256(content.encode()).hexdigest()[:16]
 
 
-def parse_file(file_path: str) -> list[Symbol]:
-    """Parse a source file and extract symbols.
+def get_tags(fname: str, rel_fname: str = None) -> list[Tag]:
+    """Extract definition and reference tags from a source file.
+
+    Uses tree-sitter query files with @name.definition.* and @name.reference.*
+    capture naming convention (Aider-style).
 
-    Returns a list of Symbol objects. Handles parse errors gracefully by
-    returning partial results and logging warnings to stderr.
+    Returns list of Tag namedtuples with kind="def" or kind="ref".
     """
-    lang = detect_language(file_path)
+    if rel_fname is None:
+        rel_fname = fname
+
+    lang = detect_language(fname)
     if not lang:
         return []
 
-    source = _read_source(file_path)
+    source = _read_source(fname)
     if source is None:
         return []
 
-    tree = _parse_source(source, lang, file_path)
+    tree = _parse_source(source, lang, fname)
     if tree is None:
         return []
 
     scm_path = os.path.join(QUERIES_DIR, f"{lang}-tags.scm")
-    if os.path.exists(scm_path):
-        symbols = _extract_with_query(tree, source, file_path, lang, scm_path)
-        # Fall back to generic walk if the query produced nothing
-        if symbols:
-            return symbols
-
-    return _extract_generic(tree, source, file_path)
-
-
-def extract_relationships(file_path: str) -> list[Relationship]:
-    """Extract relationships (calls, imports) from a source file.
-
-    Returns a list of Relationship objects.
-    """
-    lang = detect_language(file_path)
-    if not lang:
-        return []
+    if not os.path.exists(scm_path):
+        # No query file -- use Pygments fallback for refs only
+        return _pygments_ref_fallback(source, fname, rel_fname)
 
-    source = _read_source(file_path)
-    if source is None:
-        return []
+    tags = _extract_tags(tree, source, fname, rel_fname, lang, scm_path)
 
-    tree = _parse_source(source, lang, file_path)
-    if tree is None:
-        return []
+    # If we got defs but no refs, supplement with Pygments fallback for refs
+    has_defs = any(t.kind == "def" for t in tags)
+    has_refs = any(t.kind == "ref" for t in tags)
+    if has_defs and not has_refs:
+        ref_tags = _pygments_ref_fallback(source, fname, rel_fname)
+        tags.extend(ref_tags)
 
-    relationships: list[Relationship] = []
-    _extract_calls(tree.root_node, source, file_path, relationships)
-    _extract_imports(tree.root_node, source, file_path, relationships)
-    return relationships
+    return tags
 
 
 # ---------------------------------------------------------------------------
@@ -188,23 +124,15 @@ def _parse_source(source: str, lang: str, file_path: str):
 
 
 # ---------------------------------------------------------------------------
-# Query-based extraction
+# Tag extraction via tree-sitter queries
 # ---------------------------------------------------------------------------
 
-def _extract_with_query(tree, source: str, file_path: str, lang: str, scm_path: str) -> list[Symbol]:
-    """Extract symbols using a language-specific .scm query file.
-
-    Uses tree-sitter QueryCursor.matches() which returns per-pattern match dicts.
-    Each match dict contains both the @definition.X node and the @name node so
-    they are already correlated — no post-hoc joining needed.
-
-    Falls back to empty list on any error so callers can use generic extraction.
-    """
+def _extract_tags(tree, source: str, fname: str, rel_fname: str, lang: str, scm_path: str) -> list[Tag]:
+    """Extract tags using tree-sitter query with @name.definition.* / @name.reference.* convention."""
     try:
         with open(scm_path) as fh:
             query_text = fh.read()
-    except (IOError, OSError) as exc:
-        print(f"Warning: Cannot read query {scm_path}: {exc}", file=sys.stderr)
+    except (IOError, OSError):
         return []
 
     try:
@@ -212,211 +140,76 @@ def _extract_with_query(tree, source: str, file_path: str, lang: str, scm_path:
         query = ts.Query(language, query_text)
         cursor = ts.QueryCursor(query)
         matches = list(cursor.matches(tree.root_node))
-    except Exception as exc:  # noqa: BLE001
-        print(f"Warning: Query execution failed for {file_path}: {exc}", file=sys.stderr)
+    except Exception:  # noqa: BLE001
         return []
 
-    return _process_matches(matches, source, file_path)
+    tags = []
+    seen = {}  # (name, start_byte, end_byte) -> Tag for dedup
 
+    for _pattern_idx, capture_dict in matches:
+        for capture_name, nodes in capture_dict.items():
+            # Only process @name.definition.* and @name.reference.* captures
+            if capture_name.startswith("name.definition."):
+                kind = "def"
+            elif capture_name.startswith("name.reference."):
+                kind = "ref"
+            else:
+                continue
 
-def _process_matches(matches: list, source: str, file_path: str) -> list[Symbol]:
-    """Convert QueryCursor.matches() results into Symbol objects.
+            for node in nodes:
+                name_text = source[node.start_byte:node.end_byte].strip()
+                if not name_text:
+                    continue
 
-    Each match is a (pattern_index, capture_dict) tuple where capture_dict maps
-    capture name → list[Node].  Both @definition.X and @name appear in the same
-    capture_dict for each pattern match, making correlation trivial.
-    """
-    symbols: list[Symbol] = []
+                key = (name_text, node.start_byte, node.end_byte)
+                if key in seen:
+                    continue
 
-    for _pattern_idx, capture_dict in matches:
-        # Find the definition capture (e.g. "definition.function")
-        def_key = next(
-            (k for k in capture_dict if k.startswith("definition.")),
-            None,
-        )
-        if not def_key:
-            continue
-
-        kind = def_key[len("definition."):]
-        def_nodes = capture_dict[def_key]
-        name_nodes = capture_dict.get("name", [])
-
-        if not def_nodes:
-            continue
-
-        def_node = def_nodes[0]
-
-        # Prefer the @name capture; fall back to identifier child walk
-        if name_nodes:
-            sym_name = source[name_nodes[0].start_byte:name_nodes[0].end_byte].strip()
-        else:
-            sym_name = _find_name(def_node, source)
-
-        if not sym_name:
-            continue
-
-        body = source[def_node.start_byte:def_node.end_byte]
-        sig = _first_line(body)
-        doc = _find_doc_comment(def_node, source)
-
-        symbols.append(Symbol(
-            name=sym_name,
-            kind=kind,
-            file_path=file_path,
-            start_line=def_node.start_point[0] + 1,
-            end_line=def_node.end_point[0] + 1,
-            signature=sig,
-            doc_comment=doc,
-            content_hash=compute_content_hash(body),
-        ))
-
-    return symbols
+                line = node.start_point[0] + 1
+                tag = Tag(rel_fname=rel_fname, fname=fname, line=line, name=name_text, kind=kind)
+                seen[key] = tag
+                tags.append(tag)
+
+    return tags
 
 
 # ---------------------------------------------------------------------------
-# Generic AST walk extraction
+# Pygments fallback for references
 # ---------------------------------------------------------------------------
 
-def _extract_generic(tree, source: str, file_path: str) -> list[Symbol]:
-    """Walk the AST and extract symbols without a query file."""
-    symbols: list[Symbol] = []
-    _walk_node(tree.root_node, source, file_path, symbols)
-    return symbols
-
-
-def _walk_node(node, source: str, file_path: str, symbols: list[Symbol]) -> None:
-    """Recursively walk AST nodes looking for definition nodes."""
-    node_type = node.type
-
-    if node_type in DEFINITION_TYPES:
-        kind = DEFINITION_TYPES[node_type]
-
-        if kind is None:
-            # Decorated definition — unwrap inner nodes only
-            for child in node.children:
-                _walk_node(child, source, file_path, symbols)
-            return
-
-        # Skip bare arrow functions without a variable name context
-        if node_type == "arrow_function":
-            for child in node.children:
-                _walk_node(child, source, file_path, symbols)
-            return
-
-        name = _find_name(node, source)
-        if not name:
-            for child in node.children:
-                _walk_node(child, source, file_path, symbols)
-            return
-
-        body = source[node.start_byte:node.end_byte]
-        sig = _first_line(body)
-        doc = _find_doc_comment(node, source)
-
-        symbols.append(Symbol(
-            name=name,
-            kind=kind,
-            file_path=file_path,
-            start_line=node.start_point[0] + 1,
-            end_line=node.end_point[0] + 1,
-            signature=sig,
-            doc_comment=doc,
-            content_hash=compute_content_hash(body),
-        ))
-
-    # Always recurse (definitions can be nested)
-    for child in node.children:
-        _walk_node(child, source, file_path, symbols)
+def _pygments_ref_fallback(source: str, fname: str, rel_fname: str) -> list[Tag]:
+    """Use Pygments to extract reference-like tokens when tree-sitter refs are missing."""
+    try:
+        from pygments.lexers import get_lexer_for_filename
+        from pygments.token import Token
+    except ImportError:
+        return []
 
+    try:
+        lexer = get_lexer_for_filename(fname)
+    except Exception:  # noqa: BLE001
+        return []
 
-# ---------------------------------------------------------------------------
-# Relationship extraction
-# ---------------------------------------------------------------------------
+    tags = []
+    line = 1
+    for token_type, value in lexer.get_tokens(source):
+        # Count newlines for line tracking
+        newlines = value.count('\n')
+        if token_type in Token.Name and value.strip():
+            tags.append(Tag(rel_fname=rel_fname, fname=fname, line=line, name=value.strip(), kind="ref"))
+        line += newlines
 
-def _extract_calls(node, source: str, file_path: str, rels: list[Relationship]) -> None:
-    """Recursively extract function call relationships."""
-    if node.type in CALL_TYPES:
-        func_node = node.children[0] if node.children else None
-        if func_node:
-            callee_text = source[func_node.start_byte:func_node.end_byte]
-            # Simplify dotted paths to last component
-            callee_name = callee_text.split(".")[-1].split("(")[0].strip()
-            caller_name = _find_enclosing_function(node, source)
-            if callee_name and caller_name:
-                rels.append(Relationship(
-                    source_name=caller_name,
-                    target_name=callee_name,
-                    kind="calls",
-                    source_file=file_path,
-                ))
-
-    for child in node.children:
-        _extract_calls(child, source, file_path, rels)
-
-
-def _extract_imports(node, source: str, file_path: str, rels: list[Relationship]) -> None:
-    """Recursively extract import relationships."""
-    if node.type in IMPORT_TYPES:
-        module_stem = Path(file_path).stem
-        _collect_import_names(node, source, module_stem, file_path, rels)
-
-    for child in node.children:
-        _extract_imports(child, source, file_path, rels)
-
-
-def _collect_import_names(
-    node,
-    source: str,
-    module_stem: str,
-    file_path: str,
-    rels: list[Relationship],
-) -> None:
-    """Walk an import node and emit Relationship objects for each imported name."""
-    for child in node.children:
-        child_type = child.type
-
-        if child_type == "dotted_name":
-            imported = source[child.start_byte:child.end_byte]
-            rels.append(Relationship(
-                source_name=module_stem,
-                target_name=imported,
-                kind="imports",
-                source_file=file_path,
-            ))
-
-        elif child_type in ("import_clause", "named_imports", "import_specifier"):
-            for grandchild in child.children:
-                if grandchild.type in NAME_TYPES:
-                    name = source[grandchild.start_byte:grandchild.end_byte].strip()
-                    if name and name not in ("{", "}", ","):
-                        rels.append(Relationship(
-                            source_name=module_stem,
-                            target_name=name,
-                            kind="imports",
-                            source_file=file_path,
-                        ))
-
-        elif child_type == "string":
-            # import ... from "module-path"
-            raw = source[child.start_byte:child.end_byte].strip("'\"")
-            rels.append(Relationship(
-                source_name=module_stem,
-                target_name=raw,
-                kind="imports",
-                source_file=file_path,
-            ))
+    return tags
 
 
 # ---------------------------------------------------------------------------
-# Small AST utilities
+# Small utilities (kept for potential downstream use)
 # ---------------------------------------------------------------------------
 
-def _find_name(node, source: str) -> str:
-    """Find the first name-like identifier child of a node."""
-    for child in node.children:
-        if child.type in NAME_TYPES:
-            return source[child.start_byte:child.end_byte]
-    return ""
+def _first_line(text: str, max_len: int = 200) -> str:
+    """Return the first non-empty line of text, truncated to max_len."""
+    line = text.split("\n")[0].strip()
+    return line[:max_len] + "..." if len(line) > max_len else line
 
 
 def _find_doc_comment(node, source: str) -> str:
@@ -429,27 +222,3 @@ def _find_doc_comment(node, source: str) -> str:
             text = text.strip(marker)
         return text.strip()[:500]
     return ""
-
-
-def _first_line(text: str, max_len: int = 200) -> str:
-    """Return the first non-empty line of text, truncated to max_len."""
-    line = text.split("\n")[0].strip()
-    return line[:max_len] + "..." if len(line) > max_len else line
-
-
-def _find_enclosing_function(node, source: str) -> str:
-    """Walk up the AST to find the name of the nearest enclosing function."""
-    enclosing_types = {
-        "function_declaration",
-        "function_definition",
-        "method_definition",
-        "arrow_function",
-    }
-    current = node.parent
-    while current:
-        if current.type in enclosing_types:
-            name = _find_name(current, source)
-            if name:
-                return name
-        current = current.parent
-    return ""

package/ftm-map/scripts/queries/go-tags.scm

@@ -0,0 +1,20 @@
+; Functions
+(function_declaration
+  name: (identifier) @name.definition.function) @definition.function
+
+; Methods
+(method_declaration
+  name: (field_identifier) @name.definition.method) @definition.method
+
+; Type declarations
+(type_declaration
+  (type_spec
+    name: (type_identifier) @name.definition.type)) @definition.type
+
+; Call references
+(call_expression
+  function: [
+    (identifier) @name.reference.call
+    (selector_expression
+      field: (field_identifier) @name.reference.call)
+  ]) @reference.call

package/ftm-map/scripts/queries/javascript-tags.scm

@@ -1,23 +1,35 @@
 ; Functions
 (function_declaration
-  name: (identifier) @name) @definition.function
+  name: (identifier) @name.definition.function) @definition.function
 
 ; Methods
 (method_definition
-  name: (property_identifier) @name) @definition.method
+  name: (property_identifier) @name.definition.method) @definition.method
 
 ; Classes
 (class_declaration
-  name: (identifier) @name) @definition.class
+  name: (identifier) @name.definition.class) @definition.class
 
-; Arrow functions assigned to const/let variables
+; Arrow functions assigned to const/let
 (lexical_declaration
   (variable_declarator
-    name: (identifier) @name
+    name: (identifier) @name.definition.function
     value: (arrow_function))) @definition.function
 
-; Arrow functions assigned to var variables
+; Arrow functions assigned to var
 (variable_declaration
   (variable_declarator
-    name: (identifier) @name
+    name: (identifier) @name.definition.function
     value: (arrow_function))) @definition.function
+
+; Call references
+(call_expression
+  function: [
+    (identifier) @name.reference.call
+    (member_expression
+      property: (property_identifier) @name.reference.call)
+  ]) @reference.call
+
+; New expressions
+(new_expression
+  constructor: (identifier) @name.reference.class) @reference.class

package/ftm-map/scripts/queries/python-tags.scm

@@ -1,17 +1,31 @@
-; Functions
-(function_definition
-  name: (identifier) @name) @definition.function
+; Module-level constants
+(module
+  (expression_statement
+    (assignment
+      left: (identifier) @name.definition.constant) @definition.constant))
 
 ; Classes
 (class_definition
-  name: (identifier) @name) @definition.class
+  name: (identifier) @name.definition.class) @definition.class
+
+; Functions
+(function_definition
+  name: (identifier) @name.definition.function) @definition.function
 
-; Decorated functions
+; Decorated definitions (functions)
 (decorated_definition
   definition: (function_definition
-    name: (identifier) @name) @definition.function)
+    name: (identifier) @name.definition.function)) @definition.function
 
-; Decorated classes
+; Decorated definitions (classes)
 (decorated_definition
   definition: (class_definition
-    name: (identifier) @name) @definition.class)
+    name: (identifier) @name.definition.class)) @definition.class
+
+; Call references (direct function calls and attribute method calls)
+(call
+  function: [
+    (identifier) @name.reference.call
+    (attribute
+      attribute: (identifier) @name.reference.call)
+  ]) @reference.call

package/ftm-map/scripts/queries/ruby-tags.scm

@@ -0,0 +1,19 @@
+; Methods
+(method
+  name: (identifier) @name.definition.method) @definition.method
+
+; Singleton methods
+(singleton_method
+  name: (identifier) @name.definition.method) @definition.method
+
+; Classes
+(class
+  name: (constant) @name.definition.class) @definition.class
+
+; Modules
+(module
+  name: (constant) @name.definition.module) @definition.module
+
+; Call references
+(call
+  method: (identifier) @name.reference.call) @reference.call

package/ftm-map/scripts/queries/rust-tags.scm

@@ -0,0 +1,37 @@
+; Functions
+(function_item
+  name: (identifier) @name.definition.function) @definition.function
+
+; Structs
+(struct_item
+  name: (type_identifier) @name.definition.struct) @definition.struct
+
+; Enums
+(enum_item
+  name: (type_identifier) @name.definition.enum) @definition.enum
+
+; Traits
+(trait_item
+  name: (type_identifier) @name.definition.trait) @definition.trait
+
+; Impl blocks
+(impl_item
+  trait: (type_identifier) @name.definition.impl) @definition.impl
+
+; Modules
+(mod_item
+  name: (identifier) @name.definition.module) @definition.module
+
+; Macro definitions
+(macro_definition
+  name: (identifier) @name.definition.macro) @definition.macro
+
+; Call references
+(call_expression
+  function: [
+    (identifier) @name.reference.call
+    (field_expression
+      field: (field_identifier) @name.reference.call)
+    (scoped_identifier
+      name: (identifier) @name.reference.call)
+  ]) @reference.call