diff --git a/README.rst b/README.rst index 8989368..bb1dfdd 100644 --- a/README.rst +++ b/README.rst @@ -119,6 +119,49 @@ You can use resulting schema everywhere marshmallow.Schema can be used, e.g. class MyOtherSchema(m.Schema): items = f.List(f.Nested(MyUberSchema)) +Building OneOfSchema with a decorator +------------------------------------- + +When creating "one of" schemas with many options it can become cumbersome to maintain `type_schemas` dict manually. +In this case you can leave it empty at first and later register member schemas to your OneOfSchema with the decorator: + +.. code:: python + + class MyUberSchema(OneOfSchema): + type_schemas = {} + + + @MyUberSchema.register_one_of + class OneOfManyManySchemas(marshmallow.Schema): + pass + +By default schemas are named with class name with removed "Schema" suffix, which corresponds to the most common case of +"Foo" - "FooSchema" naming convention. If this is not the case, you can customize keys in `type_schemas` by overriding `get_schema_name` method: + +.. code:: python + + class Data: + def __init__(self, value): + self.value = value + + + class MyUberSchema(OneOfSchema): + type_schemas = {} + + @staticmethod + def get_schema_name(schema_class): + return schema_class._ModelClass.__name__ + + + @MyUberSchema.register_one_of + class SchemaForDataClass(marshmallow.Schema): + _ModelClass = Data + value = fields.String() + + @marshmallow.post_load + def make_data(self, data, **kwargs): + return self._ModelClass(**data) + License ------- diff --git a/marshmallow_oneofschema/one_of_schema.py b/marshmallow_oneofschema/one_of_schema.py index 7957f46..8d274bf 100644 --- a/marshmallow_oneofschema/one_of_schema.py +++ b/marshmallow_oneofschema/one_of_schema.py @@ -181,3 +181,25 @@ def validate(self, data, *, many=None, partial=None): except ValidationError as ve: return ve.messages return {} + + @staticmethod + def get_schema_name(schema_class): + """Key used for schema classes in type_schemas dict. The default is + schema class name stripped of the word 'Schema', if present. + """ + schema_class_name = schema_class.__name__ + if schema_class_name.endswith("Schema"): + schema_class_name = schema_class_name[: -len("Schema")] + return schema_class_name + + @classmethod + def register_one_of(cls, schema_class): + """A decorator to register schema class as one of the options for the + OneOfSchema without manually modifying type_schemas dict. Example: + + >>> @MyOneOfSchema.one_of + >>> class OneParticularSchema(Schema): + >>> ... + """ + cls.type_schemas[cls.get_schema_name(schema_class)] = schema_class + return schema_class diff --git a/tests/test_one_of_schema_building.py b/tests/test_one_of_schema_building.py new file mode 100644 index 0000000..4dcb7fb --- /dev/null +++ b/tests/test_one_of_schema_building.py @@ -0,0 +1,94 @@ +import marshmallow as m +import marshmallow.fields as f +from marshmallow_oneofschema import OneOfSchema + +from test_one_of_schema import Foo, Bar, Baz + + +class MySchemaWithDefaultNames(OneOfSchema): + type_schemas = {} + + +class MySchemaWithCustomNames(OneOfSchema): + type_schemas = {} + + counter = 0 + known_types = [Foo, Bar, Baz] + + def get_obj_type(self, obj): + return self.known_classes.index(obj.__class__) + + @classmethod + def get_schema_name(cls, schema_class): + cls.counter += 1 + return str(cls.counter - 1) + + +@MySchemaWithCustomNames.register_one_of +@MySchemaWithDefaultNames.register_one_of +class FooSchema(m.Schema): + value = f.String(required=True) + + @m.post_load + def make_foo(self, data, **kwargs): + return Foo(**data) + + +@MySchemaWithCustomNames.register_one_of +@MySchemaWithDefaultNames.register_one_of +class BarSchema(m.Schema): + value = f.Integer(required=True) + + @m.post_load + def make_bar(self, data, **kwargs): + return Bar(**data) + + +@MySchemaWithCustomNames.register_one_of +@MySchemaWithDefaultNames.register_one_of +class BazSchema(m.Schema): + value1 = f.Integer(required=True) + value2 = f.String(required=True) + + @m.post_load + def make_baz(self, data, **kwargs): + return Baz(**data) + + +class MyVerboseSchemaWithDefaultNames(OneOfSchema): + type_schemas = { + "Foo": FooSchema, + "Bar": BarSchema, + "Baz": BazSchema, + } + + +class MyVerboseSchemaWithCustomNames(OneOfSchema): + type_schemas = { + "0": FooSchema, + "1": BarSchema, + "2": BazSchema, + } + + +def test_schemas_building_with_register_one_of(): + assert ( + MySchemaWithDefaultNames.type_schemas + == MyVerboseSchemaWithDefaultNames.type_schemas + ) + assert ( + MySchemaWithCustomNames.type_schemas + == MyVerboseSchemaWithCustomNames.type_schemas + ) + + +def test_default_schema_naming(): + class SomeObjectSchema: + pass + + assert OneOfSchema.get_schema_name(SomeObjectSchema) == "SomeObject" + + class AnyOtherSchemaClass: + pass + + assert OneOfSchema.get_schema_name(AnyOtherSchemaClass) == "AnyOtherSchemaClass"