Repository: digitalocean/pynetbox (python)

Issue:
Add "count" functionality
There are times I want to get the total count of list data, without actually needing the dataset. This is particularly handy with filtered lists. If I was using `requests`, I would just use the `limit=1` parameter and parse out the value of `count` in the API response.

But there's no way to do this in pynetbox currently. I envision some sort of API like this:

```
>>> t1_circuits_count = nb.circuits.circuits.count(type_id=2)

>>> print(t1_circuits_count)
12
```

And behind the scenes it's just getting back a single item response, but still has the `count` key. This way I could get total circuits for each circuit type very quickly, without having to pull down a dataset of thousands of rows (and far more database queries).

Pull request:
Fixes #158: Adds count() method to Endpoint object.
Returns the count for a given query by returning only the count field
after running the query with `limit=1`.

Patch:
diff --git a/pynetbox/core/endpoint.py b/pynetbox/core/endpoint.py
index 48cff98..919e07f 100644
--- a/pynetbox/core/endpoint.py
+++ b/pynetbox/core/endpoint.py
@@ -341,6 +341,57 @@ class Endpoint(object):
 
         return self._choices
 
+    def count(self, *args, **kwargs):
+        r"""Returns the count of objects in a query.
+
+        Takes named arguments that match the usable filters on a
+        given endpoint. If an argument is passed then it's used as a
+        freeform search argument if the endpoint supports it. If no
+        arguments are passed the count for all objects on an endpoint
+        are returned.
+
+        :arg str,optional \*args: Freeform search string that's
+            accepted on given endpoint.
+        :arg str,optional \**kwargs: Any search argument the
+            endpoint accepts can be added as a keyword arg.
+
+        :Returns: Integer with count of objects returns by query.
+
+        :Examples:
+
+        To return a count of objects matching a named argument filter.
+
+        >>> nb.dcim.devices.count(site='tst1')
+        5827
+        >>>
+
+        To return a count of objects on an entire endpoint.
+
+        >>> nb.dcim.devices.count()
+        87382
+        >>>
+        """
+
+        if args:
+            kwargs.update({"q": args[0]})
+
+        if any(i in RESERVED_KWARGS for i in kwargs):
+            raise ValueError(
+                "A reserved {} kwarg was passed. Please remove it "
+                "try again.".format(RESERVED_KWARGS)
+            )
+
+        ret = Request(
+            filters=kwargs,
+            base=self.url,
+            token=self.token,
+            session_key=self.session_key,
+            ssl_verify=self.ssl_verify,
+            http_session=self.api.http_session,
+        )
+
+        return ret.get_count()
+
 
 class DetailEndpoint(object):
     """Enables read/write Operations on detail endpoints.
diff --git a/pynetbox/core/query.py b/pynetbox/core/query.py
index bc825ec..74d0f21 100644
--- a/pynetbox/core/query.py
+++ b/pynetbox/core/query.py
@@ -258,6 +258,25 @@ class Request(object):
 
         return url
 
+    def _make_call(self, url):
+
+        headers = {"accept": "application/json;"}
+        if self.token:
+            headers.update(authorization="Token {}".format(self.token))
+        if self.session_key:
+            headers.update({"X-Session-Key": self.session_key})
+
+        req = self.http_session.get(
+            url, headers=headers, verify=self.ssl_verify
+        )
+        if req.ok:
+            try:
+                return req.json()
+            except json.JSONDecodeError:
+                raise ContentError(req)
+        else:
+            raise RequestError(req)
+
     def get(self):
         """Makes a GET request.
 
@@ -270,26 +289,9 @@ class Request(object):
         :Returns: List of `Response` objects returned from the
             endpoint.
         """
-        headers = {"accept": "application/json;"}
-        if self.token:
-            headers.update(authorization="Token {}".format(self.token))
-        if self.session_key:
-            headers.update({"X-Session-Key": self.session_key})
-
-        def make_request(url):
-
-            req = self.http_session.get(url, headers=headers,
-                                        verify=self.ssl_verify)
-            if req.ok:
-                try:
-                    return req.json()
-                except json.JSONDecodeError:
-                    raise ContentError(req)
-            else:
-                raise RequestError(req)
 
         def req_all(url):
-            req = make_request(url)
+            req = self._make_call(url)
             if isinstance(req, dict) and req.get("results") is not None:
                 ret = req["results"]
                 first_run = True
@@ -304,7 +306,7 @@ class Request(object):
                         if first_run
                         else req["next"]
                     )
-                    req = make_request(next_url)
+                    req = self._make_call(next_url)
                     first_run = False
                     ret.extend(req["results"])
                 return ret
@@ -464,3 +466,19 @@ class Request(object):
                 raise ContentError(req)
         else:
             raise RequestError(req)
+
+    def get_count(self, *args, **kwargs):
+        """Returns object count for query
+
+        Makes a query to the endpoint with ``limit=1`` set and only
+        returns the value of the "count" field.
+
+        :raises: RequestError if req.ok returns false.
+        :raises: ContentError if response is not json.
+
+        :returns: Int of number of objects query returned.
+        """
+
+        return self._make_call(
+            "{}{}limit=1".format(self.url, "/?" if not self.filters else "&")
+        )["count"]


Test patch:
diff --git a/tests/unit/test_endpoint.py b/tests/unit/test_endpoint.py
index 1a4d6a6..f4c1ffd 100644
--- a/tests/unit/test_endpoint.py
+++ b/tests/unit/test_endpoint.py
@@ -5,9 +5,9 @@ import six
 from pynetbox.core.endpoint import Endpoint
 
 if six.PY3:
-    from unittest.mock import patch, Mock
+    from unittest.mock import patch, Mock, call
 else:
-    from mock import patch, Mock
+    from mock import patch, Mock, call
 
 
 class EndPointTestCase(unittest.TestCase):
@@ -38,3 +38,42 @@ class EndPointTestCase(unittest.TestCase):
         test_obj = Endpoint(api, app, "test")
         with self.assertRaises(ValueError) as _:
             test_obj.filter(id=1)
+
+    def test_count(self):
+        with patch(
+            "pynetbox.core.query.Request._make_call", return_value={"count": 42}
+        ) as mock:
+            api = Mock(base_url="http://localhost:8000/api")
+            app = Mock()
+            app.name = "test-app"
+            expected_call = call('http://localhost:8000/api/test-app/test-endpoint/?test_filter=test&limit=1')
+            test_obj = Endpoint(api, app, "test_endpoint")
+            test = test_obj.count(test_filter="test")
+            self.assertEqual(test, 42)
+            self.assertEqual(mock.call_args, expected_call)
+
+    def test_count_args(self):
+        with patch(
+            "pynetbox.core.query.Request._make_call", return_value={"count": 42}
+        ) as mock:
+            api = Mock(base_url="http://localhost:8000/api")
+            app = Mock()
+            app.name = "test-app"
+            expected_call = call('http://localhost:8000/api/test-app/test-endpoint/?q=test&limit=1')
+            test_obj = Endpoint(api, app, "test_endpoint")
+            test = test_obj.count("test")
+            self.assertEqual(test, 42)
+            self.assertEqual(mock.call_args, expected_call)
+
+    def test_count_no_args(self):
+        with patch(
+            "pynetbox.core.query.Request._make_call", return_value={"count": 42}
+        ) as mock:
+            api = Mock(base_url="http://localhost:8000/api")
+            app = Mock()
+            app.name = "test-app"
+            expected_call = call('http://localhost:8000/api/test-app/test-endpoint/?limit=1')
+            test_obj = Endpoint(api, app, "test_endpoint")
+            test = test_obj.count()
+            self.assertEqual(test, 42)
+            self.assertEqual(mock.call_args, expected_call)
diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py
new file mode 100644
index 0000000..e802503
--- /dev/null
+++ b/tests/unit/test_query.py
@@ -0,0 +1,55 @@
+import unittest
+
+import six
+
+from pynetbox.core.query import Request
+
+if six.PY3:
+    from unittest.mock import patch, Mock, call
+else:
+    from mock import patch, Mock, call
+
+
+class RequestTestCase(unittest.TestCase):
+    def test_get_count(self):
+        test_obj = Request(
+            http_session=Mock(),
+            base="http://localhost:8001/api/dcim/devices",
+            filters={"q": "abcd"},
+        )
+        test_obj.http_session.get.return_value.json.return_value = {
+            "count": 42,
+            "next": "http://localhost:8001/api/dcim/devices?limit=1&offset=1&q=abcd",
+            "previous": False,
+            "results": [],
+        }
+        expected = call(
+            "http://localhost:8001/api/dcim/devices/?q=abcd&limit=1",
+            headers={"accept": "application/json;"},
+            verify=True,
+        )
+        test_obj.http_session.get.ok = True
+        test = test_obj.get_count()
+        self.assertEqual(test, 42)
+        self.assertEqual(test_obj.http_session.get.call_args, expected)
+
+    def test_get_count_no_filters(self):
+        test_obj = Request(
+            http_session=Mock(),
+            base="http://localhost:8001/api/dcim/devices",
+        )
+        test_obj.http_session.get.return_value.json.return_value = {
+            "count": 42,
+            "next": "http://localhost:8001/api/dcim/devices?limit=1&offset=1",
+            "previous": False,
+            "results": [],
+        }
+        expected = call(
+            "http://localhost:8001/api/dcim/devices/?limit=1",
+            headers={"accept": "application/json;"},
+            verify=True,
+        )
+        test_obj.http_session.get.ok = True
+        test = test_obj.get_count()
+        self.assertEqual(test, 42)
+        self.assertEqual(test_obj.http_session.get.call_args, expected)