Implement method without_import

Peter554 · Nov 26, 2023 · b3f6699 · b3f6699
1 parent f81d056
commit b3f6699
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 1 deletion.
diff --git a/src/import_graph/errors.rs b/src/import_graph/errors.rs
@@ -15,4 +15,7 @@ pub enum Error {
 
     #[error("module not found")]
     ModuleNotFound(String),
+
+    #[error("import not found")]
+    ImportNotFound(String, String),
 }
diff --git a/src/import_graph/graph.rs b/src/import_graph/graph.rs
@@ -8,7 +8,7 @@ use super::import_discovery;
 use super::indexing;
 use super::package_discovery::{Module, Package};
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct ImportGraph {
     pub(super) packages_by_pypath: indexing::PackagesByPypath,
     pub(super) modules_by_pypath: indexing::ModulesByPypath,
@@ -259,4 +259,26 @@ impl ImportGraph {
         );
         shortest_path.map(|shortest_path| shortest_path.iter().map(|m| m.pypath.clone()).collect())
     }
+
+    pub fn without_import(&self, from_module: &str, to_module: &str) -> Result<ImportGraph> {
+        let from_module = match self.modules_by_pypath.get(from_module) {
+            Some(module) => Arc::clone(module),
+            None => Err(Error::ModuleNotFound(from_module.to_string()))?,
+        };
+        let to_module = match self.modules_by_pypath.get(to_module) {
+            Some(module) => Arc::clone(module),
+            None => Err(Error::ModuleNotFound(to_module.to_string()))?,
+        };
+        let mut import_graph = self.clone();
+        let imports = import_graph.imports.get_mut(&from_module).unwrap();
+        if imports.contains(&to_module) {
+            imports.remove(&to_module);
+            Ok(import_graph)
+        } else {
+            Err(Error::ImportNotFound(
+                from_module.pypath.clone(),
+                to_module.pypath.clone(),
+            ))?
+        }
+    }
 }
diff --git a/tests/test_import_graph.rs b/tests/test_import_graph.rs
@@ -817,3 +817,45 @@ fn test_path_exists_packages() {
         .path_exists("somesillypackage.child2", "somesillypackage.child1")
         .unwrap(),);
 }
+
+#[test]
+fn test_without_import() {
+    let root_package_path = Path::new("./testpackages/somesillypackage");
+
+    let import_graph = ImportGraphBuilder::new(root_package_path).build().unwrap();
+    assert_eq!(
+        import_graph
+            .shortest_path("somesillypackage.a", "somesillypackage.e")
+            .unwrap()
+            .unwrap(),
+        vec![
+            "somesillypackage.a",
+            "somesillypackage.c",
+            "somesillypackage.e"
+        ],
+    );
+
+    let import_graph = import_graph
+        .without_import("somesillypackage.a", "somesillypackage.c")
+        .unwrap();
+    assert_eq!(
+        import_graph
+            .shortest_path("somesillypackage.a", "somesillypackage.e")
+            .unwrap()
+            .unwrap(),
+        vec![
+            "somesillypackage.a",
+            "somesillypackage.b",
+            "somesillypackage.c",
+            "somesillypackage.e"
+        ],
+    );
+
+    let import_graph = import_graph
+        .without_import("somesillypackage.a", "somesillypackage.b")
+        .unwrap();
+    assert!(import_graph
+        .shortest_path("somesillypackage.a", "somesillypackage.e")
+        .unwrap()
+        .is_none());
+}