A few updates.

Author: Tim Band

datafaker/populate.py

@@ -1,14 +1,11 @@
-"""This file was auto-generated by datafaker but can be edited manually."""
 from collections.abc import Iterable, Mapping, MutableMapping, Sequence
-from mimesis import Generic, Numeric, Person
+from mimesis import Generic
 from mimesis.locales import Locale
 import sqlalchemy
-import sys
 from typing import Any, Callable
 import yaml
 
-from datafaker.base import FileUploader, TableGenerator, DistributionGenerator, ColumnPresence
-from datafaker.main import load_metadata
+from datafaker.base import FileUploader, TableGenerator
 from datafaker.make import TableGeneratorInfo, StoryGeneratorInfo  #TODO: move these in here!
 
 from datafaker.providers import (
@@ -23,10 +20,6 @@
 from datafaker.utils import logging, get_vocabulary_table_names
 
 generic = Generic(locale=Locale.EN_GB)
-numeric = Numeric()
-person = Person()
-dist_gen = DistributionGenerator()
-column_presence = ColumnPresence()
 
 generic.add_provider(BytesProvider)
 generic.add_provider(ColumnValueProvider)
@@ -36,10 +29,6 @@
 generic.add_provider(TimespanProvider)
 generic.add_provider(WeightedBooleanProvider)
 
-#metadata = load_metadata("{{ orm_file_name }}", "{{ config_file_name }}")
-
-#import {{ row_generator_module_name }}
-#import {{ story_generator_module_name }}
 
 def _eval_structure(config: Any, context: Mapping) -> Any:
     """
@@ -110,16 +99,34 @@ def _get_src_stats(src_stats_filename: str) -> Any:
     """
     Get the SRC_STATS object
     """
-    with open("{{ src_stats_filename }}", "r", encoding="utf-8") as f:
-        return yaml.unsafe_load(f)
+    with open(src_stats_filename, "r", encoding="utf-8") as f:
+        return yaml.load(f, yaml.SafeLoader)
 
 
 class TableGenerator:
 
-    def __init__(self, rows_per_pass: int, dst_db_conn: sqlalchemy.Connection, table_data: TableGeneratorInfo, max_unique_constraint_tries: int | None):
+    def __init__(
+        self,
+        rows_per_pass: int,
+        dst_db_conn: sqlalchemy.Connection,
+        table_data: TableGeneratorInfo,
+        max_unique_constraint_tries: int | None,
+    ) -> None:
+        """
+        Initialize a table generator.
+
+        :param rows_per_pass: How many rows to add for each call to ``__call__``.
+        :param dst_db_conn: Connection to the destination database.
+        :param table_data: Configuration for this generator.
+        :param max_unique_constraint_tries: How many times to redo generation in
+          an attempt to satisfy uniqueness constraints. None means never stop, but
+          this could cause an infinite loop if there are no solutions, or very long
+          execution if there are few solutions with many constraints.
+        """
         self.num_rows_per_pass = rows_per_pass
         self.table_data = table_data
         self.max_unique_constraint_tries = max_unique_constraint_tries
+        self.db_conn = dst_db_conn
         self.existing_constraint_hashes: MutableMapping[str, set[int]] = {}
         self.context: Mapping = {}
         for constraint in table_data.unique_constraints:
@@ -130,11 +137,13 @@ def __init__(self, rows_per_pass: int, dst_db_conn: sqlalchemy.Connection, table
                 for result in query_result
             ])
 
-    def set_context(self, context: Mapping):
+    def set_context(self, context: Mapping) -> None:
+        """Sets all the Python symbols that must be known to the configuration."""
         self.context = context
 
-    def __call__(self, dst_db_conn):
-        result = {}
+    def __call__(self):
+        """Generate some rows of the relevant table in the database."""
+        result: dict[str, Any] = {}
         columns_to_generate = set(self.table_data.nonnull_columns)
         # Which missingness patterns do we want?
         for choice in self.table_data.column_choices:
@@ -166,24 +175,33 @@ def __call__(self, dst_db_conn):
             self.existing_constraint_hashes.add(cf_hash)
         return result
 
-def get_table_generator_dict(self, rows_per_pass: int, dst_db_conn: sqlalchemy.Connection, tables_data: Iterable[TableGeneratorInfo], max_unique_constraint_tries: int | None):
+def get_table_generator_dict(
+    rows_per_pass: int,
+    dst_db_conn: sqlalchemy.Connection,
+    tables_data: Iterable[TableGeneratorInfo],
+    max_unique_constraint_tries: int | None,
+):
+    """Get a dict of table names to row generators that generate rows for that table."""
     return {
-    "{{ table_data.table_name }}": TableGenerator(rows_per_pass, dst_db_conn, table_data, max_unique_constraint_tries)
-    for table_data in tables_data
-}
+        table_data.table_name: TableGenerator(rows_per_pass, dst_db_conn, table_data, max_unique_constraint_tries)
+        for table_data in tables_data
+    }
 
 
-def get_vocab_dict(config: Mapping, metadata: sqlalchemy.MetaData) -> Mapping[str, FileUploader]: {
-    name: FileUploader[metadata.tables[name]]
-    for name in get_vocabulary_table_names(config)
-}
+def get_vocab_dict(config: Mapping, metadata: sqlalchemy.MetaData) -> Mapping[str, FileUploader]:
+    """Get a dict of table names to objects that can populate those tables from YAML files."""
+    return {
+        name: FileUploader(metadata.tables[name])
+        for name in get_vocabulary_table_names(config)
+    }
 
 def get_story_generator_list(story_generator_infos: Iterable[StoryGeneratorInfo], context: Mapping) -> list[Mapping]:
+    """Get a list of mappings describing story generators that must be run."""
     return [
         {
             "function": _call_from_context(gen_data.function_call.function_name, gen_data.function_call.argument_values, context),
-            "num_stories_per_pass": {{ gen_data.num_stories_per_pass }},
-            "name": "{{ gen_data.function_call.function_name }}",
+            "num_stories_per_pass": gen_data.num_stories_per_pass,
+            "name": gen_data.function_call.function_name,
         }
         for gen_data in story_generator_infos
     ]