enable repeated iteration of diff

Author: dlawin

data_diff/__main__.py

@@ -408,8 +408,6 @@ def _main(
         diff_iter = islice(diff_iter, int(limit))
 
     if stats:
-        # required to create this variable before get_stats
-        diff_list = list(diff_iter)
         if json_output:
             rich.print(json.dumps(diff_iter.get_stats_dict()))
         else:

data_diff/diff_tables.py

@@ -94,34 +94,31 @@ class DiffResultWrapper:
     result_list: list = field(default_factory=list)
 
     def __iter__(self):
+        yield from self.result_list
         for i in self.diff:
             self.result_list.append(i)
             yield i
 
     def _get_stats(self) -> DiffStats:
 
+        list(self)
         diff_by_key = {}
-        if len(self.result_list) > 0:
-            for sign, values in self.result_list:
-                k = values[: len(self.info_tree.info.tables[0].key_columns)]
-                if k in diff_by_key:
-                    assert sign != diff_by_key[k]
-                    diff_by_key[k] = "!"
-                else:
-                    diff_by_key[k] = sign
-
-            diff_by_sign = {k: 0 for k in "+-!"}
-            for sign in diff_by_key.values():
-                diff_by_sign[sign] += 1
-
-            table1_count = self.info_tree.info.rowcounts[1]
-            table2_count = self.info_tree.info.rowcounts[2]
-            unchanged = table1_count - diff_by_sign["-"] - diff_by_sign["!"]
-            diff_percent = 1 - unchanged / max(table1_count, table2_count)
-        else:
-            raise RuntimeError(
-                "result_list is empty, consume the diff iterator to populate values: e.g. \ndiff_iter = diff_tables(...) \ndiff_list = list(diff_iter) \ndiff_iter.print_stats(json_output)"
-            )
+        for sign, values in self.result_list:
+            k = values[: len(self.info_tree.info.tables[0].key_columns)]
+            if k in diff_by_key:
+                assert sign != diff_by_key[k]
+                diff_by_key[k] = "!"
+            else:
+                diff_by_key[k] = sign
+
+        diff_by_sign = {k: 0 for k in "+-!"}
+        for sign in diff_by_key.values():
+            diff_by_sign[sign] += 1
+
+        table1_count = self.info_tree.info.rowcounts[1]
+        table2_count = self.info_tree.info.rowcounts[2]
+        unchanged = table1_count - diff_by_sign["-"] - diff_by_sign["!"]
+        diff_percent = 1 - unchanged / max(table1_count, table2_count)
 
         return DiffStats(diff_by_sign, table1_count, table2_count, unchanged, diff_percent)
 

tests/test_api.py

@@ -81,12 +81,11 @@ def test_api_get_stats_string(self):
         t1 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_src_name)
         t2 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_dst_name)
         diff = diff_tables(t1, t2)
-        diff_list = list(diff)
         output = diff.get_stats_string()
 
         self.assertEqual(expected_string, output)
         self.assertIsNotNone(diff)
-        assert len(diff_list) == 1
+        assert len(list(diff)) == 1
 
         t1.database.close()
         t2.database.close()
@@ -96,23 +95,11 @@ def test_api_get_stats_dict(self):
         t1 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_src_name)
         t2 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_dst_name)
         diff = diff_tables(t1, t2)
-        diff_list = list(diff)
         output = diff.get_stats_dict()
 
         self.assertEqual(expected_dict, output)
         self.assertIsNotNone(diff)
-        assert len(diff_list) == 1
-
-        t1.database.close()
-        t2.database.close()
-
-    def test_api_print_error(self):
-        t1 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_src_name)
-        t2 = connect_to_table(TEST_MYSQL_CONN_STRING, (self.table_dst_name,))
-        diff = diff_tables(t1, t2)
-
-        with self.assertRaises(RuntimeError):
-            diff.get_stats_string()
+        assert len(list(diff)) == 1
 
         t1.database.close()
         t2.database.close()

tests/test_cli.py

@@ -101,3 +101,9 @@ def test_stats_json(self):
             TEST_MYSQL_CONN_STRING, self.table_src_name, TEST_MYSQL_CONN_STRING, self.table_dst_name, "-s", "--json"
         )
         assert len(diff_output) == 2
+    
+    def test_stats_no_diff(self):
+        diff_output = run_datadiff_cli(
+            TEST_MYSQL_CONN_STRING, self.table_src_name, TEST_MYSQL_CONN_STRING, self.table_src_name, "-s"
+        )
+        assert len(diff_output) == 11