diff --git a/codeanalyzer/__main__.py b/codeanalyzer/__main__.py index 6693cba..51f227f 100644 --- a/codeanalyzer/__main__.py +++ b/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) diff --git a/codeanalyzer/core.py b/codeanalyzer/core.py index 5fc4003..151446e 100644 --- a/codeanalyzer/core.py +++ b/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,6 +662,7 @@ 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. @@ -668,6 +670,10 @@ def _get_pycg_call_graph( 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 [] diff --git a/codeanalyzer/options/__init__.py b/codeanalyzer/options/__init__.py index 127a183..317d42b 100644 --- a/codeanalyzer/options/__init__.py +++ b/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"] \ No newline at end of file +__all__ = ["AnalysisOptions", "EmitTarget", "OutputFormat", "ShardStrategy"] \ No newline at end of file diff --git a/codeanalyzer/options/options.py b/codeanalyzer/options/options.py index 41eb404..4e8662c 100644 --- a/codeanalyzer/options/options.py +++ b/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 diff --git a/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py b/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py index 9a8ece9..e82c639 100644 --- a/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py +++ b/codeanalyzer/semantic_analysis/pycg/pycg_analysis.py @@ -43,12 +43,14 @@ import importlib.util # noqa: F401 import contextlib import json # noqa: F401 +import shutil import signal +import tempfile import time from collections import Counter, defaultdict from pathlib import Path -from typing import Any, Dict, Generator, List, Optional, Set, Union +from typing import Any, Dict, Generator, List, Optional, Set, Tuple, Union @contextlib.contextmanager @@ -79,20 +81,87 @@ def _handler(signum: int, frame: object) -> None: from codeanalyzer.schema.py_schema import PyCallEdge, PyModule from codeanalyzer.semantic_analysis.call_graph import iter_callables_in_symbol_table from codeanalyzer.semantic_analysis.pycg.pycg_exceptions import PyCGExceptions +from codeanalyzer.semantic_analysis.pycg.shard_planner import plan_shards from codeanalyzer.utils import ProgressBar, logger +def _materialize_shard_root( + files: List[str], + project_dir: Path, +) -> Tuple[Path, List[str]]: + """Build a temporary symlink mini-project for a shard; return ``(root, eps)``. + + PyCG bounds its import-following to the ``package`` directory — only + modules whose resolved file lives under that root are followed; everything + else becomes a ghost node (``ImportManager``: ``if self.mod_dir not in + mod.__file__: return``). A coupling-derived shard is an arbitrary set of + files that need not form a directory, so we mirror the project layout into + a temp dir holding symlinks to exactly the shard's files plus the + ``__init__.py`` chain each needs for package resolution. Running PyCG with + this mirror as the package root confines analysis to the shard while + emitting project-relative edge names (so ``prefix=""`` — no rename needed). + + The caller owns the returned *root* and must ``shutil.rmtree`` it. + """ + root = Path(tempfile.mkdtemp(prefix="canpy_pycg_shard_")) + entry_points: List[str] = [] + linked_inits: Set[Path] = set() + for f in files: + src = Path(f).resolve() + try: + rel = src.relative_to(project_dir) + except ValueError: + continue # defensively skip files outside the project + dst = root / rel + dst.parent.mkdir(parents=True, exist_ok=True) + if not dst.exists(): + dst.symlink_to(src) + entry_points.append(str(dst)) + + # Symlink the __init__.py chain from project root down to this file's + # package so PyCG/importlib can resolve the dotted module name. These + # add ~0 analysis cost (usually empty) and keep out-of-shard siblings + # unresolved → ghost nodes. + for i in range(len(rel.parent.parts) + 1): + pkg_rel = Path(*rel.parent.parts[:i]) + real_init = project_dir / pkg_rel / "__init__.py" + link_init = root / pkg_rel / "__init__.py" + if real_init.exists() and link_init not in linked_inits: + link_init.parent.mkdir(parents=True, exist_ok=True) + if not link_init.exists(): + link_init.symlink_to(real_init.resolve()) + linked_inits.add(link_init) + return root, entry_points + + +@contextlib.contextmanager +def _shard_symlink_root( + files: List[str], + project_dir: Path, +) -> Generator[Tuple[Path, List[str]], None, None]: + """Context-manager wrapper around :func:`_materialize_shard_root`. + + Yields ``(root, entry_points)`` and removes the temp tree on exit. + """ + root, entry_points = _materialize_shard_root(files, project_dir) + try: + yield root, entry_points + finally: + shutil.rmtree(root, ignore_errors=True) + + def _pycg_shard_worker( entry_points: List[str], package_dir: str, prefix: str, + max_iter: int = -1, ) -> List[tuple]: """Run PyCG on one shard; called in a Ray worker process. Returns a list of ``(source, target, weight)`` tuples that the caller converts to :class:`PyCallEdge` objects. This function is a plain module-level callable so it can be pickled by Ray without capturing any - class-level state. + class-level state. *max_iter* caps PyCG's fixpoint passes (-1 = unbounded). """ import importlib import sys @@ -123,7 +192,7 @@ def _pycg_shard_worker( cg = CallGraphGenerator( entry_points=entry_points, package=package_dir, - max_iter=-1, + max_iter=max_iter, operation="call-graph", ) cg.analyze() @@ -298,7 +367,29 @@ class PyCG: # --pycg-shard-timeout. Set to 0 to disable. _PYCG_SHARD_TIMEOUT: int = 120 - # Directory names that should never be fed to PyCG as entry points. + # Cap on PyCG's outer fixpoint passes. PyCG runs PostProcessor until the + # def/scope/MRO state stops changing; its abstract domain (field-sensitive + # access paths, no k-limiting or widening) has no ascending-chain bound, so + # on heavy metaclass/mixin code (e.g. an ORM) the def set can balloon into + # the thousands and each O(defs^2) pass costs seconds — convergence, if it + # comes, takes many passes. A finite cap turns "loop until killed" into a + # sound-but-incomplete result that still returns the edges found so far. + # 50 is generous — well-behaved code converges in well under 20 passes — + # while bounding the pathological case. Override via --pycg-max-iter; + # -1 restores PyCG's unbounded run-to-convergence behaviour. + _PYCG_MAX_ITER: int = 50 + + # Iterative decomposition of runaway (timed-out) shards: a shard that the + # wall-clock timeout kills is re-partitioned at half the budget and re-run, + # down to this file-count floor. Below the floor — or for an atomic import + # cycle that won't split — the residue falls back to Jedi-only coverage. + _PYCG_DECOMP_FLOOR: int = 10 + _PYCG_MAX_DECOMP_ROUNDS: int = 6 + + # Directory names that should never be fed to PyCG as entry points, nor + # followed into during import resolution (an in-tree .codeanalyzer venv / + # site-packages lives under project_dir and would otherwise be pulled into + # the package bound and analysed — see _shard_symlink_root). _SKIP_DIRS: frozenset = frozenset({ ".codeanalyzer", ".git", "__pycache__", "venv", ".venv", "virtualenv", "env", ".env", @@ -313,6 +404,8 @@ def __init__( shard: bool = False, shard_ceiling: Optional[int] = None, shard_timeout: Optional[int] = None, + shard_strategy: str = "jedi", + max_iter: Optional[int] = None, using_ray: bool = False, ) -> None: self.project_dir = Path(project_dir).resolve() @@ -324,8 +417,32 @@ def __init__( self.shard_timeout = ( shard_timeout if shard_timeout is not None else self._PYCG_SHARD_TIMEOUT ) + self.max_iter = max_iter if max_iter is not None else self._PYCG_MAX_ITER + # "jedi": partition the Jedi module graph (SCC + Louvain) so coupled + # modules co-compute and few edges are severed (see shard_planner). + # "package": legacy one-shard-per-package-directory grouping. + self.shard_strategy = shard_strategy self.using_ray = using_ray self._CallGraphGenerator: Optional[Any] = None + self._resolver: Optional["_PyCGCallableResolver"] = None + + @staticmethod + def _coalesce_edges(edges: List[PyCallEdge]) -> List[PyCallEdge]: + """Sum weights of duplicate ``(source, target)`` pairs across shards.""" + merged: Dict[tuple, PyCallEdge] = {} + for edge in edges: + key = (edge.source, edge.target) + if key in merged: + existing = merged[key] + merged[key] = PyCallEdge( + source=existing.source, + target=existing.target, + weight=existing.weight + edge.weight, + provenance=existing.provenance, + ) + else: + merged[key] = edge + return list(merged.values()) # ------------------------------------------------------------------ # Entry-point collection @@ -428,7 +545,7 @@ def _run_pycg_batch( cg = self._CallGraphGenerator( entry_points=entry_points, package=str(package_dir), - max_iter=-1, + max_iter=self.max_iter, operation="call-graph", ) cg.analyze() @@ -455,6 +572,201 @@ def _run_pycg_batch( # Sharded analysis # ------------------------------------------------------------------ + def _build_sharded_planned( + self, + jedi_edges: List[PyCallEdge], + symbol_table: Dict[str, PyModule], + resolver: "_PyCGCallableResolver", + ) -> List[PyCallEdge]: + """Coupling-aware sharding with iterative decomposition of runaways. + + Shards are chosen to *minimise the call edges severed between shards*: + :func:`shard_planner.plan_shards` condenses the Jedi call graph by + strongly-connected component (so import cycles never split) and clusters + it with Louvain so tightly-coupled modules land together. Each shard is + run through PyCG via a symlinked mini-project that bounds analysis to its + files. + + PyCG's fixpoint diverges on heavy metaclass/mixin clusters, and a uniform + ceiling would force *every* shard small (severing many edges) just to tame + the few that run away. Instead we start coarse (low cut, high recall on + healthy code) and **only re-decompose the shards that time out**: each + runaway's files are re-partitioned at half the budget and re-run, down to + a floor. A runaway shard contributes zero edges, so splitting it recovers + almost all of them while paying cut on its internal seams alone. The + residue that still diverges at the floor (or is an atomic cycle that won't + split) falls back to Jedi-only coverage. + """ + self._resolver = resolver + plan = plan_shards( + symbol_table, jedi_edges, budget=self.shard_ceiling, merge_small=True + ) + m = plan.metrics + logger.info( + "PyCG: planned %d shard(s) from Jedi module graph " + "(cut_ratio=%.3f, max_shard=%d files, %d modules)", + int(m["num_shards"]), m["cut_ratio"], + int(m["max_shard_files"]), int(m["modules"]), + ) + + runner = ( + self._run_fileset_shards_ray if self.using_ray + else self._run_fileset_shards_seq + ) + all_edges: List[PyCallEdge] = [] + shards = plan.shards + budget = self.shard_ceiling + converged_total = 0 + irreducible_files = 0 + round_no = 0 + + while shards: + label = "decomposition round %d (budget %d, %d shard(s))" % ( + round_no, budget, len(shards), + ) + logger.info("PyCG: %s", label) + edges, runaways = runner(shards) + all_edges.extend(edges) + converged_total += len(shards) - len(runaways) + if not runaways: + break + + next_budget = max(self._PYCG_DECOMP_FLOOR, budget // 2) + stop_decomposing = ( + round_no >= self._PYCG_MAX_DECOMP_ROUNDS or next_budget >= budget + ) + + next_shards: List[List[str]] = [] + for rf in runaways: + # Re-partition this runaway's files alone, at a tighter budget. + # An atomic cycle (or a lone file) that won't shrink is + # irreducible — accept Jedi-only rather than loop forever. + sub_st = {f: symbol_table[f] for f in rf if f in symbol_table} + if stop_decomposing or len(rf) <= 1: + irreducible_files += len(rf) + continue + sub_plan = plan_shards(sub_st, jedi_edges, budget=next_budget) + if len(sub_plan.shards) <= 1: + # did not actually split (one atomic SCC) — give up on it + irreducible_files += len(rf) + continue + next_shards.extend(sub_plan.shards) + + if not next_shards: + break + logger.info( + "PyCG: %d shard(s) ran away — decomposing into %d sub-shard(s) " + "at budget %d", len(runaways), len(next_shards), next_budget, + ) + shards, budget = next_shards, next_budget + round_no += 1 + + if irreducible_files: + logger.warning( + "PyCG: %d file(s) in irreducibly-divergent shards fall back to " + "Jedi-only coverage", irreducible_files, + ) + + result = self._coalesce_edges(all_edges) + logger.info( + "PyCG: %d edges from %d converged shard(s) over %d round(s) " + "(%d before dedup, Jedi-planned%s)", + len(result), converged_total, round_no + 1, len(all_edges), + ", Ray-parallel" if self.using_ray else "", + ) + return result + + def _run_fileset_shards_seq( + self, shards: List[List[str]], + ) -> Tuple[List[PyCallEdge], List[List[str]]]: + """Run each file-set shard sequentially; return ``(edges, runaways)``. + + A shard that times out or raises is returned in *runaways* (its file + list) for the caller to re-decompose; it contributes no edges. + """ + resolver = self._resolver + edges_all: List[PyCallEdge] = [] + runaways: List[List[str]] = [] + with ProgressBar(len(shards), "Building call graph shards", item_label="shards") as progress: + for files in shards: + try: + with _shard_symlink_root(files, self.project_dir) as (root, eps): + with _shard_timeout(self.shard_timeout): + edges = self._run_pycg_batch(eps, root, resolver, prefix="") + edges_all.extend(edges) + except (TimeoutError, PyCGExceptions.PyCGAnalysisError): + runaways.append(files) + progress.advance() + return edges_all, runaways + + def _run_fileset_shards_ray( + self, shards: List[List[str]], + ) -> Tuple[List[PyCallEdge], List[List[str]]]: + """Ray-parallel variant of :meth:`_run_fileset_shards_seq`. + + Each shard is materialised as a symlink mini-project up front (the trees + must outlive their remote tasks), submitted as a Ray task, and collected + against one wall-clock deadline — Ray workers cannot use SIGALRM, so the + timeout is enforced orchestrator-side. Timed-out/failed shards become + runaways; symlink trees are removed once the batch completes. + """ + import os + import ray + + os.environ.setdefault("RAY_IGNORE_UNHANDLED_ERRORS", "1") + remote_fn = ray.remote(_pycg_shard_worker) + + roots: List[Path] = [] + futures: List[Any] = [] + meta: Dict[Any, List[str]] = {} # ObjectRef -> shard file list + edges_all: List[PyCallEdge] = [] + runaways: List[List[str]] = [] + try: + with ProgressBar(len(shards), "Building call graph shards (parallel)", item_label="shards") as progress: + for files in shards: + root, eps = _materialize_shard_root(files, self.project_dir) + roots.append(root) + fut = remote_fn.remote(eps, str(root), "", self.max_iter) + futures.append(fut) + meta[fut] = files + + deadline = ( + time.perf_counter() + float(self.shard_timeout) + if self.shard_timeout > 0 else None + ) + pending = list(futures) + while pending: + if deadline is not None: + remaining = deadline - time.perf_counter() + if remaining <= 0: + break + else: + remaining = None + + ready, pending = ray.wait(pending, num_returns=1, timeout=remaining) + if not ready: + break + + fut = ready[0] + try: + triples = ray.get(fut) + edges_all.extend( + PyCallEdge(source=s, target=t, weight=w, provenance=["pycg"]) + for s, t, w in triples + ) + except Exception: + runaways.append(meta[fut]) + progress.advance() + + for fut in pending: # exceeded the deadline + ray.cancel(fut, force=True) + runaways.append(meta[fut]) + progress.advance() + finally: + for root in roots: + shutil.rmtree(root, ignore_errors=True) + return edges_all, runaways + def _build_sharded( self, entry_points: List[str], @@ -582,7 +894,7 @@ def _build_sharded_ray(self, shards: Dict[Path, List[str]]) -> List[PyCallEdge]: progress.advance() continue prefix = self._package_prefix(pkg_root, self.project_dir) - fut = remote_fn.remote(files, str(pkg_root), prefix) + fut = remote_fn.remote(files, str(pkg_root), prefix, self.max_iter) futures.append(fut) meta[fut] = (pkg_label, n) @@ -668,7 +980,9 @@ def _build_sharded_ray(self, shards: Dict[Path, List[str]]) -> List[PyCallEdge]: # ------------------------------------------------------------------ def build_call_graph_edges( - self, symbol_table: Dict[str, PyModule] + self, + symbol_table: Dict[str, PyModule], + jedi_edges: Optional[List[PyCallEdge]] = None, ) -> List[PyCallEdge]: """Run PyCG and return ``PyCallEdge`` entries with ``provenance=["pycg"]``. @@ -701,12 +1015,21 @@ def build_call_graph_edges( if n_files > self._PYCG_FILE_CEILING: if self.shard: - mode = "Ray-parallel" if self.using_ray else "sequential" - logger.info( - "PyCG: starting sharded call graph analysis (%d files, %s)", - n_files, mode, - ) - edges = self._build_sharded(entry_points, resolver) + if self.shard_strategy == "jedi" and jedi_edges is not None: + logger.info( + "PyCG: starting Jedi-planned sharded analysis (%d files)", + n_files, + ) + edges = self._build_sharded_planned( + jedi_edges, symbol_table, resolver + ) + else: + mode = "Ray-parallel" if self.using_ray else "sequential" + logger.info( + "PyCG: starting per-package sharded analysis (%d files, %s)", + n_files, mode, + ) + edges = self._build_sharded(entry_points, resolver) else: logger.warning( "PyCG: %d entry points exceeds ceiling of %d — " @@ -716,14 +1039,15 @@ def build_call_graph_edges( ) return [] else: - # Small project (≤ ceiling): whole-project analysis. + # Small project (≤ ceiling): whole-project analysis. Run inside a + # symlink mini-project mirroring only the (already SKIP_DIRS-filtered) + # entry points, so PyCG's package bound covers project source alone. + # Pointing PyCG at project_dir directly would put an in-tree + # .codeanalyzer venv / site-packages *under* mod_dir, and PyCG would + # follow imports into those dependencies and explode the analysis. logger.info("PyCG: starting whole-project call graph analysis (%d files)", n_files) - try: - edges = self._run_pycg_batch( - entry_points, self.project_dir, resolver, prefix="" - ) - except PyCGExceptions.PyCGAnalysisError as exc: - raise + with _shard_symlink_root(entry_points, self.project_dir) as (root, eps): + edges = self._run_pycg_batch(eps, root, resolver, prefix="") elapsed = time.perf_counter() - t0 logger.info("✅ PyCG: %d edges in %.1fs", len(edges), elapsed) diff --git a/codeanalyzer/semantic_analysis/pycg/shard_planner.py b/codeanalyzer/semantic_analysis/pycg/shard_planner.py new file mode 100644 index 0000000..a1608dc --- /dev/null +++ b/codeanalyzer/semantic_analysis/pycg/shard_planner.py @@ -0,0 +1,401 @@ +################################################################################ +# Copyright IBM Corporation 2025 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################ + +"""Coupling-aware shard planning for PyCG call-graph construction. + +Sharding splits a project so PyCG can be run on each part independently; any +call edge whose caller and callee land in *different* shards becomes a ghost +node in both, so PyCG never resolves it. The directory + file-count heuristic +is blind to coupling and can sever heavily-interacting modules. This planner +instead partitions the **module dependency graph derived from the Jedi call +graph** (already computed at analysis level 1), so the cuts fall on the weakest +seams. + +Pipeline: + +1. **Module graph** — project the Jedi ``PyCallEdge`` list (callable→callable) + down to a weighted directed graph over *modules*. ``weight(A, B)`` is the + number of Jedi call sites from a callable in module ``A`` to one in ``B``. +2. **SCC condensation** — modules in an import/call cycle must be co-computed + (splitting them breaks both shards), so each strongly-connected component is + collapsed into one indivisible unit via Tarjan's algorithm. +3. **Community detection** — Louvain modularity over the undirected weighted + condensation groups tightly-coupled units; each community is a candidate + shard. +4. **Budget enforcement** — communities over the per-shard file budget are + re-partitioned (higher Louvain resolution, then a greedy first-fit fallback + that guarantees termination); communities under budget are agglomeratively + merged by inter-shard weight to recover cut edges and reduce shard count. + +The result is a list of shards (each a list of file paths) plus a metrics +report whose headline figure is ``cut_ratio`` — the fraction of inter-module +Jedi edge weight severed by the partition. Lower is better; it is the +estimated upper bound on PyCG edges lost to sharding. +""" + +from __future__ import annotations + +import logging +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Dict, Iterator, List, Optional, Set, Tuple + +import networkx as nx + +from codeanalyzer.schema.py_schema import PyCallable, PyCallEdge, PyClass, PyModule + +logger = logging.getLogger(__name__) + + +# ---------------------------------------------------------------------------- +# Symbol-table walking: callable / class signature -> defining file +# ---------------------------------------------------------------------------- + +def _walk_callable_sigs(c: PyCallable) -> Iterator[str]: + yield c.signature + for inner in c.inner_callables.values(): + yield from _walk_callable_sigs(inner) + for inner_cls in c.inner_classes.values(): + yield from _walk_class_sigs(inner_cls) + + +def _walk_class_sigs(cls: PyClass) -> Iterator[str]: + yield cls.signature + for method in cls.methods.values(): + yield from _walk_callable_sigs(method) + for inner in cls.inner_classes.values(): + yield from _walk_class_sigs(inner) + + +def _signature_to_file(symbol_table: Dict[str, PyModule]) -> Dict[str, str]: + """Map every callable/class signature in the project to its defining file. + + Built by walking each module's full nesting tree, recording the file that + actually defines each callable. Files are the partition unit because + ``PyModule.module_name`` is only the file *stem* (``py_file.stem``), which + collides heavily across a real project (every ``__init__.py``, ``models.py`` + …) — keying the module graph by name would collapse unrelated files into + one node and silently drop files from shards. ``file_path`` is unique. + """ + sig_to_file: Dict[str, str] = {} + for module in symbol_table.values(): + for fn in module.functions.values(): + for sig in _walk_callable_sigs(fn): + sig_to_file[sig] = module.file_path + for cls in module.classes.values(): + for sig in _walk_class_sigs(cls): + sig_to_file[sig] = module.file_path + return sig_to_file + + +# ---------------------------------------------------------------------------- +# Result type +# ---------------------------------------------------------------------------- + +@dataclass +class ShardPlan: + """Output of :func:`plan_shards`. + + ``shards`` is what the PyCG executor consumes: each inner list is the set of + project file paths to analyse together as one shard. ``module_shards`` is + the same partition expressed in module names (handy for logging/tests). + """ + + shards: List[List[str]] = field(default_factory=list) + module_shards: List[List[str]] = field(default_factory=list) + metrics: Dict[str, float] = field(default_factory=dict) + + def __str__(self) -> str: + m = self.metrics + return ( + f"ShardPlan({len(self.shards)} shards, " + f"cut_ratio={m.get('cut_ratio', 0):.3f}, " + f"max_shard_files={int(m.get('max_shard_files', 0))}, " + f"oversized={int(m.get('oversized_shards', 0))})" + ) + + +# ---------------------------------------------------------------------------- +# Module dependency graph +# ---------------------------------------------------------------------------- + +def build_module_graph( + symbol_table: Dict[str, PyModule], + jedi_edges: List[PyCallEdge], +) -> nx.DiGraph: + """Project Jedi callable→callable edges onto a weighted file DiGraph. + + Nodes are file paths (the unique, collision-free partition unit — see + :func:`_signature_to_file`); each carries a ``module_name`` attribute for + readable reporting. Every project file is a node (isolated files + included). Edge weight is the summed Jedi weight of cross-file call sites; + intra-file edges and edges touching external/library symbols (no + symbol-table entry) are dropped — they cannot influence the partition. + """ + sig_to_file = _signature_to_file(symbol_table) + + g = nx.DiGraph() + for module in symbol_table.values(): + g.add_node(module.file_path, module_name=module.module_name) + + for edge in jedi_edges: + src = sig_to_file.get(edge.source) + dst = sig_to_file.get(edge.target) + if src is None or dst is None or src == dst: + continue + if g.has_edge(src, dst): + g[src][dst]["weight"] += edge.weight + else: + g.add_edge(src, dst, weight=edge.weight) + return g + + +# ---------------------------------------------------------------------------- +# Partitioning +# ---------------------------------------------------------------------------- + +def _undirected_weighted(g: nx.DiGraph) -> nx.Graph: + """Collapse a directed weighted graph to undirected, summing both directions.""" + h = nx.Graph() + h.add_nodes_from(g.nodes()) + for u, v, w in g.edges(data="weight", default=1): + if h.has_edge(u, v): + h[u][v]["weight"] += w + else: + h.add_edge(u, v, weight=w) + return h + + +def _communities(h: nx.Graph, resolution: float, seed: int) -> List[Set[str]]: + """Weighted community detection, preferring Louvain with a graceful fallback.""" + if h.number_of_nodes() == 0: + return [] + if hasattr(nx.community, "louvain_communities"): + return [ + set(c) + for c in nx.community.louvain_communities( + h, weight="weight", resolution=resolution, seed=seed + ) + ] + # networkx < 3.0 (Python 3.9/3.10 floor): greedy modularity has no seed. + return [set(c) for c in nx.community.greedy_modularity_communities(h, weight="weight")] + + +def _greedy_bin_pack( + units: List[Tuple[frozenset, int]], budget: int +) -> List[Set[str]]: + """First-fit-decreasing pack (unit -> modules, size) into <= budget bins. + + Termination guarantee for the recursive splitter: a community that Louvain + refuses to cut still gets divided here. A single unit larger than the + budget (an atomic SCC — a real import cycle) is emitted on its own; that is + unavoidable without breaking edges. + """ + bins: List[Tuple[Set[str], int]] = [] + for members, size in sorted(units, key=lambda u: u[1], reverse=True): + placed = False + for b in bins: + if b[1] + size <= budget: + b[0].update(members) + bins[bins.index(b)] = (b[0], b[1] + size) + placed = True + break + if not placed: + bins.append((set(members), size)) + return [b[0] for b in bins] + + +def _split_oversized( + h: nx.Graph, + community: Set[str], + unit_size: Dict[str, int], + budget: int, + seed: int, + depth: int = 0, +) -> List[Set[str]]: + """Recursively partition a community whose file count exceeds the budget.""" + size = sum(unit_size[n] for n in community) + if size <= budget or len(community) == 1: + # Fits, or is a single atomic SCC unit that cannot be split further. + return [community] + + sub = h.subgraph(community) + # Escalate resolution so Louvain cuts more aggressively each level. + parts = _communities(sub, resolution=2.0 * (depth + 1), seed=seed) + parts = [p for p in parts if p] + + if len(parts) <= 1: + # Louvain could not divide it (one dense blob) — fall back to bin packing + # so the budget is still honoured. + units = [(frozenset([n]), unit_size[n]) for n in community] + return _greedy_bin_pack(units, budget) + + out: List[Set[str]] = [] + for p in parts: + out.extend(_split_oversized(h, p, unit_size, budget, seed, depth + 1)) + return out + + +def _merge_small( + shards: List[Set[str]], + h: nx.Graph, + unit_size: Dict[str, int], + budget: int, +) -> List[Set[str]]: + """Agglomeratively merge under-budget shards by inter-shard weight. + + Recovers cut edges: merging two shards turns the edges between them back + into intra-shard edges. Greedy — repeatedly merge the heaviest-coupled + pair that still fits the budget — which is enough to mop up the small + fragments Louvain leaves behind without reintroducing oversized shards. + """ + shards = [set(s) for s in shards] + sizes = [sum(unit_size[n] for n in s) for s in shards] + + def pair_weight(a: Set[str], b: Set[str]) -> float: + w = 0.0 + for u in a: + for v in h.adj[u]: + if v in b: + w += h[u][v]["weight"] + return w + + while True: + best: Optional[Tuple[int, int]] = None + best_w = 0.0 + for i in range(len(shards)): + for j in range(i + 1, len(shards)): + if sizes[i] + sizes[j] > budget: + continue + w = pair_weight(shards[i], shards[j]) + if w > best_w: + best_w, best = w, (i, j) + if best is None: + break + i, j = best + shards[i] |= shards[j] + sizes[i] += sizes[j] + del shards[j] + del sizes[j] + + # Final consolidation: first-fit pack any shards that still fit together. + # Merging zero-coupling shards (e.g. isolated leaf modules) costs no cut + # weight and avoids spawning a PyCG process per trivial file. Packing + # never severs an edge, so it is strictly safe; the budget still bounds + # per-shard PyCG divergence risk. + units = [(frozenset(s), sz) for s, sz in zip(shards, sizes)] + return _greedy_bin_pack(units, budget) + + +# ---------------------------------------------------------------------------- +# Public entry point +# ---------------------------------------------------------------------------- + +def plan_shards( + symbol_table: Dict[str, PyModule], + jedi_edges: List[PyCallEdge], + budget: int = 100, + seed: int = 42, + merge_small: bool = True, +) -> ShardPlan: + """Partition a project into coupling-aware PyCG shards. + + Args: + symbol_table: The level-1 symbol table (``file_path -> PyModule``). + jedi_edges: Jedi-provenance call edges from the level-1 call graph. + budget: Maximum number of files per shard. + seed: Determinism seed for Louvain. + merge_small: Agglomeratively merge under-budget shards to cut shard + count and recover edges. + + Returns: + A :class:`ShardPlan`. Shards never silently drop files: every project + module appears in exactly one shard (an atomic SCC larger than *budget* + yields a single oversized shard, flagged in ``metrics``). + """ + g = build_module_graph(symbol_table, jedi_edges) + total_weight = float(sum(w for _, _, w in g.edges(data="weight", default=1))) + + # SCC condensation: each strongly-connected component is one atomic unit. + condensation = nx.condensation(g) # DAG; node attr 'members' = set of modules + mapping: Dict[str, int] = condensation.graph["mapping"] # module -> scc id + unit_members: Dict[int, Set[str]] = { + scc: set(condensation.nodes[scc]["members"]) for scc in condensation.nodes + } + unit_size: Dict[int, int] = {scc: len(m) for scc, m in unit_members.items()} + + # Weighted undirected graph over SCC units. + hu = nx.Graph() + hu.add_nodes_from(condensation.nodes()) + for u, v, w in g.edges(data="weight", default=1): + su, sv = mapping[u], mapping[v] + if su == sv: + continue + if hu.has_edge(su, sv): + hu[su][sv]["weight"] += w + else: + hu.add_edge(su, sv, weight=w) + + # Community detection over units, then budget enforcement. + communities = _communities(hu, resolution=1.0, seed=seed) + unit_shards: List[Set[str]] = [] # sets of SCC ids + for community in communities: + unit_shards.extend( + _split_oversized(hu, community, unit_size, budget, seed) + ) + + if merge_small: + unit_shards = _merge_small(unit_shards, hu, unit_size, budget) + + # Expand SCC units back to file paths (graph nodes are files). + file_shards: List[List[str]] = [] + for units in unit_shards: + files: Set[str] = set() + for scc in units: + files |= unit_members[scc] + if files: + file_shards.append(sorted(files)) + + # Parallel view in module names (file stems) for human-readable reporting. + module_shards = [ + [g.nodes[f].get("module_name", f) for f in files] for files in file_shards + ] + + # Metrics: how much Jedi edge weight does this partition sever? + shard_of: Dict[str, int] = {} + for idx, files in enumerate(file_shards): + for f in files: + shard_of[f] = idx + cut_weight = 0.0 + for u, v, w in g.edges(data="weight", default=1): + if shard_of.get(u) != shard_of.get(v): + cut_weight += w + + sizes = [len(s) for s in module_shards] or [0] + metrics = { + "modules": float(g.number_of_nodes()), + "module_edges": float(g.number_of_edges()), + "total_edge_weight": total_weight, + "cut_weight": cut_weight, + "cut_ratio": (cut_weight / total_weight) if total_weight else 0.0, + "num_shards": float(len(module_shards)), + "max_shard_files": float(max(sizes)), + "oversized_shards": float(sum(1 for s in sizes if s > budget)), + } + + plan = ShardPlan(shards=file_shards, module_shards=module_shards, metrics=metrics) + logger.info("Shard planner: %s", plan) + return plan diff --git a/test/test_pycg_sharding.py b/test/test_pycg_sharding.py new file mode 100644 index 0000000..c0629e6 --- /dev/null +++ b/test/test_pycg_sharding.py @@ -0,0 +1,111 @@ +"""Tests for PyCG executor scoping: dependency exclusion and the max_iter cap. + +These drive the real PyCG wrapper, so they require ``pycg`` (a level-2 install +dependency). They are deliberately tiny (a few files) so they run fast. +""" +from pathlib import Path + +import pytest + +pytest.importorskip("PyCG") + +from codeanalyzer.semantic_analysis.pycg.pycg_analysis import ( + PyCG, + _PyCGCallableResolver, + _shard_symlink_root, +) + + +def test_max_iter_default_and_override(tmp_path): + p = PyCG(tmp_path) + assert p.max_iter == PyCG._PYCG_MAX_ITER == 50 + assert PyCG(tmp_path, max_iter=7).max_iter == 7 + assert PyCG(tmp_path, max_iter=-1).max_iter == -1 + + +def test_adaptive_decomposition_splits_runaways(tmp_path, monkeypatch): + """A shard that 'runs away' is re-decomposed until its pieces converge. + + Drives the real adaptive loop + planner, but stubs the PyCG runner with a + size threshold: shards larger than the threshold time out (runaway); smaller + ones converge and yield one synthetic edge per file. A 16-file coupled + cluster must therefore be split across rounds until every piece is small + enough, with no files lost. + """ + from codeanalyzer.schema.py_schema import PyCallEdge, PyCallable, PyModule + from codeanalyzer.semantic_analysis.pycg.pycg_analysis import _PyCGCallableResolver + + # A loosely-coupled chain of 40 modules: splittable (not one atomic cycle). + st, jedi = {}, [] + for i in range(40): + path = f"/proj/m{i}.py" + st[path] = PyModule( + file_path=path, module_name=f"m{i}", + functions={"f": PyCallable(signature=f"m{i}.f", name="f", path=path)}, + ) + if i: + jedi.append(PyCallEdge(source=f"m{i-1}.f", target=f"m{i}.f", weight=1, + provenance=["jedi"])) + + # threshold >= the decomposition floor (10) so pieces can shrink enough to converge. + pycg = PyCG(tmp_path, shard_ceiling=40) + threshold = 12 # shards with > 12 files "diverge" + rounds_seen = [] + + def fake_runner(shards): + rounds_seen.append([len(s) for s in shards]) + edges, runaways = [], [] + for files in shards: + if len(files) > threshold: + runaways.append(files) + else: + edges += [PyCallEdge(source=f, target="x", weight=1, provenance=["pycg"]) + for f in files] + return edges, runaways + + monkeypatch.setattr(pycg, "_run_fileset_shards_seq", fake_runner) + edges = pycg._build_sharded_planned(jedi, st, _PyCGCallableResolver(set())) + + assert len(rounds_seen) >= 2, "runaway shard was never decomposed" + # Every shard that was finally accepted is within the convergence threshold. + assert all(sz <= threshold for sz in rounds_seen[-1]) + # No files lost: one pycg edge per file across all 16 modules. + assert len({e.source for e in edges}) == 40 + + +def test_pycg_does_not_follow_into_in_tree_dependency(tmp_path): + """An in-tree ``.codeanalyzer`` venv under project_dir must stay a ghost. + + PyCG bounds analysis to its ``package`` directory; running inside the + symlink mini-project keeps that bound on project source only, so imports + into a bundled dependency are recorded as ghost edges but never analysed. + Regression guard for the dep-reach blowup. + """ + proj = tmp_path + app = proj / "app" + app.mkdir() + (app / "__init__.py").write_text("") + (app / "main.py").write_text("import bigdep\ndef run():\n return bigdep.work()\n") + + # A bundled dependency with many internal functions: if PyCG followed into + # it, dozens of bigdep.fN definitions/edges would appear. + dep = proj / ".codeanalyzer" / "venv" / "site-packages" / "bigdep" + dep.mkdir(parents=True) + body = "".join(f"def f{i}(x):\n return f{(i + 1) % 50}(x)\n" for i in range(50)) + body += "def work():\n return f0(1)\n" + (dep / "__init__.py").write_text(body) + + pycg = PyCG(proj) + pycg._ensure_pycg_loaded() + resolver = _PyCGCallableResolver(set()) + entry_points = [str(app / "__init__.py"), str(app / "main.py")] + with _shard_symlink_root(entry_points, proj) as (root, eps): + edges = pycg._run_pycg_batch(eps, root, resolver, prefix="") + + nodes = {n for e in edges for n in (e.source, e.target)} + # bigdep is reachable as a ghost target ... + assert any(n.startswith("bigdep") for n in nodes) + # ... but none of its internals were analysed. + assert not [n for n in nodes if n.startswith("bigdep.f")] + # and the real app edge is present. + assert any(e.source == "app.main.run" and e.target == "bigdep.work" for e in edges) diff --git a/test/test_shard_planner.py b/test/test_shard_planner.py new file mode 100644 index 0000000..fe6142c --- /dev/null +++ b/test/test_shard_planner.py @@ -0,0 +1,188 @@ +"""Unit tests for the Jedi-driven PyCG shard planner. + +These exercise the pure partitioning logic (no PyCG required): module-graph +construction, SCC atomicity, budget enforcement, and the cut-ratio metric. +""" +from typing import List + +import networkx as nx +import pytest + +from codeanalyzer.schema.py_schema import PyCallEdge, PyCallable, PyClass, PyModule +from codeanalyzer.semantic_analysis.pycg.shard_planner import ( + build_module_graph, + plan_shards, +) + + +# ---------------------------------------------------------------------------- +# Builders +# ---------------------------------------------------------------------------- + +def _module(name: str, file_path: str, func_names: List[str]) -> PyModule: + fns = { + fn: PyCallable(signature=f"{name}.{fn}", name=fn, path=file_path) + for fn in func_names + } + return PyModule(file_path=file_path, module_name=name, functions=fns) + + +def _edge(src: str, dst: str, w: int = 1) -> PyCallEdge: + return PyCallEdge(source=src, target=dst, weight=w, provenance=["jedi"]) + + +def _cut_ratio(g: nx.DiGraph, file_shards: List[List[str]]) -> float: + # Graph nodes are file paths; partitions are lists of file paths. + shard_of = {f: i for i, files in enumerate(file_shards) for f in files} + total = cut = 0.0 + for u, v, w in g.edges(data="weight", default=1): + total += w + if shard_of.get(u) != shard_of.get(v): + cut += w + return cut / total if total else 0.0 + + +# ---------------------------------------------------------------------------- +# build_module_graph +# ---------------------------------------------------------------------------- + +def test_module_graph_keys_nodes_by_file_not_name(): + # Two distinct files share the stem "models" (module_name collision); the + # graph must keep them as separate nodes keyed by path, not collapse them. + st = { + "/pkg_a/models.py": _module("models", "/pkg_a/models.py", ["f"]), + "/pkg_b/models.py": _module("models", "/pkg_b/models.py", ["h"]), + } + edges = [ + _edge("models.f", "models.h", 3), # but which file? mapped by definition + ] + g = build_module_graph(st, edges) + assert set(g.nodes) == {"/pkg_a/models.py", "/pkg_b/models.py"} + assert g.nodes["/pkg_a/models.py"]["module_name"] == "models" + + +def test_module_graph_projects_callables_to_files(): + st = { + "/a.py": _module("a", "/a.py", ["f", "g"]), + "/b.py": _module("b", "/b.py", ["h"]), + } + edges = [ + _edge("a.f", "b.h", 3), # cross-file -> kept + _edge("a.f", "a.g", 5), # intra-file -> dropped + _edge("a.g", "ext.lib.x"), # external target -> dropped + ] + g = build_module_graph(st, edges) + assert set(g.nodes) == {"/a.py", "/b.py"} + assert g.has_edge("/a.py", "/b.py") and g["/a.py"]["/b.py"]["weight"] == 3 + assert g.number_of_edges() == 1 + + +def test_isolated_files_are_nodes(): + st = {"/a.py": _module("a", "/a.py", ["f"]), "/b.py": _module("b", "/b.py", ["g"])} + g = build_module_graph(st, []) # no edges + assert set(g.nodes) == {"/a.py", "/b.py"} + assert g.number_of_edges() == 0 + + +# ---------------------------------------------------------------------------- +# plan_shards +# ---------------------------------------------------------------------------- + +def _coupled_clusters_project(): + """pkg_a <-> pkg_b heavy cross-package coupling + isolated leaves.""" + st, edges = {}, [] + for i in range(4): + st[f"/pkg_a/m{i}.py"] = _module(f"pkg_a.m{i}", f"/pkg_a/m{i}.py", ["f"]) + st[f"/pkg_b/m{i}.py"] = _module(f"pkg_b.m{i}", f"/pkg_b/m{i}.py", ["g"]) + for i in range(4): + edges.append(_edge(f"pkg_a.m{i}.f", f"pkg_b.m{i}.g", 10)) + edges.append(_edge(f"pkg_b.m{i}.g", f"pkg_a.m{i}.f", 8)) + for i in range(3): + edges.append(_edge(f"pkg_a.m{i}.f", f"pkg_a.m{i+1}.f", 2)) + return st, edges + + +def test_budget_is_respected(): + st, edges = _coupled_clusters_project() + plan = plan_shards(st, edges, budget=4) + assert plan.metrics["max_shard_files"] <= 4 + assert plan.metrics["oversized_shards"] == 0 + + +def test_every_module_assigned_exactly_once(): + st, edges = _coupled_clusters_project() + plan = plan_shards(st, edges, budget=4) + assigned = [m for shard in plan.module_shards for m in shard] + assert sorted(assigned) == sorted(m.module_name for m in st.values()) + assert len(assigned) == len(set(assigned)) # no module duplicated + + +def test_beats_naive_per_package_cut_ratio(): + st, edges = _coupled_clusters_project() + g = build_module_graph(st, edges) + # Naive baseline: one shard per top-level directory (e.g. /pkg_a/...). + naive = {} + for m in st.values(): + top = m.file_path.split("/")[1] + naive.setdefault(top, []).append(m.file_path) + naive_ratio = _cut_ratio(g, list(naive.values())) + + plan = plan_shards(st, edges, budget=4) + assert plan.metrics["cut_ratio"] < naive_ratio + + +def test_import_cycle_is_never_split(): + # util.x <-> helpers.y form a cross-package cycle; must co-locate even + # though they live in different top-level packages. + st = { + "/util/x.py": _module("util.x", "/util/x.py", ["a"]), + "/helpers/y.py": _module("helpers.y", "/helpers/y.py", ["b"]), + } + edges = [_edge("util.x.a", "helpers.y.b", 5), _edge("helpers.y.b", "util.x.a", 5)] + plan = plan_shards(st, edges, budget=10) + shard_with_util = next(s for s in plan.module_shards if "util.x" in s) + assert "helpers.y" in shard_with_util + + +def test_oversized_atomic_cycle_is_flagged_not_dropped(): + # A single import cycle of 6 modules with a budget of 3 cannot be split + # without breaking edges; it must survive as one oversized shard. + st, edges = {}, [] + names = [f"cyc.m{i}" for i in range(6)] + for i, n in enumerate(names): + st[f"/cyc/m{i}.py"] = _module(n, f"/cyc/m{i}.py", ["f"]) + for i in range(6): # ring -> one big SCC + edges.append(_edge(f"{names[i]}.f", f"{names[(i + 1) % 6]}.f", 1)) + plan = plan_shards(st, edges, budget=3) + assert plan.metrics["oversized_shards"] >= 1 + assigned = [m for shard in plan.module_shards for m in shard] + assert sorted(assigned) == sorted(names) # nothing dropped + + +def test_determinism(): + st, edges = _coupled_clusters_project() + a = plan_shards(st, edges, budget=4) + b = plan_shards(st, edges, budget=4) + norm = lambda p: sorted(sorted(s) for s in p.module_shards) + assert norm(a) == norm(b) + + +def test_no_file_dropped_on_stem_collision(): + # Regression: module_name is only the file stem, so many files collide on + # name (every __init__.py, models.py, ...). Keying by name would drop all + # but one per stem. Every FILE must land in exactly one shard. + st, edges = {}, [] + for pkg in ("a", "b", "c", "d"): + for stem in ("__init__", "models", "views"): + path = f"/{pkg}/{stem}.py" + st[path] = _module(stem, path, ["f"]) + plan = plan_shards(st, edges, budget=5) + assigned = sorted(f for shard in plan.shards for f in shard) + assert assigned == sorted(st.keys()) # all 12 files present + assert len(assigned) == len(set(assigned)) # none duplicated + + +def test_empty_project(): + plan = plan_shards({}, [], budget=4) + assert plan.shards == [] + assert plan.metrics["cut_ratio"] == 0.0