Coverage for polar/oauth2/sub_type.py: 50%

1from enum import StrEnum 2 ctx1ab

2from typing import TYPE_CHECKING, Literal, TypeGuard 2 ctx1ab

3from uuid import UUID 2 ctx1ab

4 

5from sqlalchemy import ForeignKey, String, Uuid 2 ctx1ab

6from sqlalchemy.ext.hybrid import hybrid_property 2 ctx1ab

7from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship 2 ctx1ab

8 

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

10 from polar.models import Organization, User 

11 

12 

13class SubType(StrEnum): 2 ctx1ab

14 user = "user" 2 ctx1ab

15 organization = "organization" 2 ctx1ab

16 

17 

18SubTypeValue = tuple[SubType, "User | Organization"] 2 ctx1ab

19 

20 

21def is_sub_user(v: SubTypeValue) -> TypeGuard[tuple[Literal[SubType.user], "User"]]: 2 ctx1ab

22 return v[0] == SubType.user 

23 

24 

25def is_sub_organization( 2 ctx1ab

26 v: SubTypeValue, 

27) -> TypeGuard[tuple[Literal[SubType.organization], "Organization"]]: 

28 return v[0] == SubType.organization 

29 

30 

31class SubTypeModelMixin: 2 ctx1ab

32 sub_type: Mapped[SubType] = mapped_column(String, nullable=False) 2 ctx1ab

33 user_id: Mapped[UUID | None] = mapped_column( 2 ctx1ab

34 Uuid, ForeignKey("users.id", ondelete="cascade"), nullable=True 

35 ) 

36 organization_id: Mapped[UUID | None] = mapped_column( 2 ctx1ab

37 Uuid, ForeignKey("organizations.id", ondelete="cascade"), nullable=True 

38 ) 

39 

40 @declared_attr 2 ctx1ab

41 def user(cls) -> Mapped["User | None"]: 2 ctx1ab

42 return relationship("User", lazy="joined") 2 ctx1ab

43 

44 @declared_attr 2 ctx1ab

45 def organization(cls) -> Mapped["Organization | None"]: 2 ctx1ab

46 return relationship("Organization", lazy="joined") 2 ctx1ab

47 

48 @hybrid_property 2 ctx1ab

49 def sub(self) -> "User | Organization": 2 ctx1ab

50 sub: User | Organization | None = None 

51 if self.sub_type == SubType.user: 

52 sub = self.user 

53 elif self.sub_type == SubType.organization: 

54 sub = self.organization 

55 else: 

56 raise NotImplementedError() 

57 

58 if sub is None: 

59 raise ValueError("Sub is not found.") 

60 

61 return sub 

62 

63 @sub.inplace.setter 2 ctx1ab

64 def _sub_setter(self, value: "User | Organization") -> None: 2 ctx1ab

65 if self.sub_type == SubType.user: 

66 self.user_id = value.id 

67 elif self.sub_type == SubType.organization: 

68 self.organization_id = value.id 

69 else: 

70 raise NotImplementedError() 

71 

72 def get_sub_type_value(self) -> SubTypeValue: 2 ctx1ab

73 return self.sub_type, self.sub