Add self and related URL overrides to Schema

Author: RazerM

marshmallow_jsonapi/flask.py

@@ -60,6 +60,42 @@ class Meta:
         """
         pass
 
+    def __init__(self, *args, **kwargs):
+        if kwargs.get('self_url', None):
+            raise TypeError('Use `self_view` instead of `self_url` '
+                            'using the Flask extension.')
+        if kwargs.get('self_url_kwargs', None):
+            raise TypeError('Use `self_view_kwargs` instead of '
+                            '`self_url_kwargs` when using the Flask extension.')
+        if kwargs.get('related_url', None):
+            raise TypeError('Use `related_view` instead of `related_url` '
+                             'using the Flask extension.')
+        if kwargs.get('related_url_kwargs', None):
+            raise TypeError('Use `related_view_kwargs` instead of '
+                            '`related_url_kwargs` when using the Flask extension.')
+
+        self_view = kwargs.pop('self_view', None)
+        self_view_kwargs = kwargs.pop('self_view_kwargs', dict())
+        related_view = kwargs.pop('related_view', None)
+        related_view_kwargs = kwargs.pop('related_view_kwargs', dict())
+
+        if self_view_kwargs and not self_view:
+            raise ValueError('Must specify `self_view` keyword when '
+                             '`self_view_kwargs` is specified.')
+
+        if related_view_kwargs and not related_view:
+            raise ValueError('Must specify `related_view` keyword when '
+                             '`related_view_kwargs` is specified.')
+
+        if self_view:
+            kwargs['self_url'] = flask.url_for(self_view, **self_view_kwargs)
+        if related_view:
+            kwargs['related_url'] = flask.url_for(related_view,
+                                                  **related_view_kwargs)
+
+        super(Schema, self).__init__(*args, **kwargs)
+
+
     def generate_url(self, view_name, **kwargs):
         """Generate URL with any kwargs interpolated."""
         return flask.url_for(view_name, **kwargs) if view_name else None

marshmallow_jsonapi/schema.py

@@ -80,6 +80,8 @@ class Meta:
     def __init__(self, *args, **kwargs):
         self.include_data = kwargs.pop('include_data', ())
         self.resource_identifier = kwargs.pop('resource_identifier', False)
+        self.override_self_url = kwargs.pop('self_url', None)
+        self.override_related_url = kwargs.pop('related_url', None)
         super(Schema, self).__init__(*args, **kwargs)
         if self.include_data:
             self.check_relations(self.include_data)
@@ -370,10 +372,21 @@ def get_top_level_links(self, data, many):
             if self.opts.self_url_many:
                 self_link = self.generate_url(self.opts.self_url_many)
         else:
-            if self.opts.self_url:
+            if self.opts.self_url and data is not None:
                 self_link = data.get('links', {}).get('self', None)
 
-        return {'self': self_link}
+        if self.override_self_url is not None:
+            self_link = self.override_self_url
+
+        links = dict()
+
+        if self_link is not None:
+            links['self'] = self_link
+
+        if self.override_related_url:
+            links['related'] = self.override_related_url
+
+        return links
 
     def get_resource_links(self, item):
         """Hook for adding links to a resource object."""
@@ -391,7 +404,7 @@ def wrap_response(self, data, many):
         # may only be included if there is data in the ret
         if many or data:
             top_level_links = self.get_top_level_links(data, many)
-            if top_level_links['self']:
+            if top_level_links:
                 ret['links'] = top_level_links
         return ret
 

tests/test_flask.py

@@ -97,6 +97,37 @@ def test_self_link_many(self, app, posts):
         assert 'links' in data['data'][0]
         assert data['data'][0]['links']['self'] == '/posts/{}/'.format(posts[0].id)
 
+    @pytest.mark.parametrize('kwargs, exc, msg', [
+        (dict(self_url='foo'), TypeError, 'self_view'),
+        (dict(self_url_kwargs='foo'), TypeError, 'self_view_kwargs'),
+        (dict(related_url='foo'), TypeError, 'related_view'),
+        (dict(related_url_kwargs='foo'), TypeError, 'related_view_kwargs'),
+        (dict(self_view_kwargs='foo'), ValueError, 'self_view'),
+        (dict(related_view_kwargs='foo'), ValueError, 'related_view'),
+    ])
+    def test_schema_view_override_exceptions(self, kwargs, exc, msg):
+        with pytest.raises(exc) as exc_info:
+            self.PostFlaskSchema(**kwargs)
+
+        assert msg in exc_info.value.args[0]
+
+    def test_view_override(self, app):
+        # Here we use "wrong" views just to check output.
+
+        schema = self.PostFlaskSchema(self_view='posts')
+        assert schema.override_self_url == '/posts/'
+
+        schema = self.PostFlaskSchema(
+            self_view='author_detail', self_view_kwargs={'author_id': 1})
+        assert schema.override_self_url == '/authors/1'
+
+        schema = self.PostFlaskSchema(related_view='posts')
+        assert schema.override_related_url == '/posts/'
+
+        schema = self.PostFlaskSchema(
+            related_view='author_detail', related_view_kwargs={'author_id': 1})
+        assert schema.override_related_url == '/authors/1'
+
     def test_schema_with_empty_relationship(self, app, post_with_null_author):
         data = self.PostAuthorFlaskSchema().dump(post_with_null_author).data
         assert 'relationships' not in data

tests/test_schema.py

@@ -20,7 +20,14 @@ def get_top_level_links(self, data, many):
             self_link = '/authors/'
         else:
             self_link = '/authors/{}'.format(data['id'])
-        return {'self': self_link}
+
+        links = {'self': self_link}
+
+        schema_links = super(AuthorSchema, self).get_top_level_links(data, many)
+        links.update(schema_links)
+
+        return links
+
 
     class Meta:
         type_ = 'people'
@@ -855,3 +862,24 @@ def test_deserializing_relationship_errors(self):
         assert \
             self._assert_relationship_error('comments', errors['errors']) is \
             True
+
+
+class TestOverrideUrls(object):
+    def test_self_url(self, author):
+        url = '/articles/1/author'
+        schema = AuthorSchema(self_url=url)
+        data = schema.dump(author).data
+
+        assert data['links']['self'] == url
+
+    def test_related_url(self, author):
+        self_url = '/articles/1/relationships/author'
+        related_url = '/articles/1/author'
+        schema = AuthorSchema(
+            resource_identifier=True,
+            self_url=self_url,
+            related_url=related_url)
+        data = schema.dump(author).data
+
+        assert data['links']['self'] == self_url
+        assert data['links']['related'] == related_url