feat: add import graph analysis

Made-with: Cursor

Author: Sayeem3051

CHANGELOG.md

@@ -1,5 +1,14 @@
 # Changelog
 
+## [0.1.2] — 2026-04-05
+
+### Added
+
+- Python import graph: `build_import_graph()`, `expand_with_imports()`
+- `ContextEngine(use_import_graph=..., import_graph_depth=...)`
+- `ContextBuilder.use_import_graph(depth=...)` and `no_import_graph()`
+- CLI: `--import-graph` / `--no-import-graph`, `--import-graph-depth N`
+
 ## [0.1.1] — 2026-04-05
 
 ### Added

README.md

@@ -168,6 +168,60 @@ patterns = parse_ctxengignore(Path("."))
 # → list of pattern strings, or [] if no file
 ```
 
+### Import graph (Python)
+
+After files are scored, **ctxeng** parses static ``import`` / ``from … import`` statements in each discovered ``.py`` file, resolves **relative imports** from the file’s location, and can **pull in imported modules** from the same collection set before the token budget is applied.
+
+- **Default:** one hop (`import_graph_depth=1`), relevance for added files = parent score × **0.7**
+- **Edges only** to files already in the current discovery set (filesystem / git / explicit list)
+- Stdlib and third-party imports are ignored (no file under your root → no edge)
+
+```python
+from ctxeng import ContextEngine, ContextBuilder
+
+# Engine: on by default; adjust depth or turn off
+engine = ContextEngine(
+    root=".",
+    use_import_graph=True,
+    import_graph_depth=2,
+)
+
+ctx = (
+    ContextBuilder(".")
+    .for_model("claude-sonnet-4")
+    .use_import_graph(depth=2)   # follow two hops of local imports
+    # .no_import_graph()         # disable expansion
+    .build("Fix the checkout bug in orders")
+)
+```
+
+CLI (import expansion is **on** by default):
+
+```bash
+ctxeng build "Refactor auth" --no-import-graph
+ctxeng build "Refactor auth" --import-graph-depth 2
+```
+
+Lower-level API:
+
+```python
+from pathlib import Path
+from ctxeng import build_import_graph, expand_with_imports
+from ctxeng.models import ContextFile
+
+paths = [Path("src/app.py"), Path("src/lib.py")]
+graph = build_import_graph(Path("."), paths)
+# graph[path] → list of imported paths (within `paths`)
+
+expanded = expand_with_imports(
+    [ContextFile(path=paths[0], content="...", relevance_score=0.9, language="python")],
+    graph,
+    Path("."),
+    max_depth=1,
+    score_decay=0.7,
+)
+```
+
 ---
 
 ## How It Works
@@ -179,7 +233,7 @@ src/auth/login.py  ─┐
 src/models/user.py ─┤  1. Score files         2. Fit budget     <context>
 src/api/routes.py  ─┼─► vs query + git  ─►   smart truncate ─► <file>...</file>
 tests/test_auth.py ─┤     recency + AST        token-aware       <file>...</file>
-...500 more files  ─┘                                           </context>
+...500 more files  ─┘     (+ local Python import expansion)
 ```
 
 ### Scoring signals
@@ -192,6 +246,7 @@ Each file gets a relevance score from 0 → 1, combining:
 | **AST symbols** | Class/function/import names that match the query (Python) |
 | **Path relevance** | Filename and directory names matching query tokens |
 | **Git recency** | Files touched in recent commits score higher |
+| **Import expansion** | After scoring, locally imported Python modules can be added with a decayed score |
 
 ### Token budget optimization
 
@@ -270,6 +325,8 @@ ContextEngine(
     include_patterns=None,  # ["**/*.py"] — only these files
     exclude_patterns=None,  # ["tests/**"] — skip these
     use_git=True,           # Use git recency signal
+    use_import_graph=True,  # Add local Python imports of scored files
+    import_graph_depth=1,    # Hops along the import graph
 )
 ```
 
@@ -298,6 +355,7 @@ ContextBuilder(root=".")
     .with_system("You are an expert Python engineer.")
     .max_file_size(200)     # KB
     .no_git()
+    .use_import_graph(depth=2)  # optional; omit for default depth 1
     .build("query")
 # → Context
 ```
@@ -346,6 +404,10 @@ build options:
   --budget        Override total token budget
   --no-git        Disable git recency scoring
   --max-size      Max file size in KB (default: 500)
+  --import-graph / --no-import-graph
+                  Expand with local Python import graph (default: on)
+  --import-graph-depth N
+                  Import hops when import graph is on (default: 1)
 ```
 
 ---
@@ -382,9 +444,9 @@ You could. But you'll hit these problems immediately:
 - [ ] Semantic similarity scoring
 - [ ] `ctxeng watch` — auto-rebuild on file changes
 - [ ] VSCode extension
-- [ ] Import graph analysis
 - [ ] Cost estimates
 - [ ] Streaming context into LLM APIs
+- [x] Import graph analysis (local Python static imports) ✅
 - [x] `.ctxengignore` file support ✅
 
 ---

ctxeng/__init__.py

@@ -6,14 +6,17 @@
 from ctxeng.builder import ContextBuilder
 from ctxeng.core import ContextEngine
 from ctxeng.ignore import parse_ctxengignore
+from ctxeng.import_graph import build_import_graph, expand_with_imports
 from ctxeng.models import Context, ContextFile, TokenBudget
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"
 __all__ = [
     "ContextEngine",
     "ContextBuilder",
     "Context",
     "ContextFile",
     "TokenBudget",
     "parse_ctxengignore",
+    "build_import_graph",
+    "expand_with_imports",
 ]

ctxeng/builder.py

@@ -40,6 +40,8 @@ def __init__(self, root: str | Path = ".") -> None:
         self._system = ""
         self._max_file_size_kb = 500
         self._use_git = True
+        self._use_import_graph = True
+        self._import_graph_depth = 1
 
     def for_model(self, model: str) -> ContextBuilder:
         """Set the target model (determines token budget)."""
@@ -87,6 +89,17 @@ def no_git(self) -> ContextBuilder:
         self._use_git = False
         return self
 
+    def use_import_graph(self, depth: int = 1) -> ContextBuilder:
+        """Follow local Python imports from scored files (default depth 1)."""
+        self._use_import_graph = True
+        self._import_graph_depth = depth
+        return self
+
+    def no_import_graph(self) -> ContextBuilder:
+        """Disable import-graph expansion."""
+        self._use_import_graph = False
+        return self
+
     def build(self, query: str = "") -> Context:
         """
         Build and return the optimized Context.
@@ -105,6 +118,8 @@ def build(self, query: str = "") -> Context:
             include_patterns=self._include or None,
             exclude_patterns=self._exclude or None,
             use_git=self._use_git,
+            use_import_graph=self._use_import_graph,
+            import_graph_depth=self._import_graph_depth,
         )
         return engine.build(
             query=query,

ctxeng/cli.py

@@ -37,6 +37,10 @@ def cmd_build(args: argparse.Namespace) -> None:
         builder = builder.with_system(args.system)
     if args.no_git:
         builder = builder.no_git()
+    if not args.import_graph:
+        builder = builder.no_import_graph()
+    else:
+        builder = builder.use_import_graph(depth=args.import_graph_depth)
     if args.budget:
         builder = builder.with_budget(args.budget)
 
@@ -138,6 +142,19 @@ def main() -> None:
                          help="Override token budget total")
     build_p.add_argument("--max-size", type=int, default=500,
                          help="Max file size in KB (default: 500)")
+    build_p.add_argument(
+        "--import-graph",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Expand context using local Python import graph (default: on)",
+    )
+    build_p.add_argument(
+        "--import-graph-depth",
+        type=int,
+        default=1,
+        metavar="N",
+        help="Import hops when --import-graph is on (default: 1)",
+    )
     build_p.set_defaults(func=cmd_build)
 
     # --- info ---

ctxeng/core.py

@@ -4,6 +4,7 @@
 
 from pathlib import Path
 
+from ctxeng.import_graph import build_import_graph, expand_with_imports
 from ctxeng.models import Context, ContextFile, TokenBudget
 from ctxeng.optimizer import count_tokens, detect_language, optimize_budget
 from ctxeng.scorer import rank_files
@@ -31,6 +32,8 @@ class ContextEngine:
         include_patterns:   Only include files matching these glob patterns.
         exclude_patterns:   Skip files matching these glob patterns.
         use_git:            Score files using git recency signal (default: True).
+        use_import_graph:   Pull in files imported by high-scoring Python modules (default: True).
+        import_graph_depth: How many import hops to follow (default: 1).
     """
 
     def __init__(
@@ -42,6 +45,8 @@ def __init__(
         include_patterns: list[str] | None = None,
         exclude_patterns: list[str] | None = None,
         use_git: bool = True,
+        use_import_graph: bool = True,
+        import_graph_depth: int = 1,
     ) -> None:
         self.root = Path(root).resolve()
         self.model = model
@@ -50,6 +55,8 @@ def __init__(
         self.include_patterns = include_patterns
         self.exclude_patterns = exclude_patterns
         self.use_git = use_git
+        self.use_import_graph = use_import_graph
+        self.import_graph_depth = import_graph_depth
 
     def build(
         self,
@@ -108,6 +115,16 @@ def build(
             for path, content, score in ranked
         ]
 
+        if self.use_import_graph and self.import_graph_depth > 0:
+            paths_in_raw = [p for p, _ in raw]
+            graph = build_import_graph(self.root, paths_in_raw)
+            context_files = expand_with_imports(
+                context_files,
+                graph,
+                self.root,
+                max_depth=self.import_graph_depth,
+            )
+
         # 4. Optimize for token budget
         query_tokens = count_tokens(query, self.model) if query else 0
         system_tokens = count_tokens(system_prompt, self.model) if system_prompt else 0