Coverage for polar/models/benefit.py: 77%

1from enum import StrEnum 2 ctx1ab

2from typing import TYPE_CHECKING 2 ctx1ab

3from uuid import UUID 2 ctx1ab

4 

5from sqlalchemy import Boolean, ForeignKey, Text, Uuid 2 ctx1ab

6from sqlalchemy.dialects.postgresql import JSONB 2 ctx1ab

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

8 

9from polar.exceptions import PolarError 2 ctx1ab

10from polar.kit.db.models import RecordModel 2 ctx1ab

11from polar.kit.extensions.sqlalchemy.types import StringEnum 2 ctx1ab

12from polar.kit.metadata import MetadataMixin 2 ctx1ab

13 

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

15 from polar.benefit.strategies import BenefitProperties 

16 from polar.models import BenefitGrant, Organization 

17 

18 

19class TaxApplicationMustBeSpecified(PolarError): 2 ctx1ab

20 def __init__(self, type: "BenefitType") -> None: 2 ctx1ab

21 self.type = type 

22 message = "The tax application should be specified for this type." 

23 super().__init__(message) 

24 

25 

26class BenefitType(StrEnum): 2 ctx1ab

27 custom = "custom" 2 ctx1ab

28 discord = "discord" 2 ctx1ab

29 github_repository = "github_repository" 2 ctx1ab

30 downloadables = "downloadables" 2 ctx1ab

31 license_keys = "license_keys" 2 ctx1ab

32 meter_credit = "meter_credit" 2 ctx1ab

33 

34 def get_display_name(self) -> str: 2 ctx1ab

35 return { 

36 BenefitType.custom: "Custom", 

37 BenefitType.discord: "Discord", 

38 BenefitType.github_repository: "GitHub Repository", 

39 BenefitType.downloadables: "Downloadables", 

40 BenefitType.license_keys: "License Keys", 

41 BenefitType.meter_credit: "Meter Credit", 

42 }[self] 

43 

44 def is_tax_applicable(self) -> bool: 2 ctx1ab

45 try: 

46 _is_tax_applicable_map: dict[BenefitType, bool] = { 

47 BenefitType.custom: True, 

48 BenefitType.discord: True, 

49 BenefitType.github_repository: True, 

50 BenefitType.downloadables: True, 

51 BenefitType.license_keys: True, 

52 BenefitType.meter_credit: True, 

53 } 

54 return _is_tax_applicable_map[self] 

55 except KeyError as e: 

56 raise TaxApplicationMustBeSpecified(self) from e 

57 

58 

59class Benefit(MetadataMixin, RecordModel): 2 ctx1ab

60 __tablename__ = "benefits" 2 ctx1ab

61 

62 type: Mapped[BenefitType] = mapped_column( 2 ctx1ab

63 StringEnum(BenefitType), nullable=False, index=True 

64 ) 

65 description: Mapped[str] = mapped_column(Text, nullable=False) 2 ctx1ab

66 is_tax_applicable: Mapped[bool] = mapped_column( 2 ctx1ab

67 Boolean, nullable=False, default=False 

68 ) 

69 selectable: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True) 2 ctx1ab

70 deletable: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True) 2 ctx1ab

71 organization_id: Mapped[UUID] = mapped_column( 2 ctx1ab

72 Uuid, 

73 ForeignKey("organizations.id", ondelete="cascade"), 

74 nullable=False, 

75 index=True, 

76 ) 

77 

78 @declared_attr 2 ctx1ab

79 def properties(cls) -> Mapped["BenefitProperties"]: 2 ctx1ab

80 return mapped_column("properties", JSONB, nullable=False, default=dict) 2 ctx1ab

81 

82 @declared_attr 2 ctx1ab

83 def organization(cls) -> Mapped["Organization"]: 2 ctx1ab

84 return relationship("Organization", lazy="raise") 2 ctx1ab

85 

86 @declared_attr 2 ctx1ab

87 def grants(cls) -> Mapped["list[BenefitGrant]"]: 2 ctx1ab

88 return relationship( 2 ctx1ab

89 "BenefitGrant", lazy="raise", back_populates="benefit", viewonly=True 

90 )