fix tests, fill in impl details

tskit-dev · Jul 15, 2024 · 7c13513 · 7c13513
1 parent 54b1690
commit 7c13513
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 13 deletions.
diff --git a/src/traits.rs b/src/traits.rs
@@ -133,14 +133,41 @@ pub trait TableIteration: TableAccess {
     fn edges_iter(&self) -> impl Iterator<Item = crate::EdgeTableRow> + '_ {
         self.edges().iter()
     }
+    fn nodes_iter(&self) -> impl Iterator<Item = crate::NodeTableRow> + '_ {
+        self.nodes().iter()
+    }
+    fn sites_iter(&self) -> impl Iterator<Item = crate::SiteTableRow> + '_ {
+        self.sites().iter()
+    }
+    fn mutations_iter(&self) -> impl Iterator<Item = crate::MutationTableRow> + '_ {
+        self.mutations().iter()
+    }
+    fn migrations_iter(&self) -> impl Iterator<Item = crate::MigrationTableRow> + '_ {
+        self.migrations().iter()
+    }
+    fn populations_iter(&self) -> impl Iterator<Item = crate::PopulationTableRow> + '_ {
+        self.populations().iter()
+    }
 }
 
 pub trait ObjectSafeTableIteration: TableAccess {
     fn edges_iter(&self) -> Box<dyn Iterator<Item = crate::EdgeTableRow> + '_> {
         Box::new(self.edges().iter())
     }
+    fn nodes_iter(&self) -> Box<dyn Iterator<Item = crate::NodeTableRow> + '_> {
+        Box::new(self.nodes().iter())
+    }
+    fn sites_iter(&self) -> Box<dyn Iterator<Item = crate::SiteTableRow> + '_> {
+        Box::new(self.sites().iter())
+    }
+    fn mutations_iter(&self) -> Box<dyn Iterator<Item = crate::MutationTableRow> + '_> {
+        Box::new(self.mutations().iter())
+    }
+    fn migrations_iter(&self) -> Box<dyn Iterator<Item = crate::MigrationTableRow> + '_> {
+        Box::new(self.migrations().iter())
+    }
     fn populations_iter(&self) -> Box<dyn Iterator<Item = crate::PopulationTableRow> + '_> {
-        Box::new(self.populations().iter())
+        Box::new(Box::new(self.populations().iter()))
     }
 }
 

diff --git a/tests/test_table_traits.rs b/tests/test_table_traits.rs
@@ -73,12 +73,12 @@ impl IteratorOutput {
     where
         T: TableIteration,
     {
-        let edges = iterator.edges().iter().collect::<Vec<_>>();
-        let nodes = iterator.nodes().iter().collect::<Vec<_>>();
-        let sites = iterator.sites().iter().collect::<Vec<_>>();
-        let mutations = iterator.mutations().iter().collect::<Vec<_>>();
-        let populations = iterator.populations().iter().collect::<Vec<_>>();
-        let migrations = iterator.migrations().iter().collect::<Vec<_>>();
+        let edges = iterator.edges_iter().collect::<Vec<_>>();
+        let nodes = iterator.nodes_iter().collect::<Vec<_>>();
+        let sites = iterator.sites_iter().collect::<Vec<_>>();
+        let mutations = iterator.mutations_iter().collect::<Vec<_>>();
+        let populations = iterator.populations_iter().collect::<Vec<_>>();
+        let migrations = iterator.migrations_iter().collect::<Vec<_>>();
         Self {
             edges,
             nodes,
@@ -90,12 +90,13 @@ impl IteratorOutput {
     }
 
     fn new_from_dyn(dynamic: &dyn tskit::ObjectSafeTableIteration) -> Self {
-        let edges = dynamic.edges().iter().collect::<Vec<_>>();
-        let nodes = dynamic.nodes().iter().collect::<Vec<_>>();
-        let sites = dynamic.sites().iter().collect::<Vec<_>>();
-        let mutations = dynamic.mutations().iter().collect::<Vec<_>>();
-        let populations = dynamic.populations().iter().collect::<Vec<_>>();
-        let migrations = dynamic.migrations().iter().collect::<Vec<_>>();
+        let edges_iter: Box<dyn Iterator<Item = tskit::EdgeTableRow> + '_> = dynamic.edges_iter();
+        let edges = edges_iter.collect::<Vec<_>>();
+        let nodes = dynamic.nodes_iter().collect::<Vec<_>>();
+        let sites = dynamic.sites_iter().collect::<Vec<_>>();
+        let mutations = dynamic.mutations_iter().collect::<Vec<_>>();
+        let populations = dynamic.populations_iter().collect::<Vec<_>>();
+        let migrations = dynamic.migrations_iter().collect::<Vec<_>>();
         Self {
             edges,
             nodes,