Source code for simdb.database.models.types

import enum
import uuid
from typing import Any, Dict, Optional

from sqlalchemy import types as sql_types
from sqlalchemy.dialects import postgresql

from simdb.imas.utils import SimDBUrl



[docs]
class UUID(sql_types.TypeDecorator):
    """
    Platform-independent GUID type.

    Uses PostgreSQL's UUID type, otherwise uses CHAR(32), storing as stringified hex
    values.
    """

    impl = sql_types.CHAR

    cache_ok = True

    @property
    def python_type(self):
        return uuid.UUID


[docs]
    def load_dialect_impl(self, dialect):
        if dialect.name == "postgresql":
            return dialect.type_descriptor(postgresql.UUID())
        else:
            return dialect.type_descriptor(sql_types.CHAR(32))



[docs]
    def process_bind_param(self, value, dialect):
        if value is None:
            return value
        elif dialect.name == "postgresql":
            return str(value)
        elif not isinstance(value, uuid.UUID):
            try:
                return uuid.UUID(value).hex
            except ValueError:
                return value
        else:
            return value.hex



[docs]
    def process_result_value(self, value, dialect):
        if value is None:
            return value
        else:
            if not isinstance(value, uuid.UUID):
                value = uuid.UUID(value)
            return value



[docs]
    def process_literal_param(self, value, dialect):
        return self.process_result_value(value, dialect)





[docs]
class URI(sql_types.TypeDecorator):
    """
    UUID type for reading/writing to the ORM.
    """

    impl = sql_types.VARCHAR

    @property
    def python_type(self):
        return SimDBUrl


[docs]
    def process_bind_param(self, value: Optional[SimDBUrl], dialect) -> Optional[str]:
        if value is None:
            return value
        return str(value)



[docs]
    def process_result_value(self, value: Optional[str], dialect) -> Optional[SimDBUrl]:
        if value is None:
            return value
        return SimDBUrl(value)



[docs]
    def process_literal_param(self, value, dialect) -> Optional[SimDBUrl]:
        return self.process_result_value(value, dialect)





[docs]
class ChoiceType(sql_types.TypeDecorator):
    impl = sql_types.CHAR

    @property
    def python_type(self):
        return str

    def __init__(self, choices: Dict[Any, str], enum_type: type, **kw):
        if type(enum_type) is not enum.EnumMeta:
            raise ValueError("enum_type must be a class inheriting from enum.Enum.")
        self._enum_type = enum_type
        self._choices_inverse = dict(choices)
        self._choices = {v: k for k, v in self._choices_inverse.items()}
        if len(self._choices) != len(self._choices_inverse):
            raise TypeError("Values in choices dict must be unique")
        super().__init__(**kw)


[docs]
    def process_bind_param(self, value: str, dialect):
        return self._choices_inverse[self._enum_type(value)]



[docs]
    def process_result_value(self, value: str, dialect):
        return self._choices[value]



[docs]
    def process_literal_param(self, value, dialect):
        return self.process_result_value(value, dialect)