diff --git a/scripts/knowledge_base.py b/scripts/knowledge_base.py index 69e822fb..a290ea6b 100644 --- a/scripts/knowledge_base.py +++ b/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: diff --git a/tests/test_knowledge_base_ast.py b/tests/test_knowledge_base_ast.py new file mode 100644 index 00000000..287b7d00 --- /dev/null +++ b/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")