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

1from enum import StrEnum 

2from typing import TYPE_CHECKING 

3from uuid import UUID 

4 

5from sqlalchemy import Boolean, ForeignKey, Text, Uuid 

6from sqlalchemy.dialects.postgresql import JSONB 

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

8 

9from polar.exceptions import PolarError 

10from polar.kit.db.models import RecordModel 

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

12from polar.kit.metadata import MetadataMixin 

13 

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

15 from polar.benefit.strategies import BenefitProperties 

16 from polar.models import BenefitGrant, Organization 

17 

18 

19class TaxApplicationMustBeSpecified(PolarError): 

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

21 self.type = type 

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

23 super().__init__(message) 

24 

25 

26class BenefitType(StrEnum): 

27 custom = "custom" 

28 discord = "discord" 

29 github_repository = "github_repository" 

30 downloadables = "downloadables" 

31 license_keys = "license_keys" 

32 meter_credit = "meter_credit" 

33 

34 def get_display_name(self) -> str: 

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: 

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): 

60 __tablename__ = "benefits" 

61 

62 type: Mapped[BenefitType] = mapped_column( 

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

64 ) 

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

66 is_tax_applicable: Mapped[bool] = mapped_column( 

67 Boolean, nullable=False, default=False 

68 ) 

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

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

71 organization_id: Mapped[UUID] = mapped_column( 

72 Uuid, 

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

74 nullable=False, 

75 index=True, 

76 ) 

77 

78 @declared_attr 

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

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

81 

82 @declared_attr 

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

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

85 

86 @declared_attr 

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

88 return relationship( 

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

90 )