Coverage for polar/kit/extensions/sqlalchemy/types.py: 78%

1from enum import Enum, StrEnum 

2from typing import TYPE_CHECKING, Any 

3 

4import sqlalchemy as sa 

5from sqlalchemy.dialects.postgresql import UUID 

6from sqlalchemy.engine.interfaces import Dialect 

7from sqlalchemy.types import TypeDecorator as _TypeDecorator 

8 

9if TYPE_CHECKING: # pragma: no cover 9 ↛ 10line 9 didn't jump to line 10 because the condition on line 9 was never true

10 GUIDTypeDecorator = _TypeDecorator[UUID] # type: ignore 

11 TypeDecorator = _TypeDecorator[Any] 

12else: 

13 GUIDTypeDecorator = _TypeDecorator 

14 TypeDecorator = _TypeDecorator 

15 

16 

17class EnumType(TypeDecorator): 

18 def __init__(self, enum_klass: type[Enum], **kwargs: Any) -> None: 

19 super().__init__(**kwargs) 

20 self.enum_klass = enum_klass 

21 

22 def process_bind_param(self, value: Any, dialect: Dialect) -> Any: 

23 return value.value if isinstance(value, self.enum_klass) else value 

24 

25 def process_result_value(self, value: Any, dialect: Dialect) -> Any: 

26 return value if value is None else self.enum_klass(value) 

27 

28 

29class IntEnum(EnumType): 

30 impl = sa.Integer 

31 cache_ok = True 

32 

33 

34class StringEnum(EnumType): 

35 impl = sa.Unicode 

36 cache_ok = True 

37 

38 

39class StrEnumType(TypeDecorator): 

40 impl = sa.String 

41 cache_ok = True 

42 

43 def __init__(self, enum_klass: type[StrEnum], **kwargs: Any) -> None: 

44 super().__init__(**kwargs) 

45 self.enum_klass = enum_klass 

46 

47 def process_bind_param(self, value: Any, dialect: Dialect) -> Any: 

48 if isinstance(value, self.enum_klass): 48 ↛ 50line 48 didn't jump to line 50 because the condition on line 48 was always true

49 return str(value) 

50 return value 

51 

52 def process_result_value(self, value: Any, dialect: Dialect) -> Any: 

53 if value is not None: 

54 return self.enum_klass(value) 

55 return value