Source code for codegen_database.types.enums

"""Enum-backed SQLAlchemy column types for codegen_database."""

from __future__ import annotations

import enum
from typing import Any

from sqlalchemy import Integer, Text
from sqlalchemy import types as sa_types




class TextEnum(sa_types.TypeDecorator[enum.Enum]):
    """Store a Python enum as plain text in PostgreSQL.

    Values are persisted as the enum member's ``.value`` (which
    must be a string) and coerced back to the corresponding Python
    enum member on load.  No PostgreSQL ``ENUM`` type is created —
    the underlying column is ``TEXT``.

    Args:
        enum_class: The :class:`enum.Enum` subclass to coerce
            values to and from.

    Example::

        import enum
        from sqlalchemy import Column
        from codegen_database.types import TextEnum

        class Color(enum.Enum):
            RED = "red"
            GREEN = "green"
            BLUE = "blue"

        Column("color", TextEnum(Color), nullable=False)

    """

    impl = Text
    cache_ok = True

    def __init__(self, enum_class: type[enum.Enum]) -> None:
        """Store the enum class for value coercion.

        Accepts any object to surface a clear TypeError at construction
        time for callers who bypass the type annotation.
        """
        if not (
            isinstance(enum_class, type) and issubclass(enum_class, enum.Enum)
        ):
            msg = (  # type: ignore[unreachable]
                f"TextEnum expects an Enum subclass, got {enum_class!r}"
            )
            raise TypeError(msg)

        self.enum_class = enum_class
        super().__init__()



    def process_bind_param(
        self,
        value: Any,  # noqa: ANN401
        dialect: Any,  # noqa: ANN401, ARG002
    ) -> str | None:
        """Convert a Python enum member to its string value."""
        if value is None:
            return None

        if isinstance(value, self.enum_class):
            return value.value

        if isinstance(value, str):
            # Value-to-member lookup; raises ValueError if unknown.
            self.enum_class(value)  # type: ignore[misc]
            return value

        msg = (
            f"Expected {self.enum_class.__name__} or str, "
            f"got {type(value).__name__}"
        )
        raise ValueError(msg)




    def process_result_value(
        self,
        value: Any,  # noqa: ANN401
        dialect: Any,  # noqa: ANN401, ARG002
    ) -> enum.Enum | None:
        """Convert a stored text value back to a Python enum."""
        if value is None:
            return None

        return self.enum_class(value)  # type: ignore[misc]






class IntEnum(sa_types.TypeDecorator[enum.Enum]):
    """Store a Python enum as an integer in PostgreSQL.

    Values are persisted as the enum member's ``.value`` (which
    must be an integer) and coerced back to the corresponding
    Python enum member on load.  The underlying column is
    ``INTEGER``.

    Args:
        enum_class: The :class:`enum.Enum` subclass to coerce
            values to and from.  All member values must be
            integers.

    Example::

        import enum
        from sqlalchemy import Column
        from codegen_database.types import IntEnum

        class Priority(enum.Enum):
            LOW = 1
            MEDIUM = 2
            HIGH = 3

        Column("priority", IntEnum(Priority), nullable=False)

    """

    impl = Integer
    cache_ok = True

    def __init__(self, enum_class: type[enum.Enum]) -> None:
        """Store the enum class for value coercion.

        Accepts any object to surface a clear TypeError at construction
        time for callers who bypass the type annotation.
        """
        if not (
            isinstance(enum_class, type) and issubclass(enum_class, enum.Enum)
        ):
            msg = (  # type: ignore[unreachable]
                f"IntEnum expects an Enum subclass, got {enum_class!r}"
            )
            raise TypeError(msg)

        self.enum_class = enum_class
        super().__init__()



    def process_bind_param(
        self,
        value: Any,  # noqa: ANN401
        dialect: Any,  # noqa: ANN401, ARG002
    ) -> int | None:
        """Convert a Python enum member to its integer value."""
        if value is None:
            return None

        if isinstance(value, self.enum_class):
            return value.value

        if isinstance(value, int):
            # Value-to-member lookup; raises ValueError if unknown.
            self.enum_class(value)  # type: ignore[misc]
            return value

        msg = (
            f"Expected {self.enum_class.__name__} or int, "
            f"got {type(value).__name__}"
        )
        raise ValueError(msg)




    def process_result_value(
        self,
        value: Any,  # noqa: ANN401
        dialect: Any,  # noqa: ANN401, ARG002
    ) -> enum.Enum | None:
        """Convert a stored integer back to a Python enum."""
        if value is None:
            return None

        return self.enum_class(value)  # type: ignore[misc]