Source code for codegen_database.types.postgis

from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from typing import Any

from sqlalchemy import cast, func, literal_column
from sqlalchemy import types as sa_types
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.types import UserDefinedType




@dataclass(frozen=True, kw_only=True)
class StdAddr:
    building: str = ""
    house_num: str = ""
    predir: str = ""
    qual: str = ""
    pretype: str = ""
    name: str = ""
    suftype: str = ""
    sufdir: str = ""
    ruralroute: str = ""
    extra: str = ""
    city: str = ""
    state: str = ""
    country: str = ""
    postcode: str = ""
    box: str = ""
    unit: str = ""



class _StdAddrColumn(UserDefinedType[StdAddr]):
    cache_ok = True

    def get_col_spec(self, **_kw: Any) -> str:  # noqa: ANN401
        return "stdaddr"




class STDADDR(sa_types.TypeDecorator[StdAddr]):
    impl = _StdAddrColumn
    cache_ok = True



    def column_expression(self, colexpr: Any) -> Any:  # noqa: ANN401
        return func.to_jsonb(colexpr, type_=self)




    def process_result_value(
        self,
        value: Any,  # noqa: ANN401
        dialect: Any,  # noqa: ANN401, ARG002
    ) -> StdAddr | None:
        if value is None:
            return None

        return StdAddr(**{k: (v or "") for k, v in value.items()})




    def bind_expression(self, bindvalue: Any) -> Any:  # noqa: ANN401
        return func.jsonb_populate_record(
            literal_column("NULL::stdaddr"),
            cast(bindvalue, JSONB),
        )




    def process_bind_param(
        self,
        value: Any,  # noqa: ANN401
        dialect: Any,  # noqa: ANN401, ARG002
    ) -> str | None:
        if value is None:
            return None

        if isinstance(value, StdAddr):
            return json.dumps(
                {k: (v or None) for k, v in asdict(value).items()}
            )

        msg = f"Expected StdAddr, got {type(value).__name__}"
        raise ValueError(msg)