add property name based on data key

marshmallow_jsonschema/base.py

@@ -10,12 +10,10 @@
 
 from .validation import handle_length, handle_one_of, handle_range
 
-
 __all__ = (
     'JSONSchema',
 )
 
-
 TYPE_MAP = {
     dict: {
         'type': 'object',
@@ -72,7 +70,6 @@
     },
 }
 
-
 FIELD_VALIDATORS = {
     validate.Length: handle_length,
     validate.OneOf: handle_one_of,
@@ -91,6 +88,7 @@ def __init__(self, *args, **kwargs):
         """Setup internal cache of nested fields, to prevent recursion."""
         self._nested_schema_classes = {}
         self.nested = kwargs.pop('nested', False)
+        self.prefer_data_key = kwargs.pop('prefer_data_key', False)
         super(JSONSchema, self).__init__(*args, **kwargs)
 
     def _get_default_mapping(self, obj):
@@ -112,7 +110,7 @@ def get_properties(self, obj):
 
         for field_name, field in sorted(obj.fields.items()):
             schema = self._get_schema_for_field(obj, field)
-            properties[field.name] = schema
+            properties[self._get_property_name_for_field(field)] = schema
 
         return properties
 
@@ -122,16 +120,18 @@ def get_required(self, obj):
 
         for field_name, field in sorted(obj.fields.items()):
             if field.required:
-                required.append(field.name)
+                required.append(
+                    self._get_property_name_for_field(field)
+                )
 
         return required or missing
 
     def _from_python_type(self, obj, field, pytype):
         """Get schema definition from python type."""
         json_schema = {
-            'title': field.attribute or field.name,
+            'title': field.attribute or self._get_property_name_for_field(
+                field),
         }
-
         for key, val in TYPE_MAP[pytype].items():
             json_schema[key] = val
 
@@ -182,6 +182,20 @@ def _get_schema_for_field(self, obj, field):
                 )
         return schema
 
+    def _get_property_name_for_field(self, field):
+        """Get property name for field based on serialized object"""
+        name = field.name
+
+        if self.prefer_data_key:
+            # Handle change in load_from / dump_to between Marshmallow
+            # versions 2 and 3.
+            if marshmallow.__version__.split('.', 1)[0] >= '3':
+                name = field.data_key or name
+            else:
+                name = field.load_from or field.dump_to or name
+
+        return name
+
     def _from_nested_schema(self, obj, field):
         """Support nested field."""
         if isinstance(field.nested, basestring):

tests/test_dump.py

@@ -503,6 +503,7 @@ def test_use_datakey_for_title():
         class TestSchema(Schema):
             normal_field = fields.String()
             data_key = fields.String(data_key='dataKey')
+            required = fields.String(data_key='isRequired', required = True)
 
         schema = TestSchema()
         json_schema = JSONSchema(prefer_data_key=True)
@@ -516,13 +517,17 @@ class TestSchema(Schema):
         assert data_key is not None
         assert data_key.get('title') == 'dataKey'
 
+        required = dumped['definitions']['TestSchema']['required']
+        assert 'isRequired' in required
+
     # use load_from, then dump_to for marshmallow < v3
     else:
         class TestSchema(Schema):
             normal_field = fields.String()
             load_from = fields.String(load_from='loadFrom')
             dump_to = fields.String(dump_to='dumpTo')
             prefer_load_from = fields.String(load_from='both', dump_to='dumpTo')
+            required = fields.String(dump_to='isRequired', required=True)
 
         schema = TestSchema()
         json_schema = JSONSchema(prefer_data_key=True)
@@ -533,6 +538,8 @@ class TestSchema(Schema):
         assert normal_field.get('title') == 'normal_field'
 
         load_from = dumped['definitions']['TestSchema']['properties'].get('loadFrom')
+        import json
+        print(json.dumps(dumped, indent=4))
         assert load_from is not None
         assert load_from.get('title') == 'loadFrom'
 
@@ -543,3 +550,6 @@ class TestSchema(Schema):
         prefer_load_from = dumped['definitions']['TestSchema']['properties'].get('both')
         assert prefer_load_from is not None
         assert prefer_load_from.get('title') == 'both'
+
+        required = dumped['definitions']['TestSchema']['required']
+        assert 'isRequired' in required