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

1 from typing import Annotated, Any, Self

2

3 from pydantic.json_schema import (

4 DEFAULT_REF_TEMPLATE,

5 GenerateJsonSchema,

6 JsonSchemaMode,

7 WithJsonSchema,

8)

9 from sqlalchemy.ext.mutable import MutableDict

10 from sqlmodel import JSON, Field, SQLModel

11

12

13 class 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