Skip to content

Commit

Permalink
Allow registering field subclasses.
Browse files Browse the repository at this point in the history
This will allow sub-classing fields and allow marshmallow_jsonschema to
render them as the parent class would.
  • Loading branch information
s0undt3ch committed Jan 13, 2017
1 parent 6f51771 commit 96a52c8
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 5 deletions.
20 changes: 15 additions & 5 deletions marshmallow_jsonschema/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,17 +77,27 @@


class JSONSchema(Schema):

JSONSchemaTypesMapping = {
fields.Email: text_type,
fields.Dict: dict,
fields.List: list,
fields.Url: text_type,
fields.LocalDateTime: datetime.datetime
}

properties = fields.Method('get_properties')
type = fields.Constant('object')
required = fields.Method('get_required')

@classmethod
def register_jsonschema_field(cls, field, field_type):
cls.JSONSchemaTypesMapping[field] = field_type

def get_properties(self, obj):
mapping = {v: k for k, v in obj.TYPE_MAPPING.items()}
mapping[fields.Email] = text_type
mapping[fields.Dict] = dict
mapping[fields.List] = list
mapping[fields.Url] = text_type
mapping[fields.LocalDateTime] = datetime.datetime
for field, field_type in self.JSONSchemaTypesMapping.items():
mapping[field] = field_type
properties = {}

for field_name, field in sorted(obj.fields.items()):
Expand Down
31 changes: 31 additions & 0 deletions tests/test_dump.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -154,6 +155,7 @@ class TestSchema(Schema):
_validate_schema(dumped)
assert dumped['properties']['myfield']['title'] == 'Brown Cowzz'


def test_unknown_typed_field_throws_valueerror():

class Invalid(fields.Field):
Expand All @@ -168,6 +170,7 @@ class UserSchema(Schema):
with pytest.raises(ValueError):
json_schema.dump(schema).data


def test_unknown_typed_field():

class Colour(fields.Field):
Expand All @@ -192,3 +195,31 @@ class UserSchema(Schema):
json_schema = JSONSchema()
dumped = json_schema.dump(schema).data
assert dumped['properties']['favourite_colour'] == {'type': 'string'}


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)

# Register the field
JSONSchema.register_jsonschema_field(Colour, text_type)
json_schema = JSONSchema()
dumped = json_schema.dump(schema).data
try:
assert dumped['properties']['favourite_colour'] == {'default': 'red', 'title': 'favourite_colour', 'type': 'string'}
finally:
# Reset the default mapping dictionary
JSONSchema.JSONSchemaTypesMapping.pop(Colour, None)

0 comments on commit 96a52c8

Please sign in to comment.