Add tests for graph.add_module

Author: Peter554

rust/src/graph.rs

@@ -1,4 +1,4 @@
-use crate::hierarchy::ModuleHierarchy;
+use crate::hierarchy::{ModuleHierarchy, ModuleIterator};
 use crate::imports::{ImportDetails, ModuleImports};
 use anyhow::Result;
 use std::collections::HashSet;
@@ -46,7 +46,11 @@ impl Graph {
     }
 
     pub fn get_modules(&self) -> HashSet<&str> {
-        todo!()
+        self.hierarchy
+            .all()
+            .visible()
+            .names()
+            .collect()
     }
 
     pub fn count_imports(&self) -> usize {
@@ -150,3 +154,45 @@ impl Graph {
         todo!()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn add_module() -> Result<()> {
+        let mut graph = Graph::default();
+
+        graph.add_module("mypackage")?;
+
+        let result = graph.get_modules();
+
+        assert_eq!(result, HashSet::from(["mypackage"]));
+        Ok(())
+    }
+
+    #[test]
+    fn add_module_doesnt_add_parent() -> Result<()> {
+        let mut graph = Graph::default();
+
+        graph.add_module("mypackage.foo")?;
+
+        let result = graph.get_modules();
+
+        assert_eq!(result, HashSet::from(["mypackage.foo"]));
+        Ok(())
+    }
+
+    #[test]
+    fn add_modules() -> Result<()> {
+        let mut graph = Graph::default();
+
+        graph.add_module("mypackage")?;
+        graph.add_module("mypackage.foo")?;
+
+        let result = graph.get_modules();
+
+        assert_eq!(result, HashSet::from(["mypackage", "mypackage.foo"]));
+        Ok(())
+    }
+}

rust/src/hierarchy.rs

@@ -118,6 +118,9 @@ impl ModuleHierarchy {
         }
     }
 
+    // TODO(peter) Do we want this?
+    // `add_module` already has "get or insert" semantics. The difference here is just that
+    // `get_by_name_or_add` is not making the module visible.
     pub fn get_by_name_or_add(&mut self, name: &str) -> Result<&Module> {
         let token = match self.get_by_name(name) {
             Some(module) => module.token,
@@ -131,6 +134,11 @@ impl ModuleHierarchy {
         self.modules.get(module)
     }
 
+    // TODO(peter) Guarantee order?
+    pub fn all(&self) -> impl Iterator<Item = &Module> {
+        self.modules.values()
+    }
+
     pub fn get_parent(&self, module: ModuleToken) -> Option<&Module> {
         match self.module_parents.get(module) {
             Some(parent) => parent.map(|parent| self.get(parent).unwrap()),
@@ -196,6 +204,22 @@ impl<T: Sized + Clone + IntoIterator<Item = ModuleToken> + Extend<ModuleToken>>
 {
 }
 
+pub trait ModuleIterator<'a>: Iterator<Item = &'a Module> + Sized {
+    fn tokens(self) -> impl Iterator<Item = ModuleToken> {
+        self.map(|m| m.token)
+    }
+
+    fn names(self) -> impl Iterator<Item = &'a str> {
+        self.map(|m| m.name.as_str())
+    }
+
+    fn visible(self) -> impl Iterator<Item = &'a Module> {
+        self.filter(|m| !m.is_invisible)
+    }
+}
+
+impl<'a, T: Iterator<Item = &'a Module>> ModuleIterator<'a> for T {}
+
 fn parent_name(name: &str) -> Option<String> {
     match name.rsplit_once(".") {
         Some((base, _)) => Some(base.to_owned()),