Decorator for adding member Schema classes to OneOfSchema

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
 -------
 

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

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"