Merge pull request #37 from stringfellow/feature/support-functions

Nesting, recursion, fields.Function

.gitignore

@@ -1,3 +1,5 @@
+*.sw[op]
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

marshmallow_jsonschema/base.py

@@ -5,6 +5,7 @@
 from marshmallow import fields, missing, Schema, validate
 from marshmallow.class_registry import get_class
 from marshmallow.compat import text_type, binary_type, basestring
+from marshmallow.decorators import post_dump
 
 from .validation import handle_length, handle_one_of, handle_range
 
@@ -77,42 +78,44 @@
 
 
 class JSONSchema(Schema):
+    """Converts to JSONSchema as defined by http://json-schema.org/."""
+
     properties = fields.Method('get_properties')
     type = fields.Constant('object')
     required = fields.Method('get_required')
 
-    def get_properties(self, obj):
+    def __init__(self, *args, **kwargs):
+        """Setup internal cache of nested fields, to prevent recursion."""
+        self._nested_schema_classes = {}
+        self.nested = kwargs.pop('nested', False)
+        super(JSONSchema, self).__init__(*args, **kwargs)
+
+    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()}
-        mapping[fields.Email] = text_type
-        mapping[fields.Dict] = dict
-        mapping[fields.List] = list
-        mapping[fields.Url] = text_type
-        mapping[fields.LocalDateTime] = datetime.datetime
+        mapping.update({
+            fields.Email: text_type,
+            fields.Dict: dict,
+            fields.Url: text_type,
+            fields.List: list,
+            fields.LocalDateTime: datetime.datetime,
+            fields.Nested: '_from_nested_schema',
+        })
+        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()):
-            if hasattr(field, '_jsonschema_type_mapping'):
-                schema = field._jsonschema_type_mapping()
-            elif field.__class__ in mapping:
-                pytype = mapping[field.__class__]
-                schema = self.__class__._from_python_type(field, pytype)
-            elif isinstance(field, fields.Nested):
-                schema = self.__class__._from_nested_schema(field)
-            else:
-                raise ValueError('unsupported field type %s' % field)
-
-            # Apply any and all validators that field may have
-            for validator in field.validators:
-                if validator.__class__ in FIELD_VALIDATORS:
-                    schema = FIELD_VALIDATORS[validator.__class__](
-                        schema, field, validator, obj
-                    )
-
+            schema = self._get_schema_for_field(obj, field)
             properties[field.name] = schema
 
         return properties
 
     def get_required(self, obj):
+        """Fill out required field."""
         required = []
 
         for field_name, field in sorted(obj.fields.items()):
@@ -121,8 +124,8 @@ def get_required(self, obj):
 
         return required
 
-    @classmethod
-    def _from_python_type(cls, field, pytype):
+    def _from_python_type(self, obj, field, pytype):
+        """Get schema definition from python type."""
         json_schema = {
             'title': field.attribute or field.name,
         }
@@ -145,21 +148,64 @@ def _from_python_type(cls, field, pytype):
             json_schema['title'] = field.metadata['metadata'].get('title')
 
         if isinstance(field, fields.List):
-            for ptyp, mtyp in cls.TYPE_MAPPING.items():
-                if isinstance(field.container, mtyp):
-                    json_schema['items'] = cls()._from_python_type(
-                        field.container, ptyp)
-                    break
-
+            json_schema['items'] = self._get_schema_for_field(
+                obj, field.container
+            )
         return json_schema
 
-    @classmethod
-    def _from_nested_schema(cls, field):
+    def _get_schema_for_field(self, obj, field):
+        """Get schema and validators for field."""
+        mapping = self._get_default_mapping(obj)
+        if hasattr(field, '_jsonschema_type_mapping'):
+            schema = field._jsonschema_type_mapping()
+        elif '_jsonschema_type_mapping' in field.metadata:
+            schema = field.metadata['_jsonschema_type_mapping']
+        elif field.__class__ in mapping:
+            pytype = mapping[field.__class__]
+            if isinstance(pytype, basestring):
+                schema = getattr(self, pytype)(obj, field)
+            else:
+                schema = self._from_python_type(
+                    obj, field, pytype
+                )
+        else:
+            raise ValueError('unsupported field type %s' % field)
+
+        # Apply any and all validators that field may have
+        for validator in field.validators:
+            if validator.__class__ in FIELD_VALIDATORS:
+                schema = FIELD_VALIDATORS[validator.__class__](
+                    schema, field, validator, obj
+                )
+        return schema
+
+    def _from_nested_schema(self, obj, field):
+        """Support nested field."""
         if isinstance(field.nested, basestring):
             nested = get_class(field.nested)
         else:
             nested = field.nested
-        schema = cls().dump(nested()).data
+
+        name = nested.__name__
+        outer_name = obj.__class__.__name__
+
+        # 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_dumped = wrapped_nested.dump(
+                nested()
+            )
+            self._nested_schema_classes[name] = wrapped_dumped.data
+            self._nested_schema_classes.update(
+                wrapped_nested._nested_schema_classes
+            )
+
+        # and the schema is just a reference to the def
+        schema = {
+            'type': 'object',
+            '$ref': '#/definitions/{}'.format(name)
+        }
 
         if field.metadata.get('metadata', {}).get('description'):
             schema['description'] = (
@@ -176,3 +222,22 @@ def _from_nested_schema(cls, field):
             }
 
         return schema
+
+    def dump(self, obj, **kwargs):
+        """Take obj for later use: using class name to namespace definition."""
+        self.obj = obj
+        return super(JSONSchema, self).dump(obj, **kwargs)
+
+    @post_dump(pass_many=False)
+    def wrap(self, data):
+        """Wrap this with the root schema definitions."""
+        if self.nested:  # no need to wrap, will be in outer defs
+            return data
+
+        name = self.obj.__class__.__name__
+        self._nested_schema_classes[name] = data
+        root = {
+            'definitions': self._nested_schema_classes,
+            '$ref': '#/definitions/{}'.format(name)
+        }
+        return root