From 65a1cfe024d1e4269b7291ee78d90e2841368d91 Mon Sep 17 00:00:00 2001 From: Peter Byfield Date: Tue, 10 Nov 2026 17:38:50 +0100 Subject: [PATCH] Add as_package parameter to direct import methods find_modules_directly_imported_by and find_modules_that_directly_import can now treat the supplied module as a package, as find_upstream_modules and find_downstream_modules already do. The result includes the modules that the package's descendants import (or that import them), and excludes modules within the package. direct_import_exists already had as_packages. Fixes #25. Assistant-model: Claude Code --- CHANGELOG.rst | 7 +- docs/usage.rst | 12 +- rust/src/graph/direct_import_queries.rs | 134 ++++++++++++++++-- rust/src/graph/graph_manipulation.rs | 16 +-- rust/src/graph/mod.rs | 18 ++- src/grimp/application/graph.py | 22 ++- .../application/graph/test_direct_imports.py | 50 ++++++- 7 files changed, 220 insertions(+), 39 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 33d2f5c2..b436834b 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -5,9 +5,10 @@ Changelog latest ------ -* Speed up ``ImportGraph.find_matching_modules`` for expressions that don't contain wildcards. -* Speed up ``ImportGraph.find_matching_direct_imports`` by only looking at the imports of the modules that match - one of the expressions. +* Speed up ``find_matching_modules`` for expressions that don't contain wildcards. +* Speed up ``find_matching_direct_imports``. +* Add ``as_package`` parameter to ``find_modules_directly_imported_by`` and + ``find_modules_that_directly_import``. 3.19 (2026-11-06) ----------------- diff --git a/docs/usage.rst b/docs/usage.rst index df1c516e..447dd96f 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -145,15 +145,23 @@ Methods for analysing direct imports :return: Whether or not the importer directly imports the imported module. :rtype: ``True`` or ``False``. -.. py:function:: ImportGraph.find_modules_directly_imported_by(module) +.. py:function:: ImportGraph.find_modules_directly_imported_by(module, as_package=False) :param str module: A module name. + :param bool as_package: Whether or not to treat the supplied module as an individual module, + or as a package (i.e. including any descendants, if there are any). If + treating it as a package, the result will include modules *external* to the + package, and won't include modules within it. :return: Set of all modules in the graph are imported by the supplied module. :rtype: A set of strings. -.. py:function:: ImportGraph.find_modules_that_directly_import(module) +.. py:function:: ImportGraph.find_modules_that_directly_import(module, as_package=False) :param str module: A module name. + :param bool as_package: Whether or not to treat the supplied module as an individual module, + or as a package (i.e. including any descendants, if there are any). If + treating it as a package, the result will include modules *external* to the + package, and won't include modules within it. :return: Set of all modules in the graph that directly import the supplied module. :rtype: A set of strings. diff --git a/rust/src/graph/direct_import_queries.rs b/rust/src/graph/direct_import_queries.rs index 1785a717..6a617604 100644 --- a/rust/src/graph/direct_import_queries.rs +++ b/rust/src/graph/direct_import_queries.rs @@ -67,14 +67,58 @@ impl Graph { Ok(all_imports) } - pub fn modules_directly_imported_by(&self, importer: ModuleToken) -> &FxHashSet { - self.imports.get(importer).unwrap_or(&EMPTY_MODULE_TOKENS) + /// Find the modules directly imported by the module. + /// + /// If `as_package` is true, the module is treated as a package: the result is the modules + /// directly imported by it or any of its descendants, excluding modules within the package. + pub fn modules_directly_imported_by( + &self, + importer: ModuleToken, + as_package: bool, + ) -> FxHashSet { + let mut importers: FxHashSet<_> = importer.into(); + if as_package { + importers.extend_with_descendants(self); + } + + let mut imported_modules: FxHashSet<_> = importers + .iter() + .flat_map(|importer| self.imports.get(*importer).unwrap_or(&EMPTY_MODULE_TOKENS)) + .cloned() + .collect(); + if as_package { + imported_modules.retain(|imported| !importers.contains(imported)); + } + imported_modules } - pub fn modules_that_directly_import(&self, imported: ModuleToken) -> &FxHashSet { - self.reverse_imports - .get(imported) - .unwrap_or(&EMPTY_MODULE_TOKENS) + /// Find the modules that directly import the module. + /// + /// If `as_package` is true, the module is treated as a package: the result is the modules + /// that directly import it or any of its descendants, excluding modules within the package. + pub fn modules_that_directly_import( + &self, + imported: ModuleToken, + as_package: bool, + ) -> FxHashSet { + let mut importeds: FxHashSet<_> = imported.into(); + if as_package { + importeds.extend_with_descendants(self); + } + + let mut importers: FxHashSet<_> = importeds + .iter() + .flat_map(|imported| { + self.reverse_imports + .get(*imported) + .unwrap_or(&EMPTY_MODULE_TOKENS) + }) + .cloned() + .collect(); + if as_package { + importers.retain(|importer| !importeds.contains(importer)); + } + importers } pub fn get_import_details( @@ -105,10 +149,10 @@ impl Graph { importers .iter() .flat_map(|importer| { - self.modules_directly_imported_by(*importer) - .iter() - .filter(|imported| importeds.contains(*imported)) - .map(move |imported| (*importer, *imported)) + self.modules_directly_imported_by(*importer, false) + .into_iter() + .filter(|imported| importeds.contains(imported)) + .map(move |imported| (*importer, imported)) }) .collect() } @@ -118,6 +162,76 @@ impl Graph { mod test { use super::*; + #[test] + fn test_modules_directly_imported_by() { + let mut graph = Graph::default(); + + let animals = graph.get_or_add_module("pkg.animals").token; + let dog = graph.get_or_add_module("pkg.animals.dog").token; + let cat = graph.get_or_add_module("pkg.animals.cat").token; + let food = graph.get_or_add_module("pkg.food").token; + let chicken = graph.get_or_add_module("pkg.food.chicken").token; + let tesco = graph.get_or_add_module("pkg.shops.tesco").token; + + graph.add_import(dog, chicken); + graph.add_import(cat, food); + // An import within the package. + graph.add_import(dog, cat); + // An import of the package, not by it. + graph.add_import(tesco, animals); + + // Without as_package, only the module itself is looked at. + assert_eq!( + graph.modules_directly_imported_by(animals, false), + FxHashSet::default() + ); + assert_eq!( + graph.modules_directly_imported_by(dog, false), + FxHashSet::from_iter([chicken, cat]) + ); + // With as_package, descendants are looked at too, but modules within the package + // are excluded. + assert_eq!( + graph.modules_directly_imported_by(animals, true), + FxHashSet::from_iter([chicken, food]) + ); + } + + #[test] + fn test_modules_that_directly_import() { + let mut graph = Graph::default(); + + let animals = graph.get_or_add_module("pkg.animals").token; + let dog = graph.get_or_add_module("pkg.animals.dog").token; + let cat = graph.get_or_add_module("pkg.animals.cat").token; + let chicken = graph.get_or_add_module("pkg.food.chicken").token; + let shops = graph.get_or_add_module("pkg.shops").token; + let tesco = graph.get_or_add_module("pkg.shops.tesco").token; + + graph.add_import(tesco, animals); + graph.add_import(shops, cat); + // An import within the package. + graph.add_import(dog, cat); + // An import by the package, not of it. + graph.add_import(dog, chicken); + + // Without as_package, only the module itself is looked at. + assert_eq!( + graph.modules_that_directly_import(animals, false), + FxHashSet::from_iter([tesco]) + ); + assert_eq!( + graph.modules_that_directly_import(cat, false), + FxHashSet::from_iter([shops, dog]) + ); + // With as_package, descendants are looked at too, but modules within the package + // are excluded. + assert_eq!( + graph.modules_that_directly_import(animals, true), + FxHashSet::from_iter([tesco, shops]) + ); + } + #[test] fn test_find_matching_direct_imports() { let mut graph = Graph::default(); diff --git a/rust/src/graph/graph_manipulation.rs b/rust/src/graph/graph_manipulation.rs index b6c762c3..a780ab3d 100644 --- a/rust/src/graph/graph_manipulation.rs +++ b/rust/src/graph/graph_manipulation.rs @@ -91,10 +91,10 @@ impl Graph { self.module_children.remove(module); // Update imports. - for imported in self.modules_directly_imported_by(module).clone() { + for imported in self.modules_directly_imported_by(module, false) { self.remove_import(module, imported); } - for importer in self.modules_that_directly_import(module).clone() { + for importer in self.modules_that_directly_import(module, false) { self.remove_import(importer, module); } self.imports.remove(module); @@ -164,19 +164,11 @@ impl Graph { let modules_imported_by_descendants: FxHashSet<_> = descendants .iter() - .flat_map(|descendant| { - self.modules_directly_imported_by(*descendant) - .iter() - .cloned() - }) + .flat_map(|descendant| self.modules_directly_imported_by(*descendant, false)) .collect(); let modules_that_import_descendants: FxHashSet<_> = descendants .iter() - .flat_map(|descendant| { - self.modules_that_directly_import(*descendant) - .iter() - .cloned() - }) + .flat_map(|descendant| self.modules_that_directly_import(*descendant, false)) .collect(); // Add descendants and imports to parent module. diff --git a/rust/src/graph/mod.rs b/rust/src/graph/mod.rs index cf745368..40ff48cb 100644 --- a/rust/src/graph/mod.rs +++ b/rust/src/graph/mod.rs @@ -394,11 +394,16 @@ impl GraphWrapper { .direct_import_exists(importer, imported, as_packages)?) } - pub fn find_modules_directly_imported_by(&self, module: &str) -> PyResult> { + #[pyo3(signature = (module, as_package=false))] + pub fn find_modules_directly_imported_by( + &self, + module: &str, + as_package: bool, + ) -> PyResult> { let module = self.get_visible_module_by_name(module)?.token(); Ok(self ._graph - .modules_directly_imported_by(module) + .modules_directly_imported_by(module, as_package) .iter() .into_module_iterator(&self._graph) .visible() @@ -406,11 +411,16 @@ impl GraphWrapper { .collect()) } - pub fn find_modules_that_directly_import(&self, module: &str) -> PyResult> { + #[pyo3(signature = (module, as_package=false))] + pub fn find_modules_that_directly_import( + &self, + module: &str, + as_package: bool, + ) -> PyResult> { let module = self.get_visible_module_by_name(module)?.token(); Ok(self ._graph - .modules_that_directly_import(module) + .modules_that_directly_import(module, as_package) .iter() .into_module_iterator(&self._graph) .visible() diff --git a/src/grimp/application/graph.py b/src/grimp/application/graph.py index 7e53481a..3c679621 100644 --- a/src/grimp/application/graph.py +++ b/src/grimp/application/graph.py @@ -223,27 +223,41 @@ def direct_import_exists( importer=importer, imported=imported, as_packages=as_packages ) - def find_modules_directly_imported_by(self, module: str) -> set[str]: + def find_modules_directly_imported_by(self, module: str, as_package: bool = False) -> set[str]: """ Find all modules that are directly imported by the supplied module. + Args: + module: The name of the importing module. + as_package: Whether or not to treat the supplied module as an individual module, + or as a package (i.e. including any descendants, if there are any). If + treating it as a package, the result will include modules *external* to + the package, and won't include modules within it. + If the module is not present in the graph, grimp.exceptions.ModuleNotPresent will be raised. """ if not self._rustgraph.contains_module(module): msg = f'"{module}" not present in the graph.' raise ModuleNotPresent(msg) - return self._rustgraph.find_modules_directly_imported_by(module) + return self._rustgraph.find_modules_directly_imported_by(module, as_package) - def find_modules_that_directly_import(self, module: str) -> set[str]: + def find_modules_that_directly_import(self, module: str, as_package: bool = False) -> set[str]: """ Find all modules that directly import the supplied module. + Args: + module: The name of the imported module. + as_package: Whether or not to treat the supplied module as an individual module, + or as a package (i.e. including any descendants, if there are any). If + treating it as a package, the result will include modules *external* to + the package, and won't include modules within it. + If the module is not present in the graph, grimp.exceptions.ModuleNotPresent will be raised. """ if not self._rustgraph.contains_module(module): msg = f'"{module}" not present in the graph.' raise ModuleNotPresent(msg) - return self._rustgraph.find_modules_that_directly_import(module) + return self._rustgraph.find_modules_that_directly_import(module, as_package) def get_import_details(self, *, importer: str, imported: str) -> list[DetailedImport]: """ diff --git a/tests/unit/application/graph/test_direct_imports.py b/tests/unit/application/graph/test_direct_imports.py index 71270363..ac65fabd 100644 --- a/tests/unit/application/graph/test_direct_imports.py +++ b/tests/unit/application/graph/test_direct_imports.py @@ -34,20 +34,62 @@ def test_find_modules_that_directly_import(): assert {a, f} == graph.find_modules_that_directly_import("bar") -def test_find_modules_directly_imported_by_raises_module_not_present(): +def test_find_modules_directly_imported_by_as_package(): + graph = ImportGraph() + graph.add_import(importer="foo", imported="bar") + graph.add_import(importer="foo.one", imported="baz") + graph.add_import(importer="foo.one.alpha", imported="baz.one") + # Imports within the package. + graph.add_import(importer="foo.one", imported="foo.two") + graph.add_import(importer="foo.two", imported="foo") + # Imports of the package are not imports by the package. + graph.add_import(importer="qux", imported="foo") + + assert {"bar"} == graph.find_modules_directly_imported_by("foo") + assert {"bar", "baz", "baz.one"} == graph.find_modules_directly_imported_by( + "foo", as_package=True + ) + assert {"baz", "baz.one", "foo.two"} == graph.find_modules_directly_imported_by( + "foo.one", as_package=True + ) + + +def test_find_modules_that_directly_import_as_package(): + graph = ImportGraph() + graph.add_import(importer="bar", imported="foo") + graph.add_import(importer="baz", imported="foo.one") + graph.add_import(importer="baz.one", imported="foo.one.alpha") + # Imports within the package. + graph.add_import(importer="foo.two", imported="foo.one") + graph.add_import(importer="foo", imported="foo.two") + # Imports by the package are not imports of the package. + graph.add_import(importer="foo", imported="qux") + + assert {"bar"} == graph.find_modules_that_directly_import("foo") + assert {"bar", "baz", "baz.one"} == graph.find_modules_that_directly_import( + "foo", as_package=True + ) + assert {"baz", "baz.one", "foo.two"} == graph.find_modules_that_directly_import( + "foo.one", as_package=True + ) + + +@pytest.mark.parametrize("as_package", (False, True)) +def test_find_modules_directly_imported_by_raises_module_not_present(as_package): graph = ImportGraph() graph.add_module("foo") with pytest.raises(ModuleNotPresent, match=re.escape('"bar" not present in the graph.')): - graph.find_modules_directly_imported_by("bar") + graph.find_modules_directly_imported_by("bar", as_package=as_package) -def test_find_modules_that_directly_import_raises_module_not_present(): +@pytest.mark.parametrize("as_package", (False, True)) +def test_find_modules_that_directly_import_raises_module_not_present(as_package): graph = ImportGraph() graph.add_module("foo") with pytest.raises(ModuleNotPresent, match=re.escape('"bar" not present in the graph.')): - graph.find_modules_that_directly_import("bar") + graph.find_modules_that_directly_import("bar", as_package=as_package) @pytest.mark.parametrize(