Swap Dict with Counter for returning RowCounts

Author: Iain-S

sqlsynthgen/create.py

@@ -1,5 +1,6 @@
 """Functions and classes to create and populate the target database."""
-from typing import Any, Generator, Mapping, Optional, Sequence, Tuple
+from collections import Counter
+from typing import Any, Generator, Mapping, Sequence, Tuple
 
 from sqlalchemy import Connection, insert
 from sqlalchemy.exc import IntegrityError
@@ -10,7 +11,7 @@
 from sqlsynthgen.utils import create_db_engine, get_sync_engine, logger
 
 Story = Generator[Tuple[str, dict[str, Any]], dict[str, Any], None]
-RowCounts = dict[str, int]
+RowCounts = Counter[str]
 
 
 def create_db_tables(metadata: MetaData) -> None:
@@ -68,15 +69,14 @@ def create_db_data(
         create_db_engine(dst_dsn, schema_name=settings.dst_schema)
     )
 
-    row_counts: RowCounts = {}
+    row_counts: Counter[str] = Counter()
     with dst_engine.connect() as dst_conn:
         for _ in range(num_passes):
-            row_counts = populate(
+            row_counts += populate(
                 dst_conn,
                 sorted_tables,
                 table_generator_dict,
                 story_generator_list,
-                row_counts,
             )
     return row_counts
 
@@ -86,13 +86,13 @@ def _populate_story(
     table_dict: Mapping[str, Table],
     table_generator_dict: Mapping[str, TableGenerator],
     dst_conn: Connection,
-    row_counts: RowCounts,
 ) -> RowCounts:
     """Write to the database all the rows created by the given story."""
     # Loop over the rows generated by the story, insert them into their
     # respective tables. Ideally this would say
     # `for table_name, provided_values in story:`
     # but we have to loop more manually to be able to use the `send` function.
+    row_counts: Counter[str] = Counter()
     try:
         table_name, provided_values = next(story)
         while True:
@@ -129,11 +129,9 @@ def populate(
     tables: Sequence[Table],
     table_generator_dict: Mapping[str, TableGenerator],
     story_generator_list: Sequence[Mapping[str, Any]],
-    row_counts: Optional[RowCounts] = None,
 ) -> RowCounts:
     """Populate a database schema with synthetic data."""
-    if row_counts is None:
-        row_counts = {}
+    row_counts: Counter[str] = Counter()
     table_dict = {table.name: table for table in tables}
     # Generate stories
     # Each story generator returns a python generator (an unfortunate naming clash with
@@ -153,8 +151,8 @@ def populate(
         # Run the inserts for each story within a transaction.
         logger.debug('Generating data for story "%s".', name)
         with dst_conn.begin():
-            row_counts = _populate_story(
-                story, table_dict, table_generator_dict, dst_conn, row_counts
+            row_counts += _populate_story(
+                story, table_dict, table_generator_dict, dst_conn
             )
 
     # Generate individual rows, table by table.

tests/test_create.py

@@ -1,5 +1,6 @@
 """Tests for the create module."""
 import itertools as itt
+from collections import Counter
 from pathlib import Path
 from typing import Any, Generator, Tuple
 from unittest.mock import MagicMock, call, patch
@@ -82,7 +83,7 @@ def mock_story_gen(_: Any) -> Generator[Tuple[str, dict], None, None]:
                 mock_gen = MagicMock(spec=TableGenerator)
                 mock_gen.num_rows_per_pass = num_rows_per_pass
                 mock_gen.return_value = {}
-                row_counts = (
+                row_counts = Counter(
                     {table_name: num_initial_rows} if num_initial_rows > 0 else {}
                 )
 
@@ -97,20 +98,23 @@ def mock_story_gen(_: Any) -> Generator[Tuple[str, dict], None, None]:
                     if num_stories_per_pass > 0
                     else []
                 )
-                row_counts = populate(
+                row_counts += populate(
                     mock_dst_conn,
                     [mock_table],
                     {table_name: mock_gen},
                     story_generators,
-                    row_counts,
                 )
 
                 expected_row_count = (
                     num_stories_per_pass + num_rows_per_pass + num_initial_rows
                 )
                 self.assertEqual(
+                    Counter(
+                        {table_name: expected_row_count}
+                        if expected_row_count > 0
+                        else {}
+                    ),
                     row_counts,
-                    {table_name: expected_row_count} if expected_row_count > 0 else {},
                 )
                 self.assertListEqual(
                     [call(mock_dst_conn)] * (num_stories_per_pass + num_rows_per_pass),
@@ -148,7 +152,7 @@ def test_populate_diff_length(self, mock_insert: MagicMock) -> None:
             "three": mock_gen_three,
         }
 
-        row_counts = populate(mock_dst_conn, tables, row_generators, [], {})
+        row_counts = populate(mock_dst_conn, tables, row_generators, [])
         self.assertEqual(row_counts, {"two": 1, "three": 1})
         self.assertListEqual(
             [call(mock_table_two), call(mock_table_three)], mock_insert.call_args_list
@@ -221,4 +225,4 @@ def my_story() -> Story:
 
         with engine.connect() as conn:
             with conn.begin():
-                _populate_story(my_story(), dict(self.metadata.tables), {}, conn, {})
+                _populate_story(my_story(), dict(self.metadata.tables), {}, conn)

tests/test_main.py

@@ -188,7 +188,7 @@ def test_create_data(
             [
                 call("Creating data."),
                 call("Data created in %s %s.", 1, "pass"),
-                call("%s: %s %s created", "a", "row.", 1),
+                call("%s: %s %s created.", "a", 1, "row"),
             ]
         )