Merge pull request #49 from codellm-devkit/feat/jedi-shard-planner

PyCG sharding: Jedi planner + adaptive decomposition of runaways

Author: rahlk

codeanalyzer/__main__.py

@@ -7,7 +7,7 @@
 from codeanalyzer.utils import _set_log_level, logger
 from codeanalyzer.config import OutputFormat
 from codeanalyzer.schema import model_dump_json
-from codeanalyzer.options import AnalysisOptions, EmitTarget
+from codeanalyzer.options import AnalysisOptions, EmitTarget, ShardStrategy
 
 
 def main(
@@ -186,6 +186,36 @@ def main(
             min=0,
         ),
     ] = 120,
+    pycg_shard_strategy: Annotated[
+        ShardStrategy,
+        typer.Option(
+            "--pycg-shard-strategy",
+            help=(
+                "How --pycg-shard groups files (level 2 only). 'jedi' (default) "
+                "partitions the Jedi module-dependency graph (SCC + Louvain) so "
+                "tightly-coupled modules co-compute and few call edges are "
+                "severed between shards; import cycles are never split. "
+                "'package' uses the legacy one-shard-per-package-directory "
+                "grouping."
+            ),
+        ),
+    ] = ShardStrategy.JEDI,
+    pycg_max_iter: Annotated[
+        int,
+        typer.Option(
+            "--pycg-max-iter",
+            help=(
+                "Cap on PyCG's fixpoint passes per shard/project (level 2; "
+                "default 50). PyCG iterates until its points-to state stops "
+                "changing, but its access-path domain has no convergence bound, "
+                "so heavy metaclass/mixin code (e.g. an ORM) can loop with each "
+                "pass costing seconds. The cap returns a sound-but-incomplete "
+                "call graph instead of looping until the timeout kills it. "
+                "Set to -1 for PyCG's unbounded run-to-convergence behaviour."
+            ),
+            min=-1,
+        ),
+    ] = 50,
 ):
     options = AnalysisOptions(
         input=input,
@@ -209,6 +239,8 @@ def main(
         pycg_shard=pycg_shard,
         pycg_shard_ceiling=pycg_shard_ceiling,
         pycg_shard_timeout=pycg_shard_timeout,
+        pycg_shard_strategy=pycg_shard_strategy,
+        pycg_max_iter=pycg_max_iter,
     )
 
     _set_log_level(options.verbosity)

codeanalyzer/core.py

@@ -433,8 +433,9 @@ def analyze(self) -> PyApplication:
         logger.info("✅ Jedi: %d edges in %.1fs", len(call_graph), time.perf_counter() - t0_jedi)
 
         if self.analysis_level >= 2:
-            # Level 2: also add PyCG edges.
-            pycg_edges = self._get_pycg_call_graph(symbol_table)
+            # Level 2: also add PyCG edges. The Jedi edges double as the
+            # coupling graph that drives coupling-aware PyCG sharding.
+            pycg_edges = self._get_pycg_call_graph(symbol_table, jedi_edges)
             call_graph = merge_edges(call_graph, pycg_edges)
 
         call_graph = filter_external_edges(call_graph, symbol_table)
@@ -661,13 +662,18 @@ def _build_symbol_table(self, cached_symbol_table: Optional[Dict[str, PyModule]]
     def _get_pycg_call_graph(
         self,
         symbol_table: Dict[str, PyModule],
+        jedi_edges: List[PyCallEdge],
     ) -> List[PyCallEdge]:
         """Build PyCG-resolved call edges.
 
         Runs PyCG's iterative name-pointer analysis over the whole project
         and returns edges with ``provenance=["pycg"]``.  Falls back to an
         empty list and logs a warning on any failure so the caller can
         continue with Jedi-only edges.
+
+        *jedi_edges* are the level-1 call edges; under the ``jedi`` shard
+        strategy they drive coupling-aware partitioning (see
+        :func:`shard_planner.plan_shards`).
         """
         try:
             pycg = PyCG(
@@ -676,9 +682,11 @@ def _get_pycg_call_graph(
                 shard=self.options.pycg_shard,
                 shard_ceiling=self.options.pycg_shard_ceiling,
                 shard_timeout=self.options.pycg_shard_timeout,
+                shard_strategy=self.options.pycg_shard_strategy,
+                max_iter=self.options.pycg_max_iter,
                 using_ray=self.using_ray,
             )
-            return pycg.build_call_graph_edges(symbol_table)
+            return pycg.build_call_graph_edges(symbol_table, jedi_edges=jedi_edges)
         except PyCGExceptions.PyCGImportError as exc:
             logger.warning(f"PyCG not installed — level 2 edges will be Jedi-only: {exc}")
             return []

codeanalyzer/options/__init__.py

@@ -1,3 +1,3 @@
-from .options import AnalysisOptions, EmitTarget, OutputFormat
+from .options import AnalysisOptions, EmitTarget, OutputFormat, ShardStrategy
 
-__all__ = ["AnalysisOptions", "EmitTarget", "OutputFormat"]
+__all__ = ["AnalysisOptions", "EmitTarget", "OutputFormat", "ShardStrategy"]

codeanalyzer/options/options.py

@@ -23,6 +23,20 @@ class EmitTarget(str, Enum):
     SCHEMA = "schema"
 
 
+class ShardStrategy(str, Enum):
+    """How ``--pycg-shard`` groups files into shards (level 2 only).
+
+    - ``jedi``    : partition the Jedi module-dependency graph (strongly-
+                    connected-component condensation + Louvain) so tightly-
+                    coupled modules co-compute and few call edges are severed
+                    between shards. Import cycles are never split.
+    - ``package`` : legacy one-shard-per-package-directory grouping.
+    """
+
+    JEDI = "jedi"
+    PACKAGE = "package"
+
+
 @dataclass
 class AnalysisOptions:
     input: Path
@@ -46,3 +60,5 @@ class AnalysisOptions:
     pycg_shard: bool = False
     pycg_shard_ceiling: int = 100
     pycg_shard_timeout: int = 120
+    pycg_shard_strategy: ShardStrategy = ShardStrategy.JEDI
+    pycg_max_iter: int = 50