added get(x) to query /api/x

Author: James Lin

rest_framework_queryset/queryset.py

@@ -14,7 +14,7 @@ def __init__(self, url, *args, **kwargs):
         self.args = args
         self.kwargs = kwargs
 
-    def call_api(self):
+    def _call_api(self):
         """
         perform api call
         """
@@ -59,6 +59,10 @@ def page_result(self):
 class RestFrameworkQuerySet(BaseAPIQuerySet):
     page_size = 100
 
+    def __init__(self, *args, **kwargs):
+        super(RestFrameworkQuerySet, self).__init__(*args, **kwargs)
+        self.__id = None
+
     def __iter__(self):
         paginator = Paginator(self, self.page_size)
         for page in paginator.page_range:
@@ -70,19 +74,26 @@ def _clone(self):
         cloned.page_size = self.page_size
         return cloned
 
+    def _call_api(self):
+        if self.__id:
+            self.url = '{}/{}'.format(self.url.rstrip('/'), self.__id)
+        return super(RestFrameworkQuerySet, self)._call_api()
+
     def count(self):
         cloned = self._clone()
         params = cloned.kwargs.get('params', {})
         params['offset'] = 0
         params['limit'] = 0
-        resp = cloned.call_api()
+        resp = cloned._call_api()
         result = resp.json()
         return result['count']
 
     def get_result(self):
-        response = self.call_api()
+        response = self._call_api()
         result = response.json()
-        return result['results']
+        if 'results' in result:
+            return result['results']
+        return result
 
     def page_result(self, slicer):
         cloned = self._clone()
@@ -97,12 +108,15 @@ def filter(self, **kwargs):
         params.update(kwargs)
         return cloned
 
-    def get(self, **kwargs):
+    def get(self, __id=None, **kwargs):
         cloned = self.filter(**kwargs)
+        cloned.__id = __id
         result = cloned.get_result()
-        if len(result) > 1:
-            raise MultipleObjectsReturned('get() returned more than one result, it returned {}'.format(cloned.count()))
-        return result[0]
+        if isinstance(result, list):
+            if len(result) > 1:
+                raise MultipleObjectsReturned('get() returned more than one result, it returned {}'.format(cloned.count()))
+            return result[0]
+        return result
 
     def all(self):
         return self._clone()

tests.py

@@ -41,14 +41,22 @@ def test_filter_call(self):
             mock_get.assert_any_call('/api/', params={'a': 123, 'b': 234, 'offset': 0, 'limit': 10})
 
     def test_get_call(self):
-        fake_response = MagicMock(json=lambda:{'count': 10, 'results': range(10)})
+        fake_response = MagicMock(json=lambda:{'count': 10, 'results': list(range(10))})
         with patch('rest_framework_queryset.queryset.requests.Session.get', return_value=fake_response) as mock_get:
             qs = RestFrameworkQuerySet('/api/')
             with self.assertRaises(MultipleObjectsReturned):
                 qs1 = qs.get(a=123)
             self.assertEqual(list(qs), list(range(10)))
             mock_get.assert_any_call('/api/', params={'a': 123})
 
+    def test_get_call_by_id(self):
+        fake_response = MagicMock(json=lambda:{'a': 123})
+        with patch('rest_framework_queryset.queryset.requests.Session.get', return_value=fake_response) as mock_get:
+            qs = RestFrameworkQuerySet('/api/')
+            qs1 = qs.get(123)
+            self.assertEqual(qs1, {'a': 123})
+            mock_get.assert_any_call('/api/123', params={})
+
     def test_count_call(self):
         fake_response = MagicMock(json=lambda:{'count': 10, 'results': range(10)})
         with patch('rest_framework_queryset.queryset.requests.Session.get', return_value=fake_response) as mock_get: