marshmallow_dataclass/tests/test_field_for_schema.py at master · dan-starkware/marshmallow_dataclass · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import inspect
import typing
import unittest
from enum import Enum
from typing import Dict, Optional, Union, Any, List, Tuple

from marshmallow import fields, Schema

from marshmallow_dataclass import field_for_schema, dataclass, union_field


class TestFieldForSchema(unittest.TestCase):
    def assertFieldsEqual(self, a: fields.Field, b: fields.Field):
        self.assertEqual(a.__class__, b.__class__, "field class")

        def attrs(x):
            return {
                k: f"{v!r} ({v.__mro__!r})" if inspect.isclass(v) else repr(v)
                for k, v in x.__dict__.items()
                if not k.startswith("_")
            }

        self.assertEqual(attrs(a), attrs(b))

    def test_int(self):
        self.assertFieldsEqual(
            field_for_schema(int, default=9, metadata=dict(required=False)),
            fields.Integer(default=9, missing=9, required=False),
        )

    def test_any(self):
        self.assertFieldsEqual(
            field_for_schema(Any), fields.Raw(required=True, allow_none=True)
        )

    def test_dict_from_typing(self):
        self.assertFieldsEqual(
            field_for_schema(Dict[str, float]),
            fields.Dict(
                keys=fields.String(required=True),
                values=fields.Float(required=True),
                required=True,
            ),
        )

    def test_builtin_dict(self):
        self.assertFieldsEqual(
            field_for_schema(dict),
            fields.Dict(
                keys=fields.Raw(required=True, allow_none=True),
                values=fields.Raw(required=True, allow_none=True),
                required=True,
            ),
        )

    def test_builtin_list(self):
        self.assertFieldsEqual(
            field_for_schema(list, metadata=dict(required=False)),
            fields.List(fields.Raw(required=True, allow_none=True), required=False),
        )

    def test_explicit_field(self):
        explicit_field = fields.Url(required=True)
        self.assertFieldsEqual(
            field_for_schema(str, metadata={"marshmallow_field": explicit_field}),
            explicit_field,
        )

    def test_str(self):
        self.assertFieldsEqual(field_for_schema(str), fields.String(required=True))

    def test_optional_str(self):
        self.assertFieldsEqual(
            field_for_schema(Optional[str]),
            fields.String(allow_none=True, required=False, default=None, missing=None),
        )

    def test_enum(self):
        import marshmallow_enum

        class Color(Enum):
            RED: 1
            GREEN: 2
            BLUE: 3

        self.assertFieldsEqual(
            field_for_schema(Color),
            marshmallow_enum.EnumField(enum=Color, required=True),
        )

    def test_union(self):
        self.assertFieldsEqual(
            field_for_schema(Union[int, str]),
            union_field.Union(
                [
                    (int, fields.Integer(required=True)),
                    (str, fields.String(required=True)),
                ],
                required=True,
            ),
        )

    def test_newtype(self):
        self.assertFieldsEqual(
            field_for_schema(typing.NewType("UserId", int), default=0),
            fields.Integer(required=False, description="UserId", default=0, missing=0),
        )

    def test_marshmallow_dataclass(self):
        class NewSchema(Schema):
            pass

        @dataclass(base_schema=NewSchema)
        class NewDataclass:
            pass

        self.assertFieldsEqual(
            field_for_schema(NewDataclass, metadata=dict(required=False)),
            fields.Nested(NewDataclass.Schema),
        )

    def test_override_container_type_with_type_mapping(self):
        type_mapping = [
            (List, fields.List, List[int]),
            (Dict, fields.Dict, Dict[str, int]),
            (Tuple, fields.Tuple, Tuple[int, str, bytes]),
        ]
        for base_type, marshmallow_field, schema in type_mapping:

            class MyType(marshmallow_field):
                ...

            self.assertIsInstance(field_for_schema(schema), marshmallow_field)

            class BaseSchema(Schema):
                TYPE_MAPPING = {base_type: MyType}

            self.assertIsInstance(
                field_for_schema(schema, base_schema=BaseSchema), MyType
            )


if __name__ == "__main__":
    unittest.main()