Coverage for src/prisma/_fields.py: 91%

1from __future__ import annotations 

2 

3import base64 

4from typing import ( 

5 TYPE_CHECKING, 

6 Any, 

7 Dict, 

8 List, 

9 Union, 

10 Iterator, 

11 overload, 

12) 

13from typing_extensions import override 

14 

15from pydantic import Json as _PydanticJson 

16 

17from ._compat import PYDANTIC_V2, CoreSchema, core_schema 

18 

19if TYPE_CHECKING: 

20 from .types import Serializable # noqa: TID251 

21 

22__all__ = ( 

23 'Json', 

24 'Base64', 

25) 

26 

27 

28if TYPE_CHECKING: 

29 BaseJson = str 

30else: 

31 BaseJson = _PydanticJson 

32 

33_JsonKeys = Union[ 

34 None, 

35 bool, 

36 float, 

37 int, 

38 str, 

39] 

40 

41# Base64 data should only be valid ascii, we limit our encoding to ascii so that 

42# any erroneous data is caught as early on as possible. 

43BASE64_ENCODING = 'ascii' 

44 

45 

46# inherit from _PydanticJson so that pydantic will automatically 

47# transform the json string into python objects. 

48class Json(BaseJson): 

49 data: Serializable 

50 

51 def __init__(self, data: Serializable) -> None: 

52 self.data = data 

53 super().__init__() 

54 

55 @classmethod 

56 def keys(cls, **data: Serializable) -> Json: 

57 return cls(data) 

58 

59 if TYPE_CHECKING: 

60 # Fields that are of the `Json` type are automatically 

61 # de-serialized from json to the corresponding python type 

62 # when the model is created, e.g. 

63 # 

64 # User(json_obj='{"foo": null}') -> User(json_obj={'foo': None}) 

65 # 

66 # As we don't know what the type will actually be at runtime 

67 # we add methods here for convenience so that naive access 

68 # to the field is still allowed, e.g. 

69 # 

70 # user.json_obj['foo'] 

71 # user.json_obj[1] 

72 # user.json_obj[1:5] 

73 # 

74 # It should be noted that users will still have 

75 # to validate / cast fields to the type they are expecting 

76 # for any strict type binding or nested index calls to work, e.g. 

77 # 

78 # isinstance(user.json_obj, dict) 

79 # cast(Dict[str, Any], user.json_obj) 

80 # prisma.validate(ExpectedType, user.json_obj) # NOTE: not implemented yet 

81 @overload # type: ignore 

82 def __getitem__(self, i: slice) -> List[Serializable]: ... 

83 

84 @overload 

85 def __getitem__(self, i: _JsonKeys) -> Serializable: ... 

86 

87 @override 

88 def __getitem__(self, i: Union[_JsonKeys, slice]) -> Serializable: # pyright: ignore[reportIncompatibleMethodOverride] 

89 ... 

90 

91 

92class Base64: 

93 __slots__ = ('_raw',) 

94 

95 def __init__(self, raw: bytes) -> None: 

96 self._raw = raw 

97 

98 @classmethod 

99 def encode(cls, value: bytes) -> Base64: 

100 """Encode bytes into valid Base64""" 

101 return cls(base64.b64encode(value)) 

102 

103 @classmethod 

104 def fromb64(cls, value: Union[str, bytes]) -> Base64: 

105 """ 

106 Create an instance of the `Base64` class from data that has already 

107 been encoded into a valid base 64 structure. 

108 """ 

109 if isinstance(value, bytes): 

110 return cls(value) 

111 

112 return cls(value.encode(BASE64_ENCODING)) 

113 

114 def decode(self) -> bytes: 

115 """Decode from Base64 to the original bytes object""" 

116 return base64.b64decode(self._raw) 

117 

118 # NOTE: we explicitly use a different encoding here as we are decoding 

119 # to the original data provided by the user, this data does not have 

120 # the limitation of being ascii only that the raw Base64 data does 

121 def decode_str(self, encoding: str = 'utf-8') -> str: 

122 """Decode from Base64 to the original string 

123 

124 This decodes using the `utf-8` encoding by default, 

125 you can customise the encoding like so: 

126 

127 ```py 

128 value = b64.decode_str('ascii') 

129 ``` 

130 """ 

131 return self.decode().decode(encoding) 

132 

133 # Support OpenAPI schema generation 

134 if PYDANTIC_V2: 

135 

136 @classmethod 

137 def __get_pydantic_json_schema__(cls, core_schema: CoreSchema, handler: Any) -> Any: 

138 json_schema = handler(core_schema) 

139 json_schema = handler.resolve_ref_schema(json_schema) 

140 json_schema['type'] = 'string' 

141 json_schema['format'] = 'byte' 

142 return json_schema 

143 else: 

144 

145 @classmethod 

146 def __modify_schema__(cls, field_schema: Dict[str, object]) -> None: 

147 field_schema['type'] = 'string' 

148 field_schema['format'] = 'byte' 

149 

150 # Support converting objects into Base64 at the Pydantic level 

151 if PYDANTIC_V2: 

152 

153 @classmethod 

154 def __get_pydantic_core_schema__( 

155 cls, 

156 source_type: Any, 

157 *args: Any, 

158 **kwargs: Any, 

159 ) -> CoreSchema: 

160 return core_schema.no_info_before_validator_function( 

161 cls._validate, 

162 schema=core_schema.any_schema(), 

163 serialization=core_schema.plain_serializer_function_ser_schema(lambda instance: str(instance)), 

164 ) 

165 else: 

166 

167 @classmethod 

168 def __get_validators__(cls) -> Iterator[object]: 

169 yield cls._validate 

170 

171 @classmethod 

172 def _validate(cls, value: object) -> Base64: 

173 if isinstance(value, Base64): 

174 return value 

175 

176 # TODO: validate that the structure of the input is valid too? 

177 if isinstance(value, str): 

178 return cls(bytes(value, BASE64_ENCODING)) 

179 

180 if isinstance(value, bytes): 180 ↛ 183line 180 didn't jump to line 183, because the condition on line 180 was never false

181 return cls(value) 

182 

183 raise ValueError(f'Could not convert type: {type(value)} to a Base64 object; Expected a string or bytes object') 

184 

185 @override 

186 def __str__(self) -> str: 

187 return self._raw.decode(BASE64_ENCODING) 

188 

189 @override 

190 def __repr__(self) -> str: 

191 return f'{self.__class__.__name__}({self._raw})' # type: ignore[str-bytes-safe] 

192 

193 @override 

194 def __eq__(self, other: Any) -> bool: 

195 if isinstance(other, Base64): 

196 return self._raw == other._raw 

197 

198 return False