Coverage for src/fastoai/models/_metadata.py: 88%

32 statements  

« prev     ^ index     » next       coverage.py v7.6.8, created at 2024-12-06 09:34 +0800

1from typing import Annotated, Any, Self 

2 

3from pydantic.json_schema import ( 

4 DEFAULT_REF_TEMPLATE, 

5 GenerateJsonSchema, 

6 JsonSchemaMode, 

7 WithJsonSchema, 

8) 

9from sqlalchemy.ext.mutable import MutableDict 

10from sqlmodel import JSON, Field, SQLModel 

11 

12 

13class WithMetadata(SQLModel): 

14 metadata_: Annotated[ 

15 dict[ 

16 Annotated[str, Field(max_length=64)], 

17 Annotated[str, Field(max_length=512)] | Any, 

18 ] 

19 | None, 

20 Field( 

21 alias="metadata", 

22 sa_type=MutableDict.as_mutable(JSON), # type: ignore 

23 sa_column_kwargs={"name": "metadata"}, 

24 ), 

25 WithJsonSchema({"type": "object"}), 

26 ] = None 

27 """Set of 16 key-value pairs that can be attached to an object. 

28 

29 This can be useful for storing additional information about the object in a 

30 structured format. Keys can be a maximum of 64 characters long and values can be 

31 a maximum of 512 characters long. 

32 """ 

33 

34 def __init__(self, **kwargs): 

35 if "metadata" in kwargs: 

36 kwargs["metadata_"] = kwargs.pop("metadata") 

37 return super().__init__(**kwargs) 

38 

39 @classmethod 

40 def model_validate( 

41 cls, 

42 obj: Any, 

43 *, 

44 strict: bool | None = None, 

45 from_attributes: bool | None = None, 

46 context: dict[str, Any] | None = None, 

47 update: dict[str, Any] | None = None, 

48 ) -> Self: 

49 if "metadata" in obj: 

50 obj["metadata_"] = obj.pop("metadata") 

51 return super().model_validate( 

52 obj, 

53 strict=strict, 

54 from_attributes=from_attributes, 

55 context=context, 

56 update=update, 

57 ) 

58 

59 def model_dump(self, by_alias: bool = False, **kwargs) -> dict[str, Any]: 

60 result = super().model_dump(**kwargs) 

61 if by_alias and "metadata_" in result: 

62 result["metadata"] = result.pop("metadata_") 

63 return result 

64 

65 def model_dump_json(self, by_alias: bool = False, **kwargs) -> str: 

66 result = super().model_dump_json(**kwargs) 

67 if by_alias and "metadata_" in result: 

68 result.replace("metadata_", "metadata") 

69 return result 

70 

71 @classmethod 

72 def model_json_schema( 

73 cls, 

74 by_alias: bool = False, 

75 ref_template: str = DEFAULT_REF_TEMPLATE, 

76 schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema, 

77 mode: JsonSchemaMode = "validation", 

78 ) -> dict[str, Any]: 

79 schema = super().model_json_schema( 

80 by_alias=by_alias, 

81 ref_template=ref_template, 

82 schema_generator=schema_generator, 

83 mode=mode, 

84 ) 

85 if by_alias and "metadata_" in schema["properties"]: 

86 schema["properties"]["metadata"] = schema["properties"].pop("metadata_") 

87 return schema