Merge pull request 'feat: add AST-backed AST knowledge ingestion for Python files' (#504) from feat/20260413-kb-python-ast into main

scripts/knowledge_base.py

@@ -22,6 +22,7 @@ CLI:
 from __future__ import annotations
 
 import argparse
+import ast
 import json
 import os
 import sys
@@ -137,6 +138,42 @@ class KnowledgeBase:
             self._save(self._persist_path)
         return removed
 
+    def ingest_python_file(
+        self, path: Path, *, module_name: Optional[str] = None, source: str = "ast"
+    ) -> List[Fact]:
+        """Parse a Python file with ``ast`` and assert symbolic structure facts."""
+        tree = ast.parse(path.read_text(), filename=str(path))
+        module = module_name or path.stem
+        fact_source = f"{source}:{path.name}"
+        added: List[Fact] = []
+
+        def add(relation: str, *args: str) -> None:
+            added.append(self.assert_fact(relation, *args, source=fact_source))
+
+        for node in tree.body:
+            if isinstance(node, ast.Import):
+                for alias in node.names:
+                    add("imports", module, alias.name)
+            elif isinstance(node, ast.ImportFrom):
+                prefix = f"{node.module}." if node.module else ""
+                for alias in node.names:
+                    add("imports", module, f"{prefix}{alias.name}")
+            elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                add("defines_function", module, node.name)
+            elif isinstance(node, ast.ClassDef):
+                add("defines_class", module, node.name)
+                for child in node.body:
+                    if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                        add("defines_method", node.name, child.name)
+            elif isinstance(node, ast.Assign):
+                for target in node.targets:
+                    if isinstance(target, ast.Name) and target.id.isupper():
+                        add("defines_constant", module, target.id)
+            elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id.isupper():
+                add("defines_constant", module, node.target.id)
+
+        return added
+
     # ------------------------------------------------------------------
     # Query
     # ------------------------------------------------------------------
@@ -287,6 +324,12 @@ def main() -> None:
         action="store_true",
         help="Dump all facts",
     )
+    parser.add_argument(
+        "--ingest-python",
+        dest="ingest_python",
+        type=Path,
+        help="Parse a Python file with AST and assert symbolic structure facts",
+    )
     parser.add_argument(
         "--relation",
         help="Filter --dump to a specific relation",
@@ -304,6 +347,10 @@ def main() -> None:
         fact = kb.assert_fact(terms[0], *terms[1:], source="cli")
         print(f"Asserted: {fact}")
 
+    if args.ingest_python:
+        added = kb.ingest_python_file(args.ingest_python, source="cli-ast")
+        print(f"Ingested {len(added)} AST fact(s) from {args.ingest_python}")
+
     if args.retract_stmt:
         terms = _parse_terms(args.retract_stmt)
         if len(terms) < 2:

tests/test_knowledge_base_ast.py

@@ -0,0 +1,43 @@
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "scripts"))
+
+from knowledge_base import KnowledgeBase
+
+
+def test_ingest_python_file_extracts_ast_facts(tmp_path: Path) -> None:
+    source = tmp_path / "demo_module.py"
+    source.write_text(
+        "import os\n"
+        "from pathlib import Path\n\n"
+        "CONSTANT = 7\n\n"
+        "def helper(x):\n"
+        "    return x + 1\n\n"
+        "class Demo:\n"
+        "    def method(self):\n"
+        "        return helper(CONSTANT)\n"
+    )
+
+    kb = KnowledgeBase()
+    facts = kb.ingest_python_file(source)
+
+    assert facts, "AST ingestion should add symbolic facts"
+    assert kb.query("defines_function", "demo_module", "helper") == [{}]
+    assert kb.query("defines_class", "demo_module", "Demo") == [{}]
+    assert kb.query("defines_method", "Demo", "method") == [{}]
+    assert kb.query("imports", "demo_module", "os") == [{}]
+    assert kb.query("imports", "demo_module", "pathlib.Path") == [{}]
+    assert kb.query("defines_constant", "demo_module", "CONSTANT") == [{}]
+
+
+def test_ingest_python_file_rejects_invalid_syntax(tmp_path: Path) -> None:
+    broken = tmp_path / "broken.py"
+    broken.write_text("def nope(:\n    pass\n")
+
+    kb = KnowledgeBase()
+    try:
+        kb.ingest_python_file(broken)
+    except SyntaxError:
+        return
+    raise AssertionError("Expected SyntaxError for invalid Python source")