Coverage for polar/file/schemas.py: 83%

1from datetime import datetime 

2from typing import Annotated, Any, Literal, Self 

3 

4from pydantic import UUID4, Discriminator, Field, TypeAdapter, computed_field 

5 

6from polar.integrations.aws.s3.schemas import ( 

7 S3DownloadURL, 

8 S3File, 

9 S3FileCreate, 

10 S3FileDownload, 

11 S3FileUpload, 

12 S3FileUploadCompleted, 

13) 

14from polar.kit.schemas import ClassName, MergeJSONSchema, Schema, SetSchemaReference 

15from polar.models.file import File, FileServiceTypes 

16 

17from .s3 import S3_SERVICES 

18 

19 

20class FileCreateBase(S3FileCreate): 

21 service: FileServiceTypes 

22 version: str | None = None 

23 

24 

25class DownloadableFileCreate(FileCreateBase): 

26 """Schema to create a file to be associated with the downloadables benefit.""" 

27 

28 service: Literal[FileServiceTypes.downloadable] 

29 

30 

31class ProductMediaFileCreate(FileCreateBase): 

32 """Schema to create a file to be used as a product media file.""" 

33 

34 service: Literal[FileServiceTypes.product_media] 

35 mime_type: str = Field( 

36 description=( 

37 "MIME type of the file. Only images are supported for this type of file." 

38 ), 

39 pattern=r"^image\/(jpeg|png|gif|webp|svg\+xml)$", 

40 ) 

41 size: int = Field( 

42 description=( 

43 "Size of the file. A maximum of 10 MB is allowed for this type of file." 

44 ), 

45 le=10 * 1024 * 1024, 

46 ) 

47 

48 

49class OrganizationAvatarFileCreate(FileCreateBase): 

50 """Schema to create a file to be used as an organization avatar.""" 

51 

52 service: Literal[FileServiceTypes.organization_avatar] 

53 mime_type: str = Field( 

54 description=( 

55 "MIME type of the file. Only images are supported for this type of file." 

56 ), 

57 pattern=r"^image\/(jpeg|png|gif|webp|svg\+xml)$", 

58 ) 

59 size: int = Field( 

60 description=( 

61 "Size of the file. A maximum of 1 MB is allowed for this type of file." 

62 ), 

63 le=1 * 1024 * 1024, 

64 ) 

65 

66 

67FileCreate = Annotated[ 

68 DownloadableFileCreate | ProductMediaFileCreate | OrganizationAvatarFileCreate, 

69 Discriminator("service"), 

70 SetSchemaReference("FileCreate"), 

71] 

72 

73 

74class FileReadBase(S3File): 

75 version: str | None 

76 service: FileServiceTypes 

77 is_uploaded: bool 

78 created_at: datetime 

79 

80 

81class DownloadableFileRead(FileReadBase): 

82 """File to be associated with the downloadables benefit.""" 

83 

84 service: Literal[FileServiceTypes.downloadable] 

85 

86 

87class PublicFileReadBase(FileReadBase): 

88 @computed_field # type: ignore[prop-decorator] 

89 @property 

90 def public_url(self) -> str: 

91 return S3_SERVICES[self.service].get_public_url(self.path) 

92 

93 

94class ProductMediaFileRead(PublicFileReadBase): 

95 """File to be used as a product media file.""" 

96 

97 service: Literal[FileServiceTypes.product_media] 

98 

99 

100class OrganizationAvatarFileRead(PublicFileReadBase): 

101 """File to be used as an organization avatar.""" 

102 

103 service: Literal[FileServiceTypes.organization_avatar] 

104 

105 

106FileRead = Annotated[ 

107 DownloadableFileRead | ProductMediaFileRead | OrganizationAvatarFileRead, 

108 Discriminator("service"), 

109 MergeJSONSchema({"title": "FileRead"}), 

110 ClassName("FileRead"), 

111] 

112 

113FileReadAdapter: TypeAdapter[FileRead] = TypeAdapter[FileRead](FileRead) 

114 

115 

116class FileUpload(S3FileUpload): 

117 version: str | None 

118 is_uploaded: bool = False 

119 service: FileServiceTypes 

120 

121 

122class FileUploadCompleted(S3FileUploadCompleted): ... 122 ↛ 125line 122 didn't jump to line 125 because

123 

124 

125class FileDownload(S3FileDownload): 

126 version: str | None 

127 is_uploaded: bool 

128 service: FileServiceTypes 

129 

130 @classmethod 

131 def from_presigned(cls, file: File, url: str, expires_at: datetime) -> Self: 

132 file_dict: dict[str, Any] = dict( 

133 id=file.id, 

134 organization_id=file.organization_id, 

135 name=file.name, 

136 path=file.path, 

137 mime_type=file.mime_type, 

138 size=file.size, 

139 version=file.version, 

140 service=file.service, 

141 checksum_etag=file.checksum_etag, 

142 last_modified_at=file.last_modified_at, 

143 storage_version=file.storage_version, 

144 is_uploaded=file.is_uploaded, 

145 created_at=file.created_at, 

146 ) 

147 if file.checksum_sha256_base64 and file.checksum_sha256_hex: 

148 file_dict.update( 

149 checksum_sha256_base64=file.checksum_sha256_base64, 

150 checksum_sha256_hex=file.checksum_sha256_hex, 

151 ) 

152 

153 return cls( 

154 **file_dict, 

155 download=S3DownloadURL( 

156 url=url, 

157 expires_at=expires_at, 

158 ), 

159 ) 

160 

161 

162class FileUpdate(Schema): 

163 id: UUID4 

164 version: str | None 

165 checksum_etag: str 

166 last_modified_at: datetime 

167 storage_version: str | None 

168 checksum_sha256_base64: str | None 

169 checksum_sha256_hex: str | None 

170 

171 

172class FilePatch(Schema): 

173 name: str | None = None 

174 version: str | None = None