Cleaned pre-commit checks

Author: Tim Band

datafaker/base.py

@@ -2,7 +2,6 @@
 import gzip
 import os
 import random
-from abc import ABC, abstractmethod
 from collections.abc import Callable
 from dataclasses import dataclass
 from io import TextIOWrapper
@@ -12,7 +11,7 @@
 import yaml
 from sqlalchemy import Connection, insert
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.schema import MetaData, Table
+from sqlalchemy.schema import Table
 
 from datafaker.utils import (
     MAKE_VOCAB_PROGRESS_REPORT_EVERY,

datafaker/create.py

@@ -12,7 +12,7 @@
 from sqlalchemy.schema import CreateColumn, CreateSchema, CreateTable, MetaData, Table
 
 from datafaker.base import FileUploader
-from datafaker.make import StoryGeneratorInfo, get_generation_info
+from datafaker.make import FunctionCall, StoryGeneratorInfo, get_generation_info
 from datafaker.populate import (
     TableGenerator,
     call_function,
@@ -159,12 +159,12 @@ def create_db_data(
         try:
             with src_stats_filename.open(encoding="utf-8") as fh:
                 src_stats = yaml.load(fh, yaml.SafeLoader)
-        except FileNotFoundError:
+        except FileNotFoundError as exc:
             logger.error(
-                "No source stats file '%', this should be the output of the 'make-stats' command",
+                "No source stats file '%s', this should be the output of the 'make-stats' command",
                 src_stats_filename,
             )
-            raise typer.Exit(1)
+            raise typer.Exit(1) from exc
     else:
         src_stats = None
     return create_db_data_into(
@@ -205,7 +205,7 @@ def create_db_data_into(
     context = get_symbols(
         gen_info.row_generator_module_name,
         gen_info.story_generator_module_name,
-        get_property(config, "object_instantiation", dict, {}),
+        get_property(config, "object_instantiation", {}),
         src_stats,
         metadata,
     )
@@ -229,6 +229,14 @@ def create_db_data_into(
     return row_counts
 
 
+def empty_story_generator() -> (
+    Generator[tuple[str, dict[str, Any]], dict[str, Any], None]
+):
+    """Get a story generator that generates no values."""
+    empt: list[tuple[str, dict[str, Any]]] = []
+    yield from empt
+
+
 # pylint: disable=too-many-instance-attributes
 class StoryIterator:
     """Iterates through all the rows produced by all the stories."""
@@ -250,9 +258,10 @@ def __init__(
         self._final_values: dict[str, Any] | None = None
         # Number of times the current story should be run
         self._story_counts = 1
-        self._story_function_call = None
+        self._story_function_call: FunctionCall
         self._context = context
-        self._story = iter([])
+        self._story = empty_story_generator()
+        self._provided_values: dict[str, Any]
         self.next()
 
     def _get_next_story(self) -> bool:
@@ -276,6 +285,7 @@ def _get_next_story(self) -> bool:
         return True
 
     def _get_values(self) -> None:
+        """Get the values from the current story and advance the iterator."""
         if self._final_values is None:
             self._table_name, self._provided_values = next(self._story)
         else:
@@ -342,7 +352,7 @@ def next(self) -> None:
             try:
                 self._get_values()
                 return
-            except StopIteration as exc:
+            except StopIteration:
                 self._final_values = None
                 self._story_counts -= 1
                 if 0 < self._story_counts:

datafaker/dump.py

@@ -3,7 +3,6 @@
 import io
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 import pandas as pd
 import sqlalchemy

datafaker/generators/partitioned.py

@@ -445,7 +445,7 @@ def get_named_tables(self) -> Mapping[str, str]:
 
     def __init__(self, config: Mapping[str, Any]) -> None:
         """Initialize the null partitioned generator factory."""
-        tables = get_property(config, "tables", dict, {})
+        tables: dict[str, Any] = get_property(config, "tables", {})
         self._named_tables = {
             table_name: table_conf["name_column"]
             for table_name, table_conf in tables.items()

datafaker/main.py

@@ -44,7 +44,6 @@
     generated_tables,
     generators_require_stats,
     get_flag,
-    import_file,
     logger,
     read_config_file,
     sorted_non_vocabulary_tables,
@@ -263,22 +262,22 @@ def create_tables(
 
 @app.command()
 def create_generators(
-    orm_file: Path = Option(
+    _orm_file: Path = Option(
         ORM_FILENAME,
         help="The name of the ORM yaml file",
         dir_okay=False,
     ),
-    df_file: Path = Option(
+    _df_file: Path = Option(
         DF_FILENAME,
         help="Path to write Python generators to.",
         dir_okay=False,
     ),
-    config_file: Path = Option(
+    _config_file: Path = Option(
         CONFIG_FILENAME,
         help="The configuration file",
         dir_okay=False,
     ),
-    stats_file: Optional[Path] = Option(
+    _stats_file: Optional[Path] = Option(
         None,
         help=(
             "Statistics file (output of make-stats); default is src-stats.yaml if the "
@@ -287,7 +286,7 @@ def create_generators(
         show_default=False,
         dir_okay=False,
     ),
-    force: bool = Option(
+    _force: bool = Option(
         False, "--force", "-f", help="Overwrite any existing Python generators file."
     ),
 ) -> None:

datafaker/make.py

@@ -7,19 +7,25 @@
 from datetime import datetime
 from pathlib import Path
 from types import TracebackType
-from typing import Any, Final, Optional, Tuple, Type, Union
+from typing import Any, Final, Optional, Tuple, Type
 
 import pandas as pd
 import snsql
 import yaml
 from black import FileMode, format_str
 from jinja2 import Environment, FileSystemLoader, Template
 from mimesis.providers.base import BaseProvider
-from sqlalchemy import CursorResult, Engine, MetaData, UniqueConstraint, text
+from sqlalchemy import CursorResult, Engine, MetaData, text
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.engine import Connection
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
-from sqlalchemy.schema import Column, Table
+from sqlalchemy.schema import (
+    Column,
+    ColumnCollectionConstraint,
+    PrimaryKeyConstraint,
+    Table,
+    UniqueConstraint,
+)
 from sqlalchemy.sql import Executable, sqltypes
 from typing_extensions import Self
 
@@ -28,10 +34,12 @@
 from datafaker.settings import get_source_dsn, get_source_schema
 from datafaker.utils import (
     MaybeAsyncEngine,
+    constraint_name,
     create_db_engine,
     download_table,
     get_columns_assigned,
     get_property,
+    get_property_or_none,
     get_related_table_names,
     get_row_generators,
     get_sync_engine,
@@ -66,8 +74,8 @@ class FunctionCall:
     """Which function to call with what."""
 
     function_name: str
-    args: list[Any]
-    kwargs: dict[str, Any]
+    args: Sequence[Any]
+    kwargs: Mapping[str, Any]
 
 
 @dataclass
@@ -110,18 +118,6 @@ def make_column_choices(
     ]
 
 
-@dataclass
-class _PrimaryConstraint:
-    """
-    Describes a Uniqueness constraint for a multi-column primary key.
-
-    Not a real constraint, but enough to write df.py.
-    """
-
-    columns: list[Column]
-    name: str
-
-
 @dataclass
 class TableGeneratorInfo:
     """Contains the df.py content related to regular tables."""
@@ -132,7 +128,7 @@ class TableGeneratorInfo:
     column_choices: list[ColumnChoice]
     rows_per_pass: int
     row_gens: list[RowGeneratorInfo] = field(default_factory=list)
-    unique_constraints: Sequence[Union[UniqueConstraint, _PrimaryConstraint]] = field(
+    unique_constraints: Sequence[ColumnCollectionConstraint] = field(
         default_factory=list
     )
 
@@ -467,19 +463,6 @@ def _get_provider_for_column(column: Column) -> Tuple[list[str], str, dict[str,
     return variable_names, generator_function, generator_arguments
 
 
-def _constraint_sort_key(constraint: UniqueConstraint) -> str:
-    """Extract a string out of a UniqueConstraint that is unique to that constraint.
-
-    We sort the constraints so that the output of make_tables is deterministic, this is
-    the sort key.
-    """
-    return (
-        constraint.name
-        if isinstance(constraint.name, str)
-        else "_".join(map(str, constraint.columns))
-    )
-
-
 def _get_generator_for_table(
     table_config: Mapping[str, Any],
     table: Table,
@@ -491,12 +474,12 @@ def _get_generator_for_table(
             for constraint in table.constraints
             if isinstance(constraint, UniqueConstraint)
         ),
-        key=_constraint_sort_key,
+        key=constraint_name,
     )
     primary_keys = [c for c in table.columns if c.primary_key]
-    constraints: Sequence[UniqueConstraint | _PrimaryConstraint] = unique_constraints
+    constraints: Sequence[ColumnCollectionConstraint] = unique_constraints
     if 1 < len(primary_keys):
-        primary_constraint = _PrimaryConstraint(
+        primary_constraint = PrimaryKeyConstraint(
             columns=primary_keys, name=make_primary_key_name(table.name)
         )
         constraints = unique_constraints + [primary_constraint]
@@ -514,7 +497,7 @@ def _get_generator_for_table(
         class_name=table.name.title().replace(".", "") + "Generator",
         nonnull_columns=nonnull_columns,
         column_choices=column_choices,
-        rows_per_pass=get_property(table_config, "num_rows_per_pass", int, 1),
+        rows_per_pass=get_property(table_config, "num_rows_per_pass", 1),
         unique_constraints=constraints,
     )
 
@@ -583,6 +566,7 @@ def make_vocabulary_tables(
 
 
 @dataclass
+# pylint: disable=too-many-instance-attributes
 class GenerationInfo:
     """Information for the generation of all data."""
 
@@ -618,14 +602,16 @@ def get_generation_info(
 
     :return: A string that is a valid Python module, once written to file.
     """
-    row_generator_module_name = get_property(
-        config, "row_generators_module", str | None, None
+    row_generator_module_name = get_property_or_none(
+        config, "row_generators_module", str
+    )
+    story_generator_module_name = get_property_or_none(
+        config, "story_generators_module", str
     )
-    story_generator_module_name = get_property(
-        config, "story_generators_module", str | None, None
+    object_instantiation: dict[str, Any] = get_property(
+        config, "object_instantiation", {}
     )
-    object_instantiation = get_property(config, "object_instantiation", dict, {})
-    tables_config = get_property(config, "tables", dict, {})
+    tables_config: dict[str, Any] = get_property(config, "tables", {})
 
     tables: list[TableGeneratorInfo] = []
     vocabulary_tables: list[VocabularyTableGeneratorInfo] = []
@@ -654,8 +640,8 @@ def get_generation_info(
 
     story_generators = _get_story_generators(config)
 
-    max_unique_constraint_tries = get_property(
-        config, "max-unique-constraint-tries", int | None, None
+    max_unique_constraint_tries = get_property_or_none(
+        config, "max-unique-constraint-tries", int
     )
     return GenerationInfo(
         provider_imports=PROVIDER_IMPORTS,
@@ -724,7 +710,7 @@ def make_tables_file(
     if parquet_dir is not None:
         extra_meta = get_parquet_orm(parquet_dir)
         if extra_meta:
-            md_tables = get_property(meta_dict, "tables", dict, {})
+            md_tables: dict[str, Any] = get_property(meta_dict, "tables", {})
             new_tables = {**extra_meta, **md_tables}
             meta_dict["tables"] = new_tables