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

1from enum import Enum, StrEnum 2 ctx1ab

2from typing import TYPE_CHECKING, Any 2 ctx1ab

3 

4import sqlalchemy as sa 2 ctx1ab

5from sqlalchemy.dialects.postgresql import UUID 2 ctx1ab

6from sqlalchemy.engine.interfaces import Dialect 2 ctx1ab

7from sqlalchemy.types import TypeDecorator as _TypeDecorator 2 ctx1ab

8 

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

10 GUIDTypeDecorator = _TypeDecorator[UUID] # type: ignore 

11 TypeDecorator = _TypeDecorator[Any] 

12else: 

13 GUIDTypeDecorator = _TypeDecorator 2 ctx1ab

14 TypeDecorator = _TypeDecorator 2 ctx1ab

15 

16 

17class EnumType(TypeDecorator): 2 ctx1ab

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

19 super().__init__(**kwargs) 3 ctx1acb

20 self.enum_klass = enum_klass 3 ctx1acb

21 

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

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

24 

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

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

27 

28 

29class IntEnum(EnumType): 2 ctx1ab

30 impl = sa.Integer 2 ctx1ab

31 cache_ok = True 2 ctx1ab

32 

33 

34class StringEnum(EnumType): 2 ctx1ab

35 impl = sa.Unicode 2 ctx1ab

36 cache_ok = True 2 ctx1ab

37 

38 

39class StrEnumType(TypeDecorator): 2 ctx1ab

40 impl = sa.String 2 ctx1ab

41 cache_ok = True 2 ctx1ab

42 

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

44 super().__init__(**kwargs) 2 ctx1ab

45 self.enum_klass = enum_klass 2 ctx1ab

46 

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

48 if isinstance(value, self.enum_klass): 

49 return str(value) 

50 return value 

51 

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

53 if value is not None: 

54 return self.enum_klass(value) 

55 return value