feat(schema): store module source once + byte-offset spans, drop per-node code

Author: rahlk

codeanalyzer/syntactic_analysis/symbol_table_builder.py

@@ -21,6 +21,8 @@
     PyModule,
     PySymbol,
     PyVariableDeclaration,
+    Span,
+    byte_offsets,
 )
 
 
@@ -123,11 +125,12 @@ def build_pymodule_from_file(self, py_file: Path) -> PyModule:
             PyModule.builder()
             .file_path(str(py_file))
             .module_name(py_file.stem)
+            .source(source)
             .comments(self._pycomments(module, source))
             .imports(self._imports(module))
             .variables(self._module_variables(module, script))
-            .classes(self._add_class(module, script))
-            .functions(self._callables(module, script))
+            .classes(self._add_class(module, script, source))
+            .functions(self._callables(module, script, source))
             .content_hash(content_hash)
             .last_modified(last_modified)
             .file_size(file_size)
@@ -183,7 +186,7 @@ def _imports(self, module: ast.Module) -> List[PyImport]:
 
         return imports
 
-    def _add_class(self, node: AST, script: Script, prefix: str = "") -> Dict[str, PyClass]:
+    def _add_class(self, node: AST, script: Script, source: str, prefix: str = "") -> Dict[str, PyClass]:
         classes: Dict[str, PyClass] = {}
 
         for child in ast.iter_child_nodes(node):
@@ -194,6 +197,14 @@ def _add_class(self, node: AST, script: Script, prefix: str = "") -> Dict[str, P
             start_line = child.lineno
             end_line = getattr(child, "end_lineno", start_line + len(child.body))
             code = ast.unparse(child).strip()
+            span = Span(
+                start=(child.lineno, child.col_offset),
+                end=(getattr(child, "end_lineno", child.lineno),
+                     getattr(child, "end_col_offset", child.col_offset)),
+                bytes=byte_offsets(source, child.lineno, child.col_offset,
+                                   getattr(child, "end_lineno", child.lineno),
+                                   getattr(child, "end_col_offset", child.col_offset)),
+            )
 
             # Try resolving full signature with Jedi
             if prefix:
@@ -211,6 +222,7 @@ def _add_class(self, node: AST, script: Script, prefix: str = "") -> Dict[str, P
                 PyClass.builder()
                 .name(class_name)
                 .signature(signature)
+                .span(span)
                 .start_line(start_line)
                 .end_line(end_line)
                 .comments(self._pycomments(child, code))
@@ -219,9 +231,9 @@ def _add_class(self, node: AST, script: Script, prefix: str = "") -> Dict[str, P
                     for base in child.bases
                     if isinstance(base, ast.expr)
                 ])
-                .methods(self._callables(child, script, prefix=signature))  # Pass class signature as prefix
+                .methods(self._callables(child, script, source, prefix=signature))  # Pass class signature as prefix
                 .attributes(self._class_attributes(child, script))
-                .inner_classes(self._add_class(child, script, prefix=signature))  # Pass class signature as prefix
+                .inner_classes(self._add_class(child, script, source, prefix=signature))  # Pass class signature as prefix
                 .build()
             )
 
@@ -230,7 +242,7 @@ def _add_class(self, node: AST, script: Script, prefix: str = "") -> Dict[str, P
         return classes
 
 
-    def _callables(self, node: AST, script: Script, prefix: str = "") -> Dict[str, PyCallable]:
+    def _callables(self, node: AST, script: Script, source: str, prefix: str = "") -> Dict[str, PyCallable]:
         callables: Dict[str, PyCallable] = {}
 
         for child in ast.iter_child_nodes(node):
@@ -239,6 +251,14 @@ def _callables(self, node: AST, script: Script, prefix: str = "") -> Dict[str, P
                 start_line = child.lineno
                 end_line = getattr(child, "end_lineno", start_line + len(child.body))
                 code = ast.unparse(child).strip()
+                span = Span(
+                    start=(child.lineno, child.col_offset),
+                    end=(getattr(child, "end_lineno", child.lineno),
+                         getattr(child, "end_col_offset", child.col_offset)),
+                    bytes=byte_offsets(source, child.lineno, child.col_offset,
+                                       getattr(child, "end_lineno", child.lineno),
+                                       getattr(child, "end_col_offset", child.col_offset)),
+                )
                 decorators = [ast.unparse(d) for d in child.decorator_list]
 
                 if prefix:
@@ -264,6 +284,7 @@ def _callables(self, node: AST, script: Script, prefix: str = "") -> Dict[str, P
                     .name(method_name)  # Use the actual method name, not the full signature
                     .path(str(script.path))
                     .signature(signature)  # Use the full signature here
+                    .span(span)
                     .decorators(decorators)
                     .start_line(start_line)
                     .end_line(end_line)
@@ -278,8 +299,8 @@ def _callables(self, node: AST, script: Script, prefix: str = "") -> Dict[str, P
                         if child.returns else self._infer_type(script, child.lineno, child.col_offset)
                     )
                     .comments(self._pycomments(child, code))
-                    .inner_callables(self._callables(child, script, signature))  # Pass current signature as prefix
-                    .inner_classes(self._add_class(child, script, signature))    # Pass current signature as prefix
+                    .inner_callables(self._callables(child, script, source, signature))  # Pass current signature as prefix
+                    .inner_classes(self._add_class(child, script, source, signature))    # Pass current signature as prefix
                     .build()
                 )
 

test/test_v2_source_spans.py

@@ -1,4 +1,7 @@
+from pathlib import Path
+
 from codeanalyzer.schema.py_schema import byte_offsets
+from codeanalyzer.syntactic_analysis.symbol_table_builder import SymbolTableBuilder
 
 
 def test_byte_offsets_slice_source_exactly():
@@ -12,3 +15,13 @@ def test_byte_offsets_multibyte_safe():
     source = "x = 'é'\ny = 2\n"  # 'é' is 2 bytes in utf-8
     lo, hi = byte_offsets(source, 2, 0, 2, 5)
     assert source.encode("utf-8")[lo:hi].decode("utf-8") == "y = 2"
+
+
+def test_module_stores_source_and_callable_span_slices_it(tmp_path: Path):
+    f = tmp_path / "m.py"
+    f.write_text("def f(a):\n    return a\n", encoding="utf-8")
+    mod = SymbolTableBuilder(tmp_path, None).build_pymodule_from_file(f)
+    assert mod.source == "def f(a):\n    return a\n"
+    fn = next(iter(mod.functions.values()))
+    lo, hi = fn.span.bytes
+    assert mod.source.encode("utf-8")[lo:hi].decode("utf-8").startswith("def f(a):")