Improve typehint

marshmallow_dataclass/__init__.py

@@ -161,9 +161,21 @@ def decorator(clazz: Type[_U]) -> Type[_U]:
     return decorator(_cls) if _cls else decorator
 
 
-def class_schema(
-    clazz: type, base_schema: Optional[Type[marshmallow.Schema]] = None
-) -> Type[marshmallow.Schema]:
+T_Schema = TypeVar("T_Schema", bound=marshmallow.Schema)
+T_SchemaType = Type[T_Schema]
+
+
+@overload
+def class_schema(clazz: type, base_schema: None = None) -> Type[marshmallow.Schema]:
+    ...
+
+
+@overload
+def class_schema(clazz: type, base_schema: T_SchemaType) -> T_SchemaType:
+    ...
+
+
+def class_schema(clazz, base_schema=None):
 
     """
     Convert a class to a marshmallow schema

tests/test_schema_type.py

@@ -0,0 +1,30 @@
+import unittest
+
+import marshmallow
+
+import marshmallow_dataclass
+
+
+class MyBaseSchema(marshmallow.Schema):
+    pass
+
+
+class A:
+    a: int
+
+
+# Top-level so that this is checked by mypy
+schema_with_base: MyBaseSchema = marshmallow_dataclass.class_schema(
+    A, base_schema=MyBaseSchema
+)()
+
+schema_without_base: marshmallow.Schema = marshmallow_dataclass.class_schema(A)()
+
+
+# Regression test for https://github.com/lovasoa/marshmallow_dataclass/pull/77
+class TestSchemaType(unittest.TestCase):
+    def test_custom_basechema_type(self):
+        self.assertIsInstance(schema_with_base, MyBaseSchema)
+
+    def test_no_basechema_type(self):
+        self.assertIsInstance(schema_without_base, marshmallow.Schema)