Skip to content

polymorphic_schema

PolymorphicSchema

Bases: BaseSchema

PolymorphicSchema is based on https://github.com/marshmallow-code/marshmallow-oneofschema.

Source code in griptape/schemas/polymorphic_schema.py 
 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
class PolymorphicSchema(BaseSchema):
    """PolymorphicSchema is based on https://github.com/marshmallow-code/marshmallow-oneofschema."""

    def __init__(self, inner_class: Any, types_overrides: Optional[dict[str, type]] = None, **kwargs) -> None:
        super().__init__(**kwargs)

        self.inner_class = inner_class
        self.types_overrides = types_overrides

    type_field = "type"
    type_field_remove = True

    def get_obj_type(self, obj: Any) -> Any:
        """Returns name of the schema during dump() calls, given the object being dumped."""
        return obj.__class__.__name__

    def get_data_type(self, data: Any) -> Any:
        """Returns name of the schema during load() calls, given the data being loaded. Defaults to looking up `type_field` in the data."""
        data_type = data.get(self.type_field)
        if self.type_field in data and self.type_field_remove:
            data.pop(self.type_field)
        return data_type

    def dump(self, obj: Any, *, many: Any = None, **kwargs) -> Any:
        errors = {}
        result_data = []
        result_errors = {}
        many = self.many if many is None else bool(many)
        if not many:
            result = result_data = self._dump(obj, **kwargs)
        else:
            for idx, o in enumerate(obj):
                try:
                    result = self._dump(o, **kwargs)
                    result_data.append(result)
                except ValidationError as error:
                    result_errors[idx] = error.normalized_messages()
                    result_data.append(error.valid_data)

        result = result_data
        errors = result_errors

        if not errors:
            return result
        exc = ValidationError(errors, data=obj, valid_data=result)  # pyright: ignore[reportArgumentType]
        raise exc

    def _dump(self, obj: Any, *, update_fields: bool = True, **kwargs) -> Any:
        obj_type = self.get_obj_type(obj)

        if not obj_type:
            return (None, {"_schema": f"Unknown object class: {obj.__class__.__name__}"})

        type_schema = BaseSchema.from_attrs_cls(obj.__class__, types_overrides=self.types_overrides)

        if not type_schema:
            return None, {"_schema": f"Unsupported object type: {obj_type}"}

        schema = type_schema if isinstance(type_schema, Schema) else type_schema()

        schema.context.update(getattr(self, "context", {}))

        result = schema.dump(obj, many=False, **kwargs)

        if result is not None:
            result[self.type_field] = obj_type  # pyright: ignore[reportArgumentType,reportCallIssue]

        return result

    def load(self, data: Any, *, many: Any = None, partial: Any = None, unknown: Any = None, **kwargs) -> Any:
        errors = {}
        result_data = []
        result_errors = {}
        many = self.many if many is None else bool(many)
        if partial is None:
            partial = self.partial
        if not many:
            try:
                result = result_data = self._load(data, partial=partial, unknown=unknown, **kwargs)
            except ValidationError as error:
                result_errors = error.normalized_messages()
                result_data.append(error.valid_data)
        else:
            for idx, item in enumerate(data):
                try:
                    result = self._load(item, partial=partial, **kwargs)
                    result_data.append(result)
                except ValidationError as error:
                    result_errors[idx] = error.normalized_messages()
                    result_data.append(error.valid_data)

        result = result_data
        errors = result_errors

        if not errors:
            return result
        exc = ValidationError(errors, data=data, valid_data=result)
        raise exc

    def _load(self, data: Any, *, partial: Any = None, unknown: Any = None, **kwargs) -> Any:
        if not isinstance(data, dict):
            raise ValidationError({"_schema": f"Invalid data type: {data}"})

        data = dict(data)
        unknown = unknown or self.unknown
        data_type = self.get_data_type(data)

        if data_type is None:
            raise ValidationError({self.type_field: ["Missing data for required field."]})

        type_schema = self.inner_class.get_schema(data_type, module_name=data.get("module_name"))
        if not type_schema:
            raise ValidationError({self.type_field: [f"Unsupported value: {data_type}"]})

        schema = type_schema if isinstance(type_schema, Schema) else type_schema()

        schema.context.update(getattr(self, "context", {}))

        return schema.load(data, many=False, partial=partial, unknown=unknown, **kwargs)

    def validate(self, data: Any, *, many: Any = None, partial: Any = None) -> Any:  # pyright: ignore[reportIncompatibleMethodOverride]
        try:
            self.load(data, many=many, partial=partial)
        except ValidationError as ve:
            return ve.messages
        return {}

inner_class = inner_class instance-attribute

type_field = 'type' class-attribute instance-attribute

type_field_remove = True class-attribute instance-attribute

types_overrides = types_overrides instance-attribute

__init__(inner_class, types_overrides=None, **kwargs)

Source code in griptape/schemas/polymorphic_schema.py 
13
14
15
16
17
def __init__(self, inner_class: Any, types_overrides: Optional[dict[str, type]] = None, **kwargs) -> None:
    super().__init__(**kwargs)

    self.inner_class = inner_class
    self.types_overrides = types_overrides

_dump(obj, *, update_fields=True, **kwargs)

Source code in griptape/schemas/polymorphic_schema.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def _dump(self, obj: Any, *, update_fields: bool = True, **kwargs) -> Any:
    obj_type = self.get_obj_type(obj)

    if not obj_type:
        return (None, {"_schema": f"Unknown object class: {obj.__class__.__name__}"})

    type_schema = BaseSchema.from_attrs_cls(obj.__class__, types_overrides=self.types_overrides)

    if not type_schema:
        return None, {"_schema": f"Unsupported object type: {obj_type}"}

    schema = type_schema if isinstance(type_schema, Schema) else type_schema()

    schema.context.update(getattr(self, "context", {}))

    result = schema.dump(obj, many=False, **kwargs)

    if result is not None:
        result[self.type_field] = obj_type  # pyright: ignore[reportArgumentType,reportCallIssue]

    return result

_load(data, *, partial=None, unknown=None, **kwargs)

Source code in griptape/schemas/polymorphic_schema.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def _load(self, data: Any, *, partial: Any = None, unknown: Any = None, **kwargs) -> Any:
    if not isinstance(data, dict):
        raise ValidationError({"_schema": f"Invalid data type: {data}"})

    data = dict(data)
    unknown = unknown or self.unknown
    data_type = self.get_data_type(data)

    if data_type is None:
        raise ValidationError({self.type_field: ["Missing data for required field."]})

    type_schema = self.inner_class.get_schema(data_type, module_name=data.get("module_name"))
    if not type_schema:
        raise ValidationError({self.type_field: [f"Unsupported value: {data_type}"]})

    schema = type_schema if isinstance(type_schema, Schema) else type_schema()

    schema.context.update(getattr(self, "context", {}))

    return schema.load(data, many=False, partial=partial, unknown=unknown, **kwargs)

dump(obj, *, many=None, **kwargs)

Source code in griptape/schemas/polymorphic_schema.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def dump(self, obj: Any, *, many: Any = None, **kwargs) -> Any:
    errors = {}
    result_data = []
    result_errors = {}
    many = self.many if many is None else bool(many)
    if not many:
        result = result_data = self._dump(obj, **kwargs)
    else:
        for idx, o in enumerate(obj):
            try:
                result = self._dump(o, **kwargs)
                result_data.append(result)
            except ValidationError as error:
                result_errors[idx] = error.normalized_messages()
                result_data.append(error.valid_data)

    result = result_data
    errors = result_errors

    if not errors:
        return result
    exc = ValidationError(errors, data=obj, valid_data=result)  # pyright: ignore[reportArgumentType]
    raise exc

get_data_type(data)

Returns name of the schema during load() calls, given the data being loaded. Defaults to looking up type_field in the data.

Source code in griptape/schemas/polymorphic_schema.py 
26
27
28
29
30
31
def get_data_type(self, data: Any) -> Any:
    """Returns name of the schema during load() calls, given the data being loaded. Defaults to looking up `type_field` in the data."""
    data_type = data.get(self.type_field)
    if self.type_field in data and self.type_field_remove:
        data.pop(self.type_field)
    return data_type

get_obj_type(obj)

Returns name of the schema during dump() calls, given the object being dumped.

Source code in griptape/schemas/polymorphic_schema.py 
22
23
24
def get_obj_type(self, obj: Any) -> Any:
    """Returns name of the schema during dump() calls, given the object being dumped."""
    return obj.__class__.__name__

load(data, *, many=None, partial=None, unknown=None, **kwargs)

Source code in griptape/schemas/polymorphic_schema.py 
 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
def load(self, data: Any, *, many: Any = None, partial: Any = None, unknown: Any = None, **kwargs) -> Any:
    errors = {}
    result_data = []
    result_errors = {}
    many = self.many if many is None else bool(many)
    if partial is None:
        partial = self.partial
    if not many:
        try:
            result = result_data = self._load(data, partial=partial, unknown=unknown, **kwargs)
        except ValidationError as error:
            result_errors = error.normalized_messages()
            result_data.append(error.valid_data)
    else:
        for idx, item in enumerate(data):
            try:
                result = self._load(item, partial=partial, **kwargs)
                result_data.append(result)
            except ValidationError as error:
                result_errors[idx] = error.normalized_messages()
                result_data.append(error.valid_data)

    result = result_data
    errors = result_errors

    if not errors:
        return result
    exc = ValidationError(errors, data=data, valid_data=result)
    raise exc

validate(data, *, many=None, partial=None)

Source code in griptape/schemas/polymorphic_schema.py 
130
131
132
133
134
135
def validate(self, data: Any, *, many: Any = None, partial: Any = None) -> Any:  # pyright: ignore[reportIncompatibleMethodOverride]
    try:
        self.load(data, many=many, partial=partial)
    except ValidationError as ve:
        return ve.messages
    return {}