Add module expressions

CHANGELOG.rst

@@ -5,6 +5,7 @@ Changelog
 Unreleased
 ----------
 
+* Added basic functions for working with module expressions: `find_matching_modules` and `find_matching_direct_imports`.
 * Added `as_packages` option to the `find_shortest_chain` method.
 
 3.6 (2025-02-07)

docs/usage.rst

@@ -121,6 +121,16 @@ Methods for analysing the module tree
     :raises: ``ValueError`` if the module is a squashed module, as by definition it represents both itself and all
       of its descendants.
 
+.. py:function:: ImportGraph.find_matching_modules(expression)
+
+    Find all modules matching the passed expression (see :ref:`module_expressions`).
+
+    :param str expression: A module expression used for matching.
+    :return: A set of module names matching the expression.
+    :rtype: A set of strings.
+    :raises: ``grimp.exceptions.InvalidModuleExpression`` if the module expression is invalid.
+
+
 Methods for analysing direct imports
 ------------------------------------
 
@@ -184,6 +194,27 @@ Methods for analysing direct imports
         So if a module is imported twice from the same module, it will only be counted once.
     :rtype: Integer.
 
+.. py:function:: ImportGraph.find_matching_direct_imports(importer_expression, imported_expression)
+
+    Find all direct imports matching the passed expressions (see :ref:`module_expressions`).
+
+    The imports are returned are in the following form::
+
+        [
+            {
+                'importer': 'mypackage.importer',
+                'imported': 'mypackage.imported',
+            },
+            # (additional imports here)
+        ]
+
+    :param str importer_expression: A module expression used for matching importing modules.
+    :param str imported_expression: A module expression used for matching imported modules.
+    :return: An ordered list of direct imports matching the expressions (ordered alphabetically).
+    :rtype: List of dictionaries with the structure shown above. If you want to use type annotations, you may use the
+        ``grimp.Import`` TypedDict for each dictionary.
+    :raises: ``grimp.exceptions.InvalidModuleExpression`` if either of the module expressions is invalid.
+
 Methods for analysing import chains
 -----------------------------------
 
@@ -517,5 +548,25 @@ Methods for manipulating the graph
     :param str module: The name of a module, for example ``'mypackage.foo'``.
     :return: bool
 
+.. _module_expressions:
+
+Module expressions
+------------------
+
+  A module expression is used to refer to sets of modules.
+
+  - ``*`` stands in for a module name, without including subpackages.
+  - ``**`` includes subpackages too.
+
+  Examples:
+
+  - ``mypackage.foo``:  matches ``mypackage.foo`` exactly.
+  - ``mypackage.*``:  matches ``mypackage.foo`` but not ``mypackage.foo.bar``.
+  - ``mypackage.*.baz``: matches ``mypackage.foo.baz`` but not ``mypackage.foo.bar.baz``.
+  - ``mypackage.*.*``: matches ``mypackage.foo.bar`` and ``mypackage.foobar.baz``.
+  - ``mypackage.**``: matches ``mypackage.foo.bar`` and ``mypackage.foo.bar.baz``.
+  - ``mypackage.**.qux``: matches ``mypackage.foo.bar.qux`` and ``mypackage.foo.bar.baz.qux``.
+  - ``mypackage.foo*``: is not a valid expression. (The wildcard must replace a whole module name.)
+
 .. _namespace packages: https://docs.python.org/3/glossary.html#term-namespace-package
 .. _namespace portion: https://docs.python.org/3/glossary.html#term-portion

rust/Cargo.toml

@@ -20,6 +20,7 @@ itertools = "0.14.0"
 tap = "1.0.1"
 rustc-hash = "2.1.0"
 indexmap = "2.7.1"
+regex = "1.11.1"
 
 [dependencies.pyo3]
 version = "0.23.4"
@@ -29,4 +30,5 @@ extension-module = ["pyo3/extension-module"]
 default = ["extension-module"]
 
 [dev-dependencies]
+parameterized = "2.0.0"
 serde_json = "1.0.137"

rust/src/errors.rs

@@ -1,4 +1,4 @@
-use crate::exceptions::{ModuleNotPresent, NoSuchContainer};
+use crate::exceptions::{InvalidModuleExpression, ModuleNotPresent, NoSuchContainer};
 use pyo3::exceptions::PyValueError;
 use pyo3::PyErr;
 use thiserror::Error;
@@ -13,6 +13,9 @@ pub enum GrimpError {
 
     #[error("Modules have shared descendants.")]
     SharedDescendants,
+
+    #[error("{0} is not a valid module expression.")]
+    InvalidModuleExpression(String),
 }
 
 pub type GrimpResult<T> = Result<T, GrimpError>;
@@ -24,6 +27,9 @@ impl From<GrimpError> for PyErr {
             GrimpError::ModuleNotPresent(_) => ModuleNotPresent::new_err(value.to_string()),
             GrimpError::NoSuchContainer(_) => NoSuchContainer::new_err(value.to_string()),
             GrimpError::SharedDescendants => PyValueError::new_err(value.to_string()),
+            GrimpError::InvalidModuleExpression(_) => {
+                InvalidModuleExpression::new_err(value.to_string())
+            }
         }
     }
 }

rust/src/exceptions.rs

@@ -2,3 +2,8 @@ use pyo3::create_exception;
 
 create_exception!(_rustgrimp, ModuleNotPresent, pyo3::exceptions::PyException);
 create_exception!(_rustgrimp, NoSuchContainer, pyo3::exceptions::PyException);
+create_exception!(
+    _rustgrimp,
+    InvalidModuleExpression,
+    pyo3::exceptions::PyException
+);

rust/src/graph/direct_import_queries.rs

@@ -1,8 +1,9 @@
 use crate::errors::{GrimpError, GrimpResult};
 use crate::graph::{
     ExtendWithDescendants, Graph, ImportDetails, ModuleToken, EMPTY_IMPORT_DETAILS,
-    EMPTY_MODULE_TOKENS,
+    EMPTY_MODULE_TOKENS, MODULE_NAMES,
 };
+use crate::module_expressions::ModuleExpression;
 use rustc_hash::FxHashSet;
 
 impl Graph {
@@ -54,4 +55,72 @@ impl Graph {
             None => &EMPTY_IMPORT_DETAILS,
         }
     }
+
+    pub fn find_matching_direct_imports(
+        &self,
+        importer_expression: &ModuleExpression,
+        imported_expression: &ModuleExpression,
+    ) -> FxHashSet<(ModuleToken, ModuleToken)> {
+        let interner = MODULE_NAMES.read().unwrap();
+        self.imports
+            .iter()
+            .flat_map(|(importer, imports)| {
+                imports
+                    .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)
+            })
+            .collect()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_find_matching_direct_imports() {
+        let mut graph = Graph::default();
+
+        let _ = 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 _ = graph.get_or_add_module("pkg.food").token;
+        let chicken = graph.get_or_add_module("pkg.food.chicken").token;
+        let fish = graph.get_or_add_module("pkg.food.fish").token;
+        let _ = graph.get_or_add_module("pkg.colors").token;
+        let golden = graph.get_or_add_module("pkg.colors.golden").token;
+        let ginger = graph.get_or_add_module("pkg.colors.ginger").token;
+        let _ = graph.get_or_add_module("pkg.shops").token;
+        let tesco = graph.get_or_add_module("pkg.shops.tesco").token;
+        let coop = graph.get_or_add_module("pkg.shops.coop").token;
+
+        // Should match
+        graph.add_import(dog, chicken);
+        graph.add_import(cat, fish);
+        // Should not match: Imported does not match
+        graph.add_import(dog, golden);
+        graph.add_import(cat, ginger);
+        // Should not match: Importer does not match
+        graph.add_import(tesco, chicken);
+        graph.add_import(coop, fish);
+
+        let importer_expression = "pkg.animals.*".parse().unwrap();
+        let imported_expression = "pkg.food.*".parse().unwrap();
+        let matching_imports =
+            graph.find_matching_direct_imports(&importer_expression, &imported_expression);
+
+        assert_eq!(
+            matching_imports,
+            FxHashSet::from_iter([(dog, chicken), (cat, fish)])
+        );
+    }
 }