Coverage for schemas/overrides.py: 93%

1from typing import TypeVar, Annotated, Union 

2from enum import Enum as _Enum 

3from pydantic import BaseModel as _BaseModel 

4from pydantic import ConfigDict, TypeAdapter, Field 

5from pydantic.types import AnyType 

6from decouple import config, Choices 

7 

8 

9def attribute_to_camel_case(snake_str: str) -> str: 

10 components = snake_str.split("_") 

11 return components[0] + ''.join(x.title() for x in components[1:]) 

12 

13 

14def schema_extra(schema: dict, _): 

15 props = {} 

16 for k, v in schema.get('properties', {}).items(): 

17 if not v.get("doc_hidden", False): 

18 props[k] = v 

19 schema["properties"] = props 

20 

21 

22class BaseModel(_BaseModel): 

23 model_config = ConfigDict(alias_generator=attribute_to_camel_case, 

24 use_enum_values=True, 

25 json_schema_extra=schema_extra, 

26 extra=config("EXTRA_PAYLOAD_ATTRIBUTES", default="ignore", 

27 cast=Choices(["ignore", "forbid", "allow"]))) 

28 

29 

30class Enum(_Enum): 

31 @classmethod 

32 def has_value(cls, value) -> bool: 

33 return value in cls._value2member_map_ 

34 

35 

36T = TypeVar('T') 

37 

38 

39class ORUnion: 

40 def __new__(cls, union_types: Union[AnyType], discriminator: str) -> T: 

41 return lambda **args: TypeAdapter(Annotated[union_types, Field(discriminator=discriminator)]) \ 

42 .validate_python(args)