diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ea3e8ab7..d2484c47 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -5,7 +5,9 @@ Changelog latest ------ -* Speed up ``ImportGraph.find_matching_modules`` for expressions that don't contain wildcards. +* 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 b3534d23..d7a50ceb 100644 --- a/rust/src/graph/direct_import_queries.rs +++ b/rust/src/graph/direct_import_queries.rs @@ -1,6 +1,6 @@ use crate::errors::{GrimpError, GrimpResult}; use crate::graph::{ - EMPTY_IMPORT_DETAILS, EMPTY_MODULE_TOKENS, ExtendWithDescendants, Graph, MODULE_NAMES, + EMPTY_IMPORT_DETAILS, EMPTY_MODULE_TOKENS, ExtendWithDescendants, Graph, ModuleIterator, ModuleToken, PyImportDetails, }; use crate::module_expressions::ModuleExpression; @@ -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( @@ -93,22 +137,22 @@ impl Graph { importer_expression: &ModuleExpression, imported_expression: &ModuleExpression, ) -> FxHashSet<(ModuleToken, ModuleToken)> { - let interner = MODULE_NAMES.read().unwrap(); - self.imports + let importers: FxHashSet = self + .find_matching_modules(importer_expression) + .tokens() + .collect(); + let importeds: FxHashSet = self + .find_matching_modules(imported_expression) + .tokens() + .collect(); + + importers .iter() - .flat_map(|(importer, imports)| { - imports + .flat_map(|importer| { + self.modules_directly_imported_by(*importer) .iter() - .cloned() - .map(move |imported| (importer, imported)) - }) - .filter(|(importer, imported)| { - let importer = self.get_module(*importer).unwrap(); - let importer_name = interner.resolve(importer.interned_name).unwrap(); - let imported = self.get_module(*imported).unwrap(); - let imported_name = interner.resolve(imported.interned_name).unwrap(); - importer_expression.is_match(importer_name) - && imported_expression.is_match(imported_name) + .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(); @@ -155,4 +269,60 @@ mod test { FxHashSet::from_iter([(dog, chicken), (cat, fish)]) ); } + + #[test] + fn test_find_matching_direct_imports_with_and_without_wildcards() { + let mut graph = Graph::default(); + + 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 fish = graph.get_or_add_module("pkg.food.fish").token; + let golden = graph.get_or_add_module("pkg.colors.golden").token; + let tesco = graph.get_or_add_module("pkg.shops.tesco").token; + graph.get_or_add_module("pkg.shops.coop"); + + graph.add_import(dog, chicken); + graph.add_import(dog, golden); + graph.add_import(cat, fish); + graph.add_import(tesco, chicken); + + let find = |importer_expression: &str, imported_expression: &str| { + graph.find_matching_direct_imports( + &importer_expression.parse().unwrap(), + &imported_expression.parse().unwrap(), + ) + }; + + // Neither side has wildcards. + assert_eq!( + find("pkg.animals.dog", "pkg.food.chicken"), + FxHashSet::from_iter([(dog, chicken)]) + ); + // Only the importer is a single module. + assert_eq!( + find("pkg.animals.dog", "pkg.*.*"), + FxHashSet::from_iter([(dog, chicken), (dog, golden)]) + ); + // Only the imported is a single module. + assert_eq!( + find("pkg.*.*", "pkg.food.chicken"), + FxHashSet::from_iter([(dog, chicken), (tesco, chicken)]) + ); + // Both sides have wildcards. + assert_eq!( + find("pkg.shops.*", "pkg.food.*"), + FxHashSet::from_iter([(tesco, chicken)]) + ); + // A module with no imports. + assert!(find("pkg.shops.coop", "pkg.*.*").is_empty()); + // A module that has no importers. + assert!(find("pkg.*.*", "pkg.shops.coop").is_empty()); + // Modules that aren't in the graph. + assert!(find("pkg.animals.missing", "pkg.*.*").is_empty()); + assert!(find("pkg.*.*", "pkg.food.missing").is_empty()); + assert!(find("pkg.animals.missing", "pkg.food.missing").is_empty()); + // Wildcards that match nothing. + assert!(find("pkg.nothing.*", "pkg.*.*").is_empty()); + } } 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 d7d724a3..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( @@ -344,6 +386,58 @@ def test_finds_matching_direct_imports(self, import_line_number, import_line_con {"importer": "pkg.animals.dog", "imported": "pkg.food.chicken"}, ] + @pytest.mark.parametrize( + "expression, expected_imports", + [ + # Neither side has wildcards. + ( + "pkg.animals.dog -> pkg.food.chicken", + [("pkg.animals.dog", "pkg.food.chicken")], + ), + # Only the importer has no wildcards. + ( + "pkg.animals.dog -> pkg.*.*", + [ + ("pkg.animals.dog", "pkg.colors.golden"), + ("pkg.animals.dog", "pkg.food.chicken"), + ], + ), + # Only the imported has no wildcards. + ( + "pkg.*.* -> pkg.food.chicken", + [("pkg.animals.dog", "pkg.food.chicken"), ("pkg.shops.tesco", "pkg.food.chicken")], + ), + # Both sides have wildcards. + ( + "pkg.shops.* -> pkg.food.*", + [("pkg.shops.tesco", "pkg.food.chicken")], + ), + # A module that doesn't import anything. + ("pkg.shops.coop -> pkg.*.*", []), + # A module that isn't imported by anything. + ("pkg.*.* -> pkg.shops.coop", []), + # Modules that aren't in the graph. + ("pkg.animals.missing -> pkg.*.*", []), + ("pkg.*.* -> pkg.food.missing", []), + ("pkg.animals.missing -> pkg.food.missing", []), + # Wildcards that match nothing. + ("pkg.nothing.* -> pkg.*.*", []), + ], + ) + def test_finds_matching_direct_imports_with_and_without_wildcards( + self, expression, expected_imports + ): + graph = ImportGraph() + graph.add_module("pkg.shops.coop") + graph.add_import(importer="pkg.animals.dog", imported="pkg.food.chicken") + graph.add_import(importer="pkg.animals.dog", imported="pkg.colors.golden") + graph.add_import(importer="pkg.animals.cat", imported="pkg.food.fish") + graph.add_import(importer="pkg.shops.tesco", imported="pkg.food.chicken") + + assert graph.find_matching_direct_imports(expression) == [ + {"importer": importer, "imported": imported} for importer, imported in expected_imports + ] + def test_deduplicates_imports(self): graph = ImportGraph() graph.add_import(