Coverage for mongodb_odm / types.py: 93%

1from collections.abc import Mapping, Sequence 6 ctx1fabcde

2from typing import Any, TypeVar, Union 6 ctx1fabcde

3 

4from bson import ObjectId 6 ctx1fabcde

5from pydantic import GetCoreSchemaHandler 6 ctx1fabcde

6from pydantic_core import core_schema 6 ctx1fabcde

7from pymongo.operations import ( 6 ctx1fabcde

8 DeleteMany, 

9 DeleteOne, 

10 InsertOne, 

11 ReplaceOne, 

12 UpdateMany, 

13 UpdateOne, 

14) 

15 

16# Common types 

17DICT_TYPE = dict[str, Any] 6 ctx1fabcde

18SORT_TYPE = Union[str, Sequence[tuple[str, Union[int, str, Mapping[str, Any]]]]] 6 ctx1fabcde

19 

20DocumentType = TypeVar("DocumentType", bound=Mapping[str, Any]) 6 ctx1fabcde

21 

22WriteOp = Union[ 6 ctx1fabcde

23 InsertOne[DocumentType], 

24 DeleteOne, 

25 DeleteMany, 

26 ReplaceOne[DocumentType], 

27 UpdateOne, 

28 UpdateMany, 

29] 

30 

31 

32class ODMObjectId(ObjectId): 6 ctx1fabcde

33 @classmethod 6 ctx1fabcde

34 def __get_pydantic_core_schema__( 6 ctx1fabcde

35 cls, _source_type: Any, _handler: Any 

36 ) -> core_schema.CoreSchema: 

37 return core_schema.json_or_python_schema( 6 ctx1fabcde

38 json_schema=core_schema.chain_schema( 

39 [ 

40 core_schema.str_schema(), 

41 core_schema.no_info_plain_validator_function(cls.validate), 

42 ] 

43 ), 

44 python_schema=core_schema.chain_schema( 

45 [ 

46 core_schema.union_schema( 

47 [ 

48 core_schema.is_instance_schema(ObjectId), 

49 core_schema.str_schema(), 

50 ] 

51 ), 

52 core_schema.no_info_plain_validator_function(cls.validate), 

53 ] 

54 ), 

55 serialization=core_schema.plain_serializer_function_ser_schema( 

56 lambda x: ObjectId(x) 

57 ), 

58 ) 

59 

60 @classmethod 6 ctx1fabcde

61 def validate(cls, v: Union[str, ObjectId]) -> ObjectId: 6 ctx1fabcde

62 if isinstance(v, ObjectId): 6 ctx1fabcde

63 """No conversion needed if type are already ObjectId""" 5 ctx1abcde

64 return v 6 ctx1fabcde

65 try: 6 ctx1fabcde

66 return ObjectId(v) 6 ctx1fabcde

67 except Exception as e: 6 ctx1fabcde

68 raise ValueError("Invalid data. ObjectId required") from e 6 ctx1fabcde

69 

70 

71class ObjectIdStr(str): 6 ctx1fabcde

72 _example_object_id = str(ObjectId()) 6 ctx1fabcde

73 

74 @classmethod 6 ctx1fabcde

75 def __get_pydantic_core_schema__( 6 ctx1fabcde

76 cls, _source_type: Any, handler: GetCoreSchemaHandler 

77 ) -> core_schema.CoreSchema: 

78 return core_schema.json_or_python_schema( 6 ctx1fabcde

79 json_schema=core_schema.chain_schema( 

80 [ 

81 core_schema.str_schema(), 

82 core_schema.no_info_plain_validator_function(cls.validate), 

83 ] 

84 ), 

85 python_schema=core_schema.chain_schema( 

86 [ 

87 core_schema.union_schema( 

88 [ 

89 core_schema.is_instance_schema(ObjectId), 

90 core_schema.str_schema(), 

91 ] 

92 ), 

93 core_schema.no_info_plain_validator_function(cls.validate), 

94 ] 

95 ), 

96 serialization=core_schema.plain_serializer_function_ser_schema( 

97 lambda x: str(x), # Always serialize as string 

98 return_schema=core_schema.str_schema(), 

99 ), 

100 ) 

101 

102 @classmethod 6 ctx1fabcde

103 def validate(cls, v: Union[str, ObjectId]) -> str: 6 ctx1fabcde

104 if isinstance(v, ObjectId): 6 ctx1fabcde

105 return cls(str(v)) 6 ctx1fabcde

106 elif isinstance(v, str): 6 ctx1fabcde

107 try: 6 ctx1fabcde

108 return cls(str(ObjectId(v))) 6 ctx1fabcde

109 except Exception as e: 6 ctx1fabcde

110 raise ValueError("Invalid ObjectId") from e 6 ctx1fabcde

111 

112 raise ValueError("Invalid data. ObjectId required") 6 ctx1fabcde

113 

114 @classmethod 6 ctx1fabcde

115 def __get_pydantic_json_schema__( 6 ctx1fabcde

116 cls, schema: core_schema.CoreSchema, handler: GetCoreSchemaHandler 

117 ) -> Any: 

118 """Provide OpenAPI/JSON Schema metadata.""" 

119 json_schema: Any = handler(schema) 

120 json_schema.update( 

121 type="string", 

122 format="objectid", 

123 example=cls._example_object_id, 

124 ) 

125 return json_schema