Pydantic models parser

pydantic_parser.py

"""
This is free and unencumbered software released into the public domain.
Anyone is free to copy, modify, publish, use, compile, sell, or
distribute this software, either in source code form or as a compiled
binary, for any purpose, commercial or non-commercial, and by any
means.
In jurisdictions that recognize copyright laws, the author or authors
of this software dedicate any and all copyright interest in the
software to the public domain. We make this dedication for the benefit
of the public at large and to the detriment of our heirs and
successors. We intend this dedication to be an overt act of
relinquishment in perpetuity of all present and future rights to this
software under copyright law.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.
For more information, please refer to <http://unlicense.org/>
"""
from __future__ import annotations
from collections.abc import Mapping as _Mapping, Iterable as _Iterable
from typing import (
    Any,
    Dict,
    Iterator,
    List,
    Mapping,
    Iterable,
    Optional,
    Tuple,
    Type,
    TypeVar,
    Union,
    get_args,
    get_origin,
    overload,
)

# I don't know why but importing typevar T from mypy works and defining it ourself does not work 😐
from typing import T  # type: ignore[attr-defined]

from pydantic import BaseModel, create_model
from pydantic.main import BaseConfig


ModelT = TypeVar("ModelT", bound=BaseModel)
DataT = TypeVar("DataT")
DefaultT = TypeVar("DefaultT")


@overload
def validate(
    _type: Type[DataT],
    data: Any,
) -> DataT:
    ...


@overload
def validate(
    _type: Type[DataT],
    data: None,
    default: DefaultT,
) -> DefaultT:
    ...


@overload
def validate(
    _type: Type[DataT],
    data: Any,
    default: DefaultT,
) -> DataT:
    ...


@overload
def validate(
    _type: T,
    data: Any,
) -> T:
    ...


@overload
def validate(_type: T, data: Any, default: DefaultT) -> T:
    ...


def validate(_type: Any, data: Any, default: Any = ...) -> Any:
    """Perform validation according to given type"""
    is_optional, data_type = _check_optional(_type)
    # Handle case when annotation is optional
    # Default value is always used when provided
    if data is None:
        # Handle case when default is available
        if default is not ...:
            return default
        # Handle wase when annotation is optional
        if is_optional:
            return None
    # In all other case we must try to parse the given data into data_type
    # Handle case when data_type is a pydantic model
    if _check_pydantic_model(data_type):
        return parse_model(data_type, data)
    # Handle all other cases by generating a dynamic model
    # At this point we can dynamically create a pydantic model
    DataModel: Type[BaseModel] = create_model(
        "DataModel", __root__=(_type, ...), __config__=_BaseConfig
    )
    # We can reuse the parse_model function and return __root__ attribute
    return parse_model(DataModel, data).__root__


def parse_model(
    model: Type[ModelT],
    data: Any = None,
) -> ModelT:
    """Get a pydantic model with given type from message.

    Arguments:
        model: a BaseModel class that should be used to parse the data
        data: data to parse, can be of any type

    Returns:
        An instance of the model class given as argument

    Raises:
        ValidationError: When given data cannot be parsed into desired pydantic model
    """
    # Handle bytes and strings
    if isinstance(data, (bytes, str)):
        return model.parse_raw(data)
    # Handle mapping and sequences
    elif isinstance(data, (_Mapping, _Iterable, Mapping, Iterable)):
        return model.parse_obj(data)
    # Handle objects
    return model.from_orm(data)


class _BaseConfig(BaseConfig):
    arbitrary_types_allowed = True


def _check_optional(annotation: Any) -> Tuple[bool, Any]:
    """Check if an annotation is optional"""
    if get_origin(annotation) is Union:
        args = get_args(annotation)
        if type(None) or None in args:
            return True, annotation
    return False, annotation


def _check_pydantic_model(annotation: Any) -> bool:
    """Check that given object is a valid pydantic model class."""
    try:
        return issubclass(annotation, BaseModel)
    except TypeError:
        return False


### UNIT TESTS ###
import pytest  # noqa: E402
from pydantic import ValidationError  # noqa: E402


def test_parse_pydantic_model_from_dict() -> None:
    """Check that we can parse a pydantic model from a dict."""

    class Foo(BaseModel):
        a: int
        b: float

    model = validate(Foo, {"a": 1, "b": 2})
    assert model.a == 1
    assert model.b == 2


def test_parse_pydantic_model_from_string() -> None:
    """Check that we can parse a pydantic model from a string."""

    class Foo(BaseModel):
        a: int
        b: float

    model = validate(Foo, '{"a": 1, "b": 2}')
    assert model.a == 1
    assert model.b == 2


def test_parse_pydantic_model_from_bytes() -> None:
    """Check that we can parse a pydantic model from bytes."""

    class Foo(BaseModel):
        a: int
        b: float

    model = validate(Foo, b'{"a": 1, "b": 2}')
    assert model.a == 1
    assert model.b == 2


def test_parse_pydantic_model_from_list() -> None:
    """Check that we can parse a pydantic model from a list.

    Pydantic support parsing lists. Values can be retrieved using the special __root__ attribute.
    """

    class Bar(BaseModel):
        __root__: List[int]

    model = validate(Bar, [1.1, 2.2, 3.3])
    assert model.__root__ == [1, 2, 3]


def test_parse_pydantic_model_from_empty_dict() -> None:
    """Check that default pydantic model is indeed parsed from provided with an emty dict"""

    class Baz(BaseModel):
        a: int = 0
        b: float = 0

    model = validate(Baz, {})
    assert model.a == 0
    assert model.b == 0


def test_parse_pydantic_model_from_none_default_none() -> None:
    """Check that None is always returned when data is None and default is None"""

    class Baz(BaseModel):
        a: int = 0
        b: float

    model = validate(Baz, None, None)
    # model here is typed as None because given data is None and default is None
    assert model is None


def test_parse_pydantic_validation_error() -> None:
    """parse_model should raise an error when no data is given and optional is False"""

    class Baz(BaseModel):
        a: int = 1
        b: float

    with pytest.raises(ValidationError):
        # Note that mypy detects the error here (remove the "type: ignore[call-overload]" comment below to see error)
        # No overload variant of "parse_model" matches argument types "Type[Baz]", "bool"
        # It also displays a possible usage:
        # def parse_model(model: Type[BaseModel], *, optional: Literal[True]) -> None
        # It means that only optional=True is supported when no data is provided
        model = validate(Baz, {})
        model


def test_parse_pydantic_validation_error_even_with_default() -> None:
    """Check that if a dict is provided, even if a default value is specified, validation is performed and errors are raised if needed"""

    class Baz(BaseModel):
        a: int
        b: float

    # model here is typed a Baz because data is not None so validation will happen no matter default value
    with pytest.raises(ValidationError):
        model = validate(Baz, {}, default=None)
        # Assertion will never be run, it's used to perform type checks
        # as_integer_ratio method exists on integers
        assert model.a.as_integer_ratio()
        # is_integer method exists on floatss
        assert model.b.is_integer()


def test_parse_pydantic_model_from_object() -> None:
    """Pydantic support parsing objects when orm_mode is set to True in config."""

    class Bar(BaseModel):
        a: int
        b: float

        class Config:
            orm_mode = True

    class MyObject:
        def __init__(self, a: int, b: float) -> None:
            self.a = a
            self.b = b

    model = validate(Bar, MyObject(1, 3))
    assert model.a == 1
    assert model.b == 3


def test_parse_pydantic_model_from_custom_mapping_typing() -> None:
    """Pydantic support parsing dict.

    Starting python 3.7 it's possible to declare classes as Mapping using typing.Mapping generic class
    """

    class Foo(BaseModel):
        a: int

    class MyMapping(Mapping[str, float]):
        def __len__(self) -> int:
            return 1

        def __iter__(self) -> Iterator[str]:
            yield "a"

        def __getitem__(self, key: str) -> float:
            return 2.2

    model = validate(Foo, MyMapping())
    assert model.a == 2


def test_parse_pydantic_model_from_custom_mapping_abc() -> None:
    """Pydantic support parsing dict, so a custom mapping that inherits from collections.abc.Mapping shoud work.

    Annotation can be used with collections.abc.Mapping stating python 3.9 only.

    This includes DefaultDict, OrderedDict, Counter classes and many more.
    """

    class Foo(BaseModel):
        a: int

    class MyMapping(_Mapping):  # type: ignore[type-arg]
        def __len__(self) -> int:
            return 1

        def __iter__(self) -> Iterator[str]:
            yield "a"

        def __getitem__(self, key: str) -> float:
            return 2.2

    model = validate(Foo, MyMapping())
    assert model.a == 2


def test_parse_list_of_int() -> None:
    """Check that we can parse a list of integers"""
    result = validate(List[int], [2.1, 3.4, 4.5])
    assert result == [2, 3, 4]


def test_parse_list_of_int_from_set() -> None:
    """Check that we can parse a list of integers from a set"""
    result = validate(List[int], set([2.1, 3.4, 4.5]))
    assert result == [2, 3, 4]


def test_parse_list_of_int_from_string() -> None:
    """Check that we can parse a list of integers from a set"""
    result = validate(List[int], "[2.1, 3.4, 4.5]")
    assert result == [2, 3, 4]


def test_parse_list_of_int_from_bytes() -> None:
    """Check that we can parse a list of integers from a set"""
    result = validate(List[int], b"[2.1, 3.4, 4.5]")
    assert result == [2, 3, 4]


def test_parse_dict() -> None:
    """Check that we can parse a dict as is"""
    result = validate(Dict[str, str], {"hello": "world"})
    assert result["hello"] == "world"


def test_parse_dict_from_string() -> None:
    """Check that we can parse a dict from a string"""
    result = validate(Dict[str, str], '{"hello": "world"}')
    assert result["hello"] == "world"


def test_parse_dict_from_bytes() -> None:
    """Check that we can parse a dict from some bytes"""
    result = validate(Dict[str, str], b'{"hello": "world"}')
    assert result["hello"] == "world"


def test_parse_union() -> None:
    """Check that we can parse a dict from some bytes"""
    result = validate(Union[List[int], Dict[str, str]], [1, 2, 3])
    assert result[0] == 1
    result = validate(Union[List[int], Dict[str, str]], {"hello": "world"})
    assert result["hello"] == "world"
    result = validate(Union[List[int], Dict[str, str], None], None)
    assert result is None


def test_parse_optional_list_int() -> None:
    """Check that we can parse a dict from some bytes"""
    result = validate(Optional[List[int]], [1, 2, 3])
    assert result[0] == 1


def test_parse_optional_base_model() -> None:
    """Check that we can parse optionnal base model"""

    class Foo(BaseModel):
        hello: str

    result = validate(Optional[Foo], {"hello": "world"})
    assert result.hello == "world"

    with pytest.raises(ValidationError):
        result = validate(Optional[Foo], {})


def test_parse_optional_base_model_none() -> None:
    """Check that we do not raise validation error when optionnal"""

    class Foo(BaseModel):
        hello: str

    result = validate(Optional[Foo], None, None)
    assert result is None