Coverage for polar/custom_field/data.py: 69%

1import datetime 2 ctx1ab

2from collections.abc import Sequence 2 ctx1ab

3from typing import TYPE_CHECKING, Any 2 ctx1ab

4 

5from pydantic import BaseModel, Field, ValidationError, create_model 2 ctx1ab

6from sqlalchemy import event 2 ctx1ab

7from sqlalchemy.dialects.postgresql import JSONB 2 ctx1ab

8from sqlalchemy.orm import Mapped, Mapper, ORMDescriptor, mapped_column 2 ctx1ab

9 

10from polar.exceptions import PolarRequestValidationError 2 ctx1ab

11 

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

13 from polar.models import CustomField, Organization 

14 

15 from .attachment import AttachedCustomFieldMixin 

16 

17 

18class CustomFieldDataMixin: 2 ctx1ab

19 custom_field_data: Mapped[ 2 ctx1ab

20 dict[str, str | int | bool | datetime.datetime | None] 

21 ] = mapped_column(JSONB, nullable=False, default=dict) 

22 

23 # Make the type checker happy, but we should make sure actual models 

24 # declare an organization attribute by themselves. 

25 if TYPE_CHECKING: 25 ↛ anywhereline 25 didn't jump anywhere: it always raised an exception.2 ctx1ab

26 organization: ORMDescriptor["Organization"] 2 ctx1ab

27 

28 

29custom_field_data_models: set[type[CustomFieldDataMixin]] = set() 2 ctx1ab

30 

31 

32# Event listener to track models inheriting from CustomFieldDataMixin 

33@event.listens_for(Mapper, "mapper_configured") 2 ctx1ab

34def track_attached_custom_field_mixin(_mapper: Mapper[Any], class_: type) -> None: 2 ctx1ab

35 if issubclass(class_, CustomFieldDataMixin): 1 ctx1c

36 custom_field_data_models.add(class_) 1 ctx1c

37 

38 

39class CustomFieldDataInputMixin(BaseModel): 2 ctx1ab

40 custom_field_data: dict[str, str | int | bool | datetime.datetime | None] = Field( 2 ctx1ab

41 default_factory=dict, 

42 description="Key-value object storing custom field values.", 

43 ) 

44 

45 

46class CustomFieldDataOutputMixin(BaseModel): 2 ctx1ab

47 custom_field_data: dict[str, str | int | bool | datetime.datetime | None] = Field( 2 ctx1ab

48 default_factory=dict, 

49 description="Key-value object storing custom field values.", 

50 ) 

51 

52 

53def build_custom_field_data_schema( 2 ctx1ab

54 custom_fields: Sequence[tuple["CustomField", bool]], 

55) -> type[BaseModel]: 

56 fields_definitions: Any = { 

57 custom_field.slug: custom_field.get_field_definition(required) 

58 for custom_field, required in custom_fields 

59 } 

60 return create_model("CustomFieldDataInput", **fields_definitions) 

61 

62 

63def validate_custom_field_data( 2 ctx1ab

64 attached_custom_fields: Sequence["AttachedCustomFieldMixin"], 

65 data: dict[str, Any], 

66 *, 

67 error_loc_prefix: Sequence[str] = ("body", "custom_field_data"), 

68 validate_required: bool = True, 

69) -> dict[str, Any]: 

70 schema = build_custom_field_data_schema( 

71 [ 

72 (f.custom_field, validate_required and f.required) 

73 for f in attached_custom_fields 

74 ] 

75 ) 

76 try: 

77 return schema.model_validate(data).model_dump(mode="json") 

78 except ValidationError as e: 

79 raise PolarRequestValidationError( 

80 [{**err, "loc": (*error_loc_prefix, *err["loc"])} for err in e.errors()] # pyright: ignore 

81 )