Source code for codegen_database.types.numeric

"""Numeric / decimal column types for codegen_database.

``"numeric"`` and ``"decimal"`` are synonyms for
:class:`~sqlalchemy.Numeric` (PostgreSQL ``NUMERIC``); ``"decimal"``
exists because that is the name most callers reach for when they want
an exact fixed-point amount with an explicit precision/scale.
``"integer"`` stays the minor-unit (e.g. cents) default.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal

from sqlalchemy import Integer, Numeric

from codegen_database.errors import CodegenDatabaseValidationError

if TYPE_CHECKING:
    from sqlalchemy.types import TypeEngine

ValueType = Literal["integer", "numeric", "decimal"]
"""Named scalar column types accepted by the ledger factories.

``"numeric"`` and ``"decimal"`` are synonyms for
:class:`~sqlalchemy.Numeric`; ``"integer"`` is the minor-unit
(e.g. cents) default.  Used to type the ``value_type`` argument of
:class:`~codegen_database.factory.ledger.LedgerTablePlugin`,
:class:`~codegen_database.plugins.snapshot.LedgerSnapshotPlugin`,
and :func:`build_value_type`.
"""


VALUE_TYPES: dict[ValueType, type[TypeEngine]] = {
    "integer": Integer,
    "numeric": Numeric,
    "decimal": Numeric,
}
"""Maps each :data:`ValueType` to its SQLAlchemy type class.

``"numeric"`` and ``"decimal"`` both resolve to
:class:`~sqlalchemy.Numeric`.  Used to type the ledger ``value`` /
snapshot ``balance`` columns.
"""




def build_value_type(
    value_type: ValueType,
    *,
    precision: int | None = None,
    scale: int | None = None,
) -> TypeEngine:
    """Instantiate the SQLAlchemy type for a named value column.

    Args:
        value_type: One of :data:`ValueType`.
        precision: Total number of digits for a ``NUMERIC`` column
            (the ``NUMERIC(precision, scale)`` first argument).  Only
            valid for the ``"numeric"`` / ``"decimal"`` types.
        scale: Number of digits after the decimal point.  Requires
            *precision* -- ``NUMERIC`` cannot fix a scale without a
            precision.

    Returns:
        A configured :class:`~sqlalchemy.types.TypeEngine` instance
        ready to drop into a :class:`~sqlalchemy.Column`.

    Raises:
        CodegenDatabaseValidationError: If *value_type* is unknown,
            if precision/scale are supplied for a non-``NUMERIC``
            type, or if *scale* is given without *precision*.

    """
    if value_type not in VALUE_TYPES:
        msg = (
            f"Unknown value_type {value_type!r}. "
            f"Must be one of: {sorted(VALUE_TYPES)}"
        )
        raise CodegenDatabaseValidationError(msg)

    sa_type = VALUE_TYPES[value_type]

    if precision is None and scale is None:
        return sa_type()

    if sa_type is not Numeric:
        msg = (
            f"precision/scale are only valid for the 'numeric' / "
            f"'decimal' value types, not {value_type!r}."
        )
        raise CodegenDatabaseValidationError(msg)

    if precision is None:
        msg = (
            "scale requires precision "
            "(NUMERIC needs a precision to fix a scale)."
        )
        raise CodegenDatabaseValidationError(msg)

    return Numeric(precision, scale)