rowsmyth 0.1.0__py3-none-any.whl

rowsmyth/__init__.py

@@ -0,0 +1,16 @@
+from importlib.metadata import version
+
+from ._base import TableSpecMixin, declarative_base
+from ._builder import FactoryBuilder
+from ._dataset import Dataset
+from ._variant import variant
+
+__version__ = version("rowsmyth")
+__all__ = [
+    "Dataset",
+    "FactoryBuilder",
+    "TableSpecMixin",
+    "__version__",
+    "declarative_base",
+    "variant",
+]

rowsmyth/_base.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.orm import registry as _registry
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from sqlalchemy import MetaData
+
+    from ._builder import FactoryBuilder
+    from ._dataset import Dataset
+
+
+# SQLAlchemy's declared_attr only works for instance-level access, so we use a
+# custom descriptor to expose table metadata as class-level read-only properties.
+class _classproperty:  # noqa: N801
+    def __init__(self, func: Callable[..., Any]) -> None:
+        self.func = func
+
+    def __get__(self, obj: Any, cls: type | None = None) -> Any:
+        return self.func(cls if cls is not None else type(obj))
+
+
+class TableSpecMixin:
+    __variants__: dict[str, Callable[..., Any]]
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+        cls.__variants__ = {
+            name: val
+            for name, val in vars(cls).items()
+            if callable(val) and getattr(val, "_is_variant", False)
+        }
+
+    @_classproperty
+    def __comment__(cls) -> str | None:  # noqa: N805
+        return cls.__table__.comment  # ty: ignore[unresolved-attribute]
+
+    @_classproperty
+    def __table_info__(cls) -> dict[str, Any]:  # noqa: N805
+        return cls.__table__.info  # ty: ignore[unresolved-attribute]
+
+    @_classproperty
+    def __column_info__(cls) -> dict[str, dict[str, Any]]:  # noqa: N805
+        return {
+            prop.key: col.info
+            for prop in cls.__mapper__.column_attrs  # ty: ignore[unresolved-attribute]
+            for col in prop.columns
+        }
+
+    @_classproperty
+    def __expectations__(cls) -> dict[str, str]:  # noqa: N805
+        from sqlalchemy import CheckConstraint
+
+        return {
+            c.name: str(c.sqltext)
+            for c in cls.__table__.constraints  # ty: ignore[unresolved-attribute]
+            if isinstance(c, CheckConstraint) and c.name is not None
+        }
+
+    @_classproperty
+    def __spark_schema__(cls) -> Any:  # noqa: N805
+        from ._spark import to_spark_schema
+
+        return to_spark_schema(cls)  # ty: ignore[invalid-argument-type]
+
+    def __repr__(self) -> str:
+        attrs = ", ".join(
+            f"{c.key}={getattr(self, c.key)!r}"
+            for c in self.__mapper__.columns  # ty: ignore[unresolved-attribute]
+        )
+        return f"{self.__class__.__name__}({attrs})"
+
+    @classmethod
+    def generators(cls) -> dict[Any, Any]:
+        """Override to supply factory-boy declarations for column generation.
+
+        Returns:
+            A dict mapping column attributes (or string names) to factory-boy
+            declarations (e.g. ``factory.Faker(...)``). Keys can be the column
+            attribute itself (``cls.name``) or a plain string (``"name"``).
+
+        Example:
+            >>> TableSpecMixin.generators()
+            {}
+        """
+        return {}
+
+    @classmethod
+    def factory(cls, count: int = 1, max_count: int | None = None) -> FactoryBuilder:
+        """Return a FactoryBuilder for this model.
+
+        Args:
+            count: Exact number of instances to generate, or minimum when
+                ``max_count`` is also provided.
+            max_count: Upper bound for a random count in
+                ``[count, max_count]`` inclusive. ``None`` means exact.
+
+        Returns:
+            A ``FactoryBuilder`` ready to be configured with ``.has()``,
+            ``.mix()``, ``.where()`` and ``.create()``.
+
+        Example:
+            >>> # User.factory(10)         # exactly 10
+            >>> # User.factory(1, 5)       # random 1-5 per parent
+            >>> # User.factory(5).create() # generate and persist
+        """  # doctest: +SKIP
+        from ._builder import FactoryBuilder
+
+        return FactoryBuilder(cls, count, max_count)
+
+    @classmethod
+    def dataset(cls, *builders: FactoryBuilder) -> Dataset:
+        """Return a Dataset generating one row per registered model.
+
+        Args:
+            *builders: Optional ``FactoryBuilder`` overrides for specific
+                models. Models without an override get one instance using
+                their default generators.
+
+        Returns:
+            A ``Dataset`` whose ``.create()`` returns a dict keyed by table
+            name, values are lists of persisted model instances.
+
+        Raises:
+            TypeError: If any argument is not a ``FactoryBuilder``.
+            ValueError: If a builder's model is not registered to this Base.
+
+        Example:
+            >>> # result = Base.dataset(User.factory(10)).create()
+            >>> # result["users"]  # -> list of 10 User instances
+        """  # doctest: +SKIP
+        from ._builder import FactoryBuilder
+        from ._dataset import Dataset
+
+        base_models = {mapper.class_ for mapper in cls.registry.mappers}  # ty: ignore[unresolved-attribute]
+        for b in builders:
+            if not isinstance(b, FactoryBuilder):
+                msg = f"Expected a FactoryBuilder, got {type(b).__name__}"
+                raise TypeError(msg)
+            if b.model not in base_models:
+                msg = f"{b.model.__name__} is not registered to this Base."
+                raise ValueError(msg)
+        overrides = {b.model: b for b in builders}
+        return Dataset(cls, overrides)
+
+
+def declarative_base(
+    metadata: MetaData | None = None,
+    type_annotation_map: dict[Any, Any] | None = None,
+    registry: _registry | None = None,
+) -> type[TableSpecMixin]:
+    """Create a new SQLAlchemy declarative base with rowsmyth capabilities.
+
+    Args:
+        metadata: Optional ``MetaData`` instance shared by all models that
+            subclass the returned base. A new ``MetaData`` is created if
+            not provided.
+        type_annotation_map: Optional mapping of Python types to SQLAlchemy
+            ``TypeEngine`` classes or instances, used by ``Mapped[]``
+            annotations.
+        registry: Optional pre-existing ``registry`` instance. When provided,
+            models using this base share the mapper registry with other bases
+            that reference the same registry.
+
+    Returns:
+        A new declarative base class with ``TableSpecMixin`` mixed in.
+
+    Example:
+        >>> from rowsmyth import declarative_base
+        >>> Base = declarative_base()
+        >>> issubclass(Base, TableSpecMixin)
+        True
+    """
+    attrs: dict[str, Any] = {}
+    if metadata is not None:
+        attrs["metadata"] = metadata
+    if type_annotation_map is not None:
+        attrs["type_annotation_map"] = type_annotation_map
+    if registry is not None:
+        attrs["registry"] = registry
+    return type("Base", (TableSpecMixin, DeclarativeBase), attrs)

rowsmyth/_builder.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import random
+from typing import Any
+
+
+class FactoryBuilder:
+    """Fluent builder for generating and persisting model instances.
+
+    Obtain via ``Model.factory(count)``. Chain configuration methods,
+    then call ``.create()`` to generate and persist to an in-memory database.
+
+    Example:
+        >>> # users = User.factory(10).has(Order.factory(1, 3)).create()
+    """  # doctest: +SKIP
+
+    def __init__(
+        self, model: type, count: int = 1, max_count: int | None = None
+    ) -> None:
+        self.model = model
+        self.min_count = count
+        self.max_count = max_count
+        self._children: list[tuple[FactoryBuilder, str | None]] = []
+        self._mix: dict[str, float] = {}
+        self._overrides: dict[str, Any] = {}
+        self._seed: int | None = None
+
+    def has(self, *builders: FactoryBuilder, via: str | None = None) -> FactoryBuilder:
+        """Attach child builders whose rows are created for each parent instance.
+
+        Args:
+            *builders: ``FactoryBuilder`` instances for child models.
+            via: Relationship attribute name on the child model to use when
+                wiring the FK. Only required when the child has multiple
+                relationships pointing to the same parent model.
+
+        Returns:
+            ``self`` for chaining.
+
+        Example:
+            >>> # User.factory(5).has(Order.factory(1, 3))
+        """  # doctest: +SKIP
+        for builder in builders:
+            self._children.append((builder, via))
+        return self
+
+    def mix(self, **proportions: float) -> FactoryBuilder:
+        """Set variant proportions for this builder.
+
+        Args:
+            **proportions: Variant name to proportion in ``(0, 1]``. Total
+                must not exceed 1.0. The remainder (``1 - sum``) is the
+                proportion of un-varied (default) instances.
+
+        Returns:
+            ``self`` for chaining.
+
+        Raises:
+            ValueError: If proportions sum to more than 1.0, or a named
+                variant is not defined on the model.
+
+        Example:
+            >>> # User.factory(100).mix(admin=0.05, premium=0.20)
+        """  # doctest: +SKIP
+        for name, proportion in proportions.items():
+            if proportion <= 0.0:
+                msg = f"Proportion for '{name}' must be > 0, got {proportion}"
+                raise ValueError(msg)
+        total = sum(proportions.values())
+        if total > 1.0:
+            msg = f"Proportions sum to {total:.2f}, must be ≤ 1.0"
+            raise ValueError(msg)
+        for name in proportions:
+            if name not in self.model.__variants__:  # ty: ignore[unresolved-attribute]
+                msg = (
+                    f"Variant '{name}' not defined on {self.model.__name__}. "
+                    f"Defined: {list(self.model.__variants__)}"  # ty: ignore[unresolved-attribute]
+                )
+                raise ValueError(msg)
+        self._mix = proportions
+        return self
+
+    def _resolve_count(self) -> int:
+        if self.max_count is None:
+            return self.min_count
+        return random.randint(self.min_count, self.max_count)
+
+    def _pick_variant(self) -> str | None:
+        if not self._mix:
+            return None
+        r = random.random()
+        cumulative = 0.0
+        for name, proportion in self._mix.items():
+            cumulative += proportion
+            if r < cumulative:
+                return name
+        # Remainder proportion (1 - sum) maps to no variant. When proportions
+        # sum to exactly 1.0, floating-point imprecision may leave a tiny gap
+        # here for a negligible fraction of calls - that's acceptable.
+        return None
+
+    def random_seed(self, value: int) -> FactoryBuilder:
+        """Fix the random seed for reproducible output.
+
+        Args:
+            value: Integer seed passed to both ``random`` and ``Faker``.
+
+        Returns:
+            ``self`` for chaining.
+
+        Example:
+            >>> # User.factory(10).random_seed(42).create()
+        """  # doctest: +SKIP
+        self._seed = value
+        return self
+
+    def where(self, overrides: dict[Any, Any]) -> FactoryBuilder:
+        """Apply fixed column overrides to every generated instance.
+
+        Args:
+            overrides: Dict mapping column attributes (or string names) to
+                fixed values. These take precedence over generators and variants.
+
+        Returns:
+            ``self`` for chaining.
+
+        Example:
+            >>> # User.factory(5).where({User.tier: "premium"})
+        """  # doctest: +SKIP
+        self._overrides.update(overrides)
+        return self
+
+    def create(self) -> list[Any]:
+        """Generate and persist all instances to an in-memory SQLite database.
+
+        Returns:
+            List of root model instances. Child instances (added via ``.has()``)
+            are persisted and accessible via SQLAlchemy relationships but are
+            not included in the returned list.
+
+        Example:
+            >>> # users = User.factory(10).has(Order.factory(2)).create()
+            >>> # len(users)  # -> 10
+        """  # doctest: +SKIP
+        from ._execute import execute_builder
+
+        return execute_builder(self, self._overrides, self._seed)

rowsmyth/_dataset.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import graphlib
+import random
+from typing import Any
+
+from sqlalchemy.orm import MANYTOONE
+
+from ._execute import (
+    _build_fk_dependency_graph,
+    _resolve_variant,
+    _seed_random,
+    build_session,
+    make_factory,
+)
+
+
+class Dataset:
+    """Generates one instance per registered model, respecting FK order.
+
+    Obtain via ``Base.dataset(*builders)``. Use when you need a minimal
+    coherent dataset across all tables rather than a deep hierarchy rooted
+    at one model.
+
+    Example:
+        >>> # result = Base.dataset(User.factory(5), Order.factory(10)).create()
+        >>> # result["users"]   # -> 5 User instances
+        >>> # result["orders"]  # -> 10 Order instances
+    """  # doctest: +SKIP
+
+    def __init__(self, base: type, overrides: dict[type, Any] | None = None) -> None:
+        self._base = base
+        self._overrides: dict[type, Any] = overrides or {}
+        self._seed: int | None = None
+
+    def random_seed(self, value: int) -> Dataset:
+        """Fix the random seed for reproducible output.
+
+        Args:
+            value: Integer seed passed to both ``random`` and ``Faker``.
+
+        Returns:
+            ``self`` for chaining.
+        """
+        self._seed = value
+        return self
+
+    def create(self) -> dict[str, list[Any]]:
+        """Generate and persist one instance per registered model.
+
+        Instances are created in FK dependency order (parents before children).
+        FK columns are wired automatically by sampling randomly from the parent pool.
+
+        Returns:
+            Dict keyed by table name mapping to lists of persisted instances.
+
+        Raises:
+            ValueError: If a parent model has 0 instances but a child model
+                has an FK pointing to it.
+        """
+        from ._builder import FactoryBuilder
+
+        if self._seed is not None:
+            _seed_random(self._seed)
+
+        builders = []
+        for mapper in self._base.registry.mappers:  # ty: ignore[unresolved-attribute]
+            model = mapper.class_
+            builders.append(self._overrides.get(model, FactoryBuilder(model, 1)))
+
+        ordered = self._topo_sort(builders)
+        models = {b.model for b in ordered}
+        session = build_session(models)
+        factories = {m: make_factory(m, session) for m in models}
+
+        pool: dict[type, list[Any]] = {}
+        result: dict[str, list[Any]] = {}
+
+        for builder in ordered:
+            instances: list[Any] = []
+            for _ in range(builder._resolve_count()):
+                fk_overrides: dict[str, Any] = {}
+                for rel in builder.model.__mapper__.relationships:
+                    if rel.direction is not MANYTOONE:
+                        continue
+                    related = rel.mapper.class_
+                    if related not in pool:
+                        continue
+                    if not pool[related]:
+                        model_name = builder.model.__name__
+                        related_name = related.__name__
+                        msg = (
+                            f"Cannot wire {model_name}.{rel.key}: "
+                            f"{related_name} has 0 instances in the dataset. "
+                            f"Use {related_name}.factory(0) only for models "
+                            f"with no FK dependents."
+                        )
+                        raise ValueError(msg)
+                    fk_overrides[rel.key] = random.choice(pool[related])
+
+                inst_overrides = _resolve_variant(builder)
+                inst_overrides.update(fk_overrides)
+
+                instance = factories[builder.model](**inst_overrides)
+                instances.append(instance)
+
+            pool[builder.model] = instances
+            result[builder.model.__tablename__] = instances
+
+        return result
+
+    def _topo_sort(self, builders: list[Any]) -> list[Any]:
+        model_to_builder = {b.model: b for b in builders}
+        graph = _build_fk_dependency_graph(builders)
+        ts = graphlib.TopologicalSorter(graph)
+        ordered_models = list(ts.static_order())
+        return [model_to_builder[m] for m in ordered_models if m in model_to_builder]

rowsmyth/_execute.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+import functools
+import graphlib
+import random
+from typing import TYPE_CHECKING, Any
+
+from factory.alchemy import SQLAlchemyModelFactory
+from sqlalchemy import Table, create_engine
+from sqlalchemy.orm import MANYTOONE, Session
+
+if TYPE_CHECKING:
+    from ._builder import FactoryBuilder
+
+
+def _seed_random(seed: int) -> None:
+    from faker import Faker
+
+    random.seed(seed)
+    Faker.seed(seed)
+
+
+def check_cycles(builder: FactoryBuilder) -> None:
+    graph: dict[type, set[type]] = {}
+
+    def _collect(b: FactoryBuilder) -> None:
+        if b.model in graph:
+            return
+        graph[b.model] = set()
+        for child_builder, _ in b._children:
+            graph[b.model].add(child_builder.model)
+            _collect(child_builder)
+
+    _collect(builder)
+    ts = graphlib.TopologicalSorter(graph)
+    ts.prepare()
+
+
+def _build_fk_dependency_graph(builders: list[FactoryBuilder]) -> dict[type, set[type]]:
+    model_set = {b.model for b in builders}
+    return {
+        b.model: {
+            rel.mapper.class_
+            for rel in b.model.__mapper__.relationships  # ty: ignore[unresolved-attribute]
+            if rel.direction is MANYTOONE
+            and rel.mapper.class_ in model_set
+            and rel.mapper.class_ is not b.model
+        }
+        for b in builders
+    }
+
+
+@functools.cache
+def _get_default_builders(model: type) -> dict[str, FactoryBuilder]:
+    from ._builder import FactoryBuilder
+
+    return {
+        getattr(key, "key", key): val
+        for key, val in model.generators().items()  # ty: ignore[unresolved-attribute]
+        if isinstance(val, FactoryBuilder)
+    }
+
+
+def collect_models(builder: FactoryBuilder) -> set[type]:
+    models: set[type] = set()
+
+    def _collect(b: FactoryBuilder) -> None:
+        models.add(b.model)
+        for child_builder, _ in b._children:
+            _collect(child_builder)
+        for default_builder in _get_default_builders(b.model).values():
+            _collect(default_builder)
+
+    _collect(builder)
+    return models
+
+
+def collect_all_tables(models: set[type]) -> set[Table]:
+    seen: set[Table] = {m.__table__ for m in models}  # ty: ignore[unresolved-attribute]
+    queue: list[Table] = list(seen)
+    while queue:
+        table = queue.pop()
+        for fk in table.foreign_key_constraints:
+            for col in fk.elements:
+                ref_table = col.column.table
+                if ref_table not in seen:
+                    seen.add(ref_table)
+                    queue.append(ref_table)
+    return seen
+
+
+def _resolve_variant(builder: FactoryBuilder) -> dict[str, Any]:
+    variant_name = builder._pick_variant()
+    if not variant_name:
+        return {}
+    raw = builder.model.__variants__[variant_name](builder.model)  # ty: ignore[unresolved-attribute]
+    return _resolve_attr_names(raw)
+
+
+def _resolve_attr_names(d: dict[Any, Any]) -> dict[str, Any]:
+    result: dict[str, Any] = {}
+    for key, val in d.items():
+        attr_name = getattr(key, "key", key)
+        result[attr_name] = val
+    return result
+
+
+def build_session(models: set[type]) -> Session:
+    tables = collect_all_tables(models)
+    metadata = next(iter(models)).__table__.metadata  # ty: ignore[unresolved-attribute]
+    engine = create_engine("sqlite:///:memory:")
+    metadata.create_all(engine, tables=list(tables))
+    return Session(engine)
+
+
+def make_factory(model: type, session: Session) -> type[SQLAlchemyModelFactory]:
+    from ._builder import FactoryBuilder
+
+    raw_generators = {
+        k: v
+        for k, v in model.generators().items()  # ty: ignore[unresolved-attribute]
+        if not isinstance(v, FactoryBuilder)
+    }
+    generators = _resolve_attr_names(raw_generators)
+    meta = type(
+        "Meta",
+        (),
+        {
+            "model": model,
+            "sqlalchemy_session": session,
+            "sqlalchemy_session_persistence": "commit",
+        },
+    )
+    attrs: dict[str, Any] = {"Meta": meta, **generators}
+    return type(f"{model.__name__}Factory", (SQLAlchemyModelFactory,), attrs)
+
+
+@functools.cache
+def find_relationship(
+    child_model: type, parent_model: type, via: str | None = None
+) -> str:
+    if via is not None:
+        return via
+    matches = [
+        rel.key
+        for rel in child_model.__mapper__.relationships  # ty: ignore[unresolved-attribute]
+        if rel.mapper.class_ is parent_model
+    ]
+    if not matches:
+        msg = (
+            f"No relationship from {child_model.__name__} to {parent_model.__name__}. "
+            f"Define one via SQLAlchemy relationship() or pass via= to .has()"
+        )
+        raise ValueError(msg)
+    if len(matches) > 1:
+        msg = (
+            f"Ambiguous - found {matches} on {child_model.__name__} pointing to "
+            f"{parent_model.__name__}. Pass via='rel_name' to .has()"
+        )
+        raise ValueError(msg)
+    return matches[0]
+
+
+def _create_instance(
+    builder: FactoryBuilder,
+    extra_overrides: dict[str, Any],
+    factories: dict[type, type[SQLAlchemyModelFactory]],
+) -> Any:
+    inst_overrides = _resolve_variant(builder)
+    inst_overrides.update(extra_overrides)
+
+    default_builders = _get_default_builders(builder.model)
+    for rel in builder.model.__mapper__.relationships:  # ty: ignore[unresolved-attribute]
+        if rel.direction is not MANYTOONE:
+            continue
+        if rel.key not in inst_overrides and rel.key in default_builders:
+            inst_overrides[rel.key] = _create_instance(
+                default_builders[rel.key], {}, factories
+            )
+
+    instance = factories[builder.model](**inst_overrides)
+
+    for child_builder, via in builder._children:
+        rel_key = find_relationship(child_builder.model, builder.model, via)
+        for _ in range(child_builder._resolve_count()):
+            _create_instance(child_builder, {rel_key: instance}, factories)
+
+    return instance
+
+
+def execute_builder(
+    builder: FactoryBuilder,
+    overrides: dict[Any, Any],
+    seed: int | None = None,
+) -> list[Any]:
+    if seed is not None:
+        _seed_random(seed)
+
+    check_cycles(builder)
+    models = collect_models(builder)
+    session = build_session(models)
+    factories = {m: make_factory(m, session) for m in models}
+    root_overrides = _resolve_attr_names(overrides)
+
+    return [
+        _create_instance(builder, root_overrides, factories)
+        for _ in range(builder._resolve_count())
+    ]