Allow providing custom fields mappings

marshmallow_jsonschema/base.py

@@ -90,6 +90,12 @@ def __init__(self, *args, **kwargs):
         self.nested = kwargs.pop('nested', False)
         super(JSONSchema, self).__init__(*args, **kwargs)
 
+    def get_custom_mappings(self):
+        """Return a dict of field class for keys, type for values for any
+        custom types.
+        """
+        return None
+
     def _get_default_mapping(self, obj):
         """Return default mapping if there are no special needs."""
         mapping = {v: k for k, v in obj.TYPE_MAPPING.items()}
@@ -101,11 +107,13 @@ def _get_default_mapping(self, obj):
             fields.LocalDateTime: datetime.datetime,
             fields.Nested: '_from_nested_schema',
         })
+        custom_mappings = self.get_custom_mappings()
+        if custom_mappings:
+            mapping.update(self.get_custom_mappings())
         return mapping
 
     def get_properties(self, obj):
         """Fill out properties field."""
-        mapping = self._get_default_mapping(obj)
         properties = {}
 
         for field_name, field in sorted(obj.fields.items()):
@@ -196,7 +204,7 @@ def _from_nested_schema(self, obj, field):
         # If this is not a schema we've seen, and it's not this schema,
         # put it in our list of schema defs
         if name not in self._nested_schema_classes and name != outer_name:
-            wrapped_nested = JSONSchema(nested=True)
+            wrapped_nested = self.__class__(nested=True)
             wrapped_dumped = wrapped_nested.dump(
                 nested()
             )

tests/test_dump.py

@@ -1,4 +1,5 @@
 from marshmallow import Schema, fields, validate
+from marshmallow.compat import text_type
 from marshmallow_jsonschema import JSONSchema
 from jsonschema import Draft4Validator
 import pytest
@@ -266,6 +267,7 @@ class TestSchema(Schema):
         'title'
     ] == 'Brown Cowzz'
 
+
 def test_unknown_typed_field_throws_valueerror():
 
     class Invalid(fields.Field):
@@ -280,6 +282,7 @@ class UserSchema(Schema):
     with pytest.raises(ValueError):
         json_schema.dump(schema).data
 
+
 def test_unknown_typed_field():
 
     class Colour(fields.Field):
@@ -340,3 +343,68 @@ class TestSchema(Schema):
         'type': 'string',
         'description': 'Directly on the field!',
     }
+
+
+def test_custom_field_by_subclassing():
+
+    class Colour(fields.String):
+        def __init__(self, default='red', **kwargs):
+            super(Colour, self).__init__(default=default, **kwargs)
+
+    class UserSchema(Schema):
+        name = fields.String(required=True)
+        favourite_colour = Colour()
+
+    schema = UserSchema()
+    json_schema = JSONSchema()
+
+    with pytest.raises(ValueError):
+        # The custom Color field is not registered
+        dumped = json_schema.dump(schema)
+
+    # Provide the custom field to the default mappings
+    class CustomJSONSchema(JSONSchema):
+        def get_custom_mappings(self):
+            return {Colour:  text_type}
+
+    # Register the field
+    json_schema = CustomJSONSchema()
+    dumped = json_schema.dump(schema).data
+    assert dumped['definitions']['UserSchema']['properties']['favourite_colour'] == {
+        'default': 'red',
+        'title': 'favourite_colour',
+        'type': 'string'}
+
+
+def test_nested_custom_field_by_subclassing():
+
+    class Colour(fields.String):
+        def __init__(self, default='red', **kwargs):
+            super(Colour, self).__init__(default=default, **kwargs)
+
+    class ColoursSchema(Schema):
+        favourite_colour = Colour()
+
+    class UserSchema(Schema):
+        name = fields.String(required=True)
+        colours = fields.Nested(ColoursSchema)
+
+    schema = UserSchema()
+    json_schema = JSONSchema()
+
+    with pytest.raises(ValueError):
+        # The custom Color field is not registered
+        dumped = json_schema.dump(schema)
+
+    # Provide the custom field to the default mappings
+    class CustomJSONSchema(JSONSchema):
+        def get_custom_mappings(self):
+            return {Colour:  text_type}
+
+    # Register the field
+    json_schema = CustomJSONSchema()
+    dumped = json_schema.dump(schema).data
+    assert dumped['definitions']['ColoursSchema']['properties']['favourite_colour'] == {
+        'default': 'red',
+        'title': 'favourite_colour',
+        'type': 'string'}