diff --git a/marshmallow_oneofschema/one_of_schema.py b/marshmallow_oneofschema/one_of_schema.py index 7957f46..983ef86 100644 --- a/marshmallow_oneofschema/one_of_schema.py +++ b/marshmallow_oneofschema/one_of_schema.py @@ -57,8 +57,46 @@ def get_obj_type(self, obj): type_field = "type" type_field_remove = True + type_field_dump = True type_schemas = {} + def __init__(self, **kwargs): + only = kwargs.get("only") + exclude = kwargs.get("exclude", ()) + if only is not None: + self.type_field_dump = self.type_field in only + kwargs["only"] = () + if exclude: + self.type_field_dump = self.type_field not in exclude + kwargs["exclude"] = () + super().__init__(**kwargs) + self._init_type_schemas(only, exclude) + + def _init_type_schemas(self, only, exclude): + self.type_schemas = { + k: self._create_type_schema_instance(v, only, exclude) + for k, v in self.type_schemas.items() + } + + def _create_type_schema_instance(self, SchemaCls, only, exclude): + if only or exclude: + if SchemaCls.opts.fields: + available_field_names = self.set_class(SchemaCls.opts.fields) + else: + available_field_names = self.set_class( + SchemaCls._declared_fields.keys() + ) + if SchemaCls.opts.additional: + available_field_names |= self.set_class( + SchemaCls.opts.additional + ) + if only: + only = self.set_class(only) & available_field_names + if exclude: + exclude = self.set_class(exclude) & available_field_names + + return SchemaCls(only=only, exclude=exclude) + def get_obj_type(self, obj): """Returns name of object schema""" return obj.__class__.__name__ @@ -96,16 +134,14 @@ def _dump(self, obj, *, update_fields=True, **kwargs): {"_schema": "Unknown object class: %s" % obj.__class__.__name__}, ) - type_schema = self.type_schemas.get(obj_type) - if not type_schema: + schema = self.type_schemas.get(obj_type) + if not schema: return None, {"_schema": "Unsupported object type: %s" % 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: + if result is not None and self.type_field_dump: result[self.type_field] = obj_type return result @@ -160,17 +196,17 @@ def _load(self, data, *, partial=None, unknown=None, **kwargs): ) try: - type_schema = self.type_schemas.get(data_type) + schema = self.type_schemas.get(data_type) except TypeError: # data_type could be unhashable - raise ValidationError({self.type_field: ["Invalid value: %s" % data_type]}) - if not type_schema: + raise ValidationError( + {self.type_field: ["Invalid value: %s" % data_type]} + ) + if not schema: raise ValidationError( {self.type_field: ["Unsupported value: %s" % 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) diff --git a/tests/test_one_of_schema.py b/tests/test_one_of_schema.py index adbebe8..f581f40 100644 --- a/tests/test_one_of_schema.py +++ b/tests/test_one_of_schema.py @@ -115,6 +115,22 @@ def test_dump_with_empty_keeps_type(self): result = MySchema().dump(Empty()) assert {"type": "Empty"} == result + def test_dump_only(self): + result = MySchema(only=("type", "value1")).dump( + [Foo("hello"), Bar(123), Baz(456, 789)], many=True + ) + assert [ + {"type": "Foo"}, + {"type": "Bar"}, + {"type": "Baz", "value1": 456}, + ] == result + + def test_dump_exclude(self): + result = MySchema(exclude=("type", "value2")).dump( + [Foo("hello"), Bar(123), Baz(456, 789)], many=True + ) + assert [{"value": "hello"}, {"value": 123}, {"value1": 456}] == result + def test_load(self): foo_result = MySchema().load({"type": "Foo", "value": "world"}) assert Foo("world") == foo_result