A few pre-commit fixes

Author: Tim Band

datafaker/base.py

@@ -1,15 +1,15 @@
 """Base table generator classes."""
+import gzip
+import os
+import random
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from dataclasses import dataclass
-import gzip
 from io import TextIOWrapper
-import os
 from pathlib import Path
-import random
 from typing import Any
-import yaml
 
+import yaml
 from sqlalchemy import Connection, insert
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.schema import MetaData, Table

datafaker/create.py

@@ -2,17 +2,17 @@
 from collections import Counter
 from pathlib import Path
 from typing import Any, Generator, Iterable, Iterator, Mapping, Sequence, Tuple
-import yaml
 
+import typer
+import yaml
 from sqlalchemy import Connection, insert, inspect
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.orm import Session
 from sqlalchemy.schema import CreateColumn, CreateSchema, CreateTable, MetaData, Table
-import typer
 
 from datafaker.base import FileUploader
-from datafaker.make import get_generation_info, StoryGeneratorInfo
+from datafaker.make import StoryGeneratorInfo, get_generation_info
 from datafaker.populate import (
     TableGenerator,
     call_function,
@@ -255,7 +255,7 @@ def __init__(
         self._story = iter([])
         self.next()
 
-    def _get_next_story(self) -> None:
+    def _get_next_story(self) -> bool:
         """
         Iterate to the next ``_story_infos``.
 
@@ -265,7 +265,9 @@ def _get_next_story(self) -> None:
             sgi = next(self._story_infos)
             self._story_counts = sgi.num_stories_per_pass
             self._story_function_call = sgi.function_call
-            logger.info("Generating data for story '%s'", sgi.function_call.function_name)
+            logger.info(
+                "Generating data for story '%s'", sgi.function_call.function_name
+            )
             self._story = call_function(sgi.function_call, self._context)
             self._final_values = None
         except StopIteration:
@@ -289,7 +291,6 @@ def is_ended(self) -> bool:
         """
         return self._story_counts == -1
 
-
     def has_table(self, table_name: str) -> bool:
         """Check if we have a row for table ``table_name``."""
         return table_name == self._table_name
@@ -346,7 +347,9 @@ def next(self) -> None:
                 self._story_counts -= 1
                 if 0 < self._story_counts:
                     # Reinitialize the same story again
-                    self._story = call_function(self._story_function_call, self._context)
+                    self._story = call_function(
+                        self._story_function_call, self._context
+                    )
                 elif not self._get_next_story():
                     self._story_counts = -1
                     return

datafaker/dump.py

@@ -1,10 +1,9 @@
 """Data dumping functions."""
 import csv
 import io
-from typing import TYPE_CHECKING
-
 from abc import ABC, abstractmethod
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import pandas as pd
 import sqlalchemy

datafaker/main.py

@@ -29,11 +29,7 @@
     update_missingness,
 )
 from datafaker.interactive.base import DbCmd
-from datafaker.make import (
-    make_src_stats,
-    make_tables_file,
-    make_vocabulary_tables,
-)
+from datafaker.make import make_src_stats, make_tables_file, make_vocabulary_tables
 from datafaker.remove import remove_db_data, remove_db_tables, remove_db_vocab
 from datafaker.settings import (
     SettingsError,

datafaker/make.py

@@ -66,8 +66,8 @@ class FunctionCall:
     """Which function to call with what."""
 
     function_name: str
-    args: list[str]
-    kwargs: dict[str, str]
+    args: list[Any]
+    kwargs: dict[str, Any]
 
 
 @dataclass
@@ -585,6 +585,7 @@ def make_vocabulary_tables(
 @dataclass
 class GenerationInfo:
     """Information for the generation of all data."""
+
     provider_imports: list[str]
     row_generator_module_name: str | None
     story_generator_module_name: str | None
@@ -623,9 +624,7 @@ def get_generation_info(
     story_generator_module_name = get_property(
         config, "story_generators_module", str | None, None
     )
-    object_instantiation = get_property(
-        config, "object_instantiation", dict, {}
-    )
+    object_instantiation = get_property(config, "object_instantiation", dict, {})
     tables_config = get_property(config, "tables", dict, {})
 
     tables: list[TableGeneratorInfo] = []

datafaker/populate.py

@@ -1,13 +1,13 @@
 from collections.abc import Iterable, Mapping, MutableMapping, Sequence
 from pathlib import Path
+from typing import Any, Callable
+
+import sqlalchemy
 from mimesis import Generic
 from mimesis.locales import Locale
-import sqlalchemy
-from typing import Any, Callable
 
-from datafaker.base import FileUploader, ColumnPresence
+from datafaker.base import ColumnPresence, FileUploader
 from datafaker.make import FunctionCall, TableGeneratorInfo
-
 from datafaker.providers import (
     BytesProvider,
     ColumnValueProvider,
@@ -20,7 +20,8 @@
 )
 from datafaker.utils import get_vocabulary_table_names, import_file
 
-def make_generic():
+
+def make_generic() -> Generic:
     g = Generic(locale=Locale.EN_GB)
     g.add_providers(
         BytesProvider,
@@ -38,7 +39,7 @@ def make_generic():
 generic = make_generic()
 
 
-def reset_generic():
+def reset_generic() -> None:
     """
     Reset all the generators.
 
@@ -63,10 +64,7 @@ def _eval_structure(config: Any, context: Mapping) -> Any:
         except NameError as exc:
             raise exc
     if isinstance(config, Mapping):
-        return {
-            k: _eval_structure(v, context)
-            for k, v in config.items()
-        }
+        return {k: _eval_structure(v, context) for k, v in config.items()}
     if isinstance(config, Sequence):
         return [_eval_structure(v, context) for v in config]
     return config
@@ -96,10 +94,7 @@ def _get_object(class_name: str, context: Mapping) -> Any:
 
 
 def _call_from_context(
-    callable_name: str,
-    args: list[Any],
-    kwargs: dict[str, Any],
-    context: Mapping
+    callable_name: str, args: list[Any], kwargs: dict[str, Any], context: Mapping
 ) -> Any:
     """
     Call a callable from the classes (or functions) in the context.
@@ -111,14 +106,8 @@ def _call_from_context(
     cls = _get_object(callable_name, context)
     if not isinstance(cls, Callable):
         return None
-    arg_objs = [
-        _eval_structure(arg, context)
-        for arg in args
-    ]
-    kwarg_objs = {
-        k: _eval_structure(v, context)
-        for k, v in kwargs.items()
-    }
+    arg_objs = [_eval_structure(arg, context) for arg in args]
+    kwarg_objs = {k: _eval_structure(v, context) for k, v in kwargs.items()}
     return cls(*arg_objs, **kwarg_objs)
 
 
@@ -191,7 +180,6 @@ def _get_symbols_instantiation(symbols: dict[str, Any], objs: dict[str, Any]) ->
 
 
 class TableGenerator:
-
     def __init__(
         self,
         dst_db_conn: sqlalchemy.Connection,
@@ -217,10 +205,9 @@ def __init__(
             for constraint in table_data.unique_constraints:
                 expr = sqlalchemy.select(constraint.columns)
                 query_result = dst_db_conn.execute(expr).fetchall()
-                self.existing_constraint_hashes[constraint.name] = set([
-                    hash(tuple(result))
-                    for result in query_result
-                ])
+                self.existing_constraint_hashes[constraint.name] = set(
+                    [hash(tuple(result)) for result in query_result]
+                )
 
     @property
     def num_rows_per_pass(self):
@@ -242,13 +229,17 @@ def __call__(self, db_conn: sqlalchemy.Connection):
         columns_to_generate = set(self.table_data.nonnull_columns)
         # Which missingness patterns do we want?
         for choice in self.table_data.column_choices:
-            cols = _call_from_context(choice.function_name, choice.args, choice.kwargs, self.context)
+            cols = _call_from_context(
+                choice.function_name, choice.args, choice.kwargs, self.context
+            )
             columns_to_generate.update(cols)
 
         max_tries = self.max_unique_constraint_tries
         while columns_to_generate:
             if max_tries == 0:
-                raise RuntimeError(f"Failed to satisfy unique constraints for table {self.table_data.table_name} after {self.max_unique_constraint_tries} attempts.")
+                raise RuntimeError(
+                    f"Failed to satisfy unique constraints for table {self.table_data.table_name} after {self.max_unique_constraint_tries} attempts."
+                )
             if max_tries is not None:
                 max_tries -= 1
             for row_gen in self.table_data.row_gens:
@@ -264,15 +255,11 @@ def __call__(self, db_conn: sqlalchemy.Connection):
                             result[variable_name] = values[index]
             columns_to_generate = set()
             for constraint in self.table_data.unique_constraints:
-                cf_hash = hash(tuple(
-                    result[col.name] for col in constraint.columns
-                ))
+                cf_hash = hash(tuple(result[col.name] for col in constraint.columns))
                 if cf_hash in self.existing_constraint_hashes[constraint.name]:
                     columns_to_generate.update(c.name for c in constraint.columns)
         for constraint in self.table_data.unique_constraints:
-            cf_hash = hash(tuple(
-                result[col.name] for col in constraint.columns
-            ))
+            cf_hash = hash(tuple(result[col.name] for col in constraint.columns))
             self.existing_constraint_hashes[constraint.name].add(cf_hash)
         return result
 

tests/test_create.py

@@ -13,7 +13,6 @@
 from sqlalchemy import Connection, Engine, select
 from sqlalchemy.schema import MetaData, Table
 
-from datafaker.populate import TableGenerator
 from datafaker.create import (
     create_db_data_into,
     create_db_tables,
@@ -22,7 +21,8 @@
     populate,
 )
 from datafaker.make import FunctionCall, StoryGeneratorInfo
-from datafaker.serialize_metadata import metadata_to_dict, dict_to_metadata
+from datafaker.populate import TableGenerator
+from datafaker.serialize_metadata import dict_to_metadata, metadata_to_dict
 from datafaker.utils import sorted_non_vocabulary_tables
 from tests.utils import DatafakerTestCase, GeneratesDBTestCase, RequiresDBTestCase
 
@@ -128,7 +128,7 @@ def story() -> Generator[Tuple[str, dict], None, None]:
                     {table_name: num_initial_rows} if num_initial_rows > 0 else {}
                 )
 
-                story_generators: list[dict[str, Any]] = (
+                story_generators: list[StoryGeneratorInfo] = (
                     [
                         StoryGeneratorInfo(
                             "mock_story_gen name",
@@ -345,6 +345,7 @@ def __call__(self, connection: Connection, base_path: Path) -> None:
 
 class CreateDataTestCase(RequiresDBTestCase):
     """Tests for create-data."""
+
     dump_file_path = "empty.sql"
     database_name = "empty"
     schema_name = "public"
@@ -379,21 +380,23 @@ def test_create_data_minimal(self) -> None:
         with self.sync_engine.connect() as connection:
             stmt = select(metadata.tables["one"])
             rows = connection.execute(stmt).fetchall()
-            self.assertListEqual(rows, [(1,), (2,), (3,), (4,)])
-        self.assertListEqual(list(row_counts.keys()), ['one'])
+            self.assertEqual(rows, [(1,), (2,), (3,), (4,)])
+        self.assertListEqual(list(row_counts.keys()), ["one"])
         self.assertEqual(row_counts["one"], generate_count)
 
     def test_unique_constraint_minimal(self) -> None:
         config = {
             "tables": {
                 "one": {
-                    "row_generators": [{
-                        "name": "dist_gen.constant",
-                        "kwargs": {
-                            "value": 123,
-                        },
-                        "columns_assigned": ["tiger"],
-                    }]
+                    "row_generators": [
+                        {
+                            "name": "dist_gen.constant",
+                            "kwargs": {
+                                "value": 123,
+                            },
+                            "columns_assigned": ["tiger"],
+                        }
+                    ]
                 }
             },
             "max-unique-constraint-tries": 20,
@@ -410,9 +413,7 @@ def test_unique_constraint_minimal(self) -> None:
                             "type": "INTEGER",
                         },
                     },
-                    "unique": [
-                        {"name": "tiger_uniq", "columns": ["tiger"]}
-                    ]
+                    "unique": [{"name": "tiger_uniq", "columns": ["tiger"]}],
                 }
             }
         }