Merge branch 'master' into master.tracing-fix

Author: zmoody

.github/workflows/py2.yml

@@ -148,6 +148,12 @@ class PluginsApp(object):
     def __init__(self, api):
         self.api = api
 
+    def __getstate__(self):
+        return self.__dict__
+
+    def __setstate__(self, d):
+        self.__dict__.update(d)
+
     def __getattr__(self, name):
         return App(self.api, "plugins/{}".format(name))
 

pynetbox/core/app.py

@@ -14,15 +14,9 @@
 limitations under the License.
 """
 from pynetbox.core.query import Request, RequestError
-from pynetbox.core.response import Record
+from pynetbox.core.response import Record, RecordSet
 
-RESERVED_KWARGS = ("id", "pk", "limit", "offset")
-
-
-def response_loader(req, return_obj, endpoint):
-    if isinstance(req, list):
-        return [return_obj(i, endpoint.api, endpoint) for i in req]
-    return return_obj(req, endpoint.api, endpoint)
+RESERVED_KWARGS = ("offset",)
 
 
 class Endpoint(object):
@@ -77,11 +71,14 @@ def _lookup_ret_obj(self, name, model):
             ret = Record
         return ret
 
-    def all(self):
+    def all(self, limit=0):
         """Queries the 'ListView' of a given endpoint.
 
         Returns all objects from an endpoint.
 
+        :arg int,optional limit: Overrides the max page size on
+            paginated returns.
+
         :Returns: List of :py:class:`.Record` objects.
 
         :Examples:
@@ -96,9 +93,10 @@ def all(self):
             session_key=self.session_key,
             http_session=self.api.http_session,
             threading=self.api.threading,
+            limit=limit,
         )
 
-        return response_loader(req.get(), self.return_obj, self)
+        return RecordSet(self, req)
 
     def get(self, *args, **kwargs):
         r"""Queries the DetailsView of a given endpoint.
@@ -135,17 +133,7 @@ def get(self, *args, **kwargs):
             key = None
 
         if not key:
-            filter_lookup = self.filter(**kwargs)
-            if filter_lookup:
-                if len(filter_lookup) > 1:
-                    raise ValueError(
-                        "get() returned more than one result. "
-                        "Check that the kwarg(s) passed are valid for this "
-                        "endpoint or use filter() or all() instead."
-                    )
-                else:
-                    return filter_lookup[0]
-            return None
+            return next(iter(self.filter(**kwargs)), None)
 
         req = Request(
             key=key,
@@ -154,17 +142,14 @@ def get(self, *args, **kwargs):
             session_key=self.session_key,
             http_session=self.api.http_session,
         )
-
         try:
-            resp = req.get()
+            return next(iter(RecordSet(self, req)), None)
         except RequestError as e:
             if e.req.status_code == 404:
                 return None
             else:
                 raise e
 
-        return response_loader(resp, self.return_obj, self)
-
     def filter(self, *args, **kwargs):
         r"""Queries the 'ListView' of a given endpoint.
 
@@ -176,6 +161,8 @@ def filter(self, *args, **kwargs):
             accepted on given endpoint.
         :arg str,optional \**kwargs: Any search argument the
             endpoint accepts can be added as a keyword arg.
+        :arg int,optional limit: Overrides the max page size on
+            paginated returns.
 
         :Returns: A list of :py:class:`.Record` objects.
 
@@ -214,8 +201,8 @@ def filter(self, *args, **kwargs):
             raise ValueError("filter must be passed kwargs. Perhaps use all() instead.")
         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)
+                "A reserved kwarg was passed ({}). Please remove it "
+                "and try again.".format(RESERVED_KWARGS)
             )
 
         req = Request(
@@ -225,9 +212,10 @@ def filter(self, *args, **kwargs):
             session_key=self.session_key,
             http_session=self.api.http_session,
             threading=self.api.threading,
+            limit=kwargs.get("limit", 0),
         )
 
-        return response_loader(req.get(), self.return_obj, self)
+        return RecordSet(self, req)
 
     def create(self, *args, **kwargs):
         r"""Creates an object on an endpoint.
@@ -287,7 +275,7 @@ def create(self, *args, **kwargs):
             http_session=self.api.http_session,
         ).post(args[0] if args else kwargs)
 
-        return response_loader(req, self.return_obj, self)
+        return self.return_obj(req, self.api, self)
 
     def choices(self):
         """ Returns all choices from the endpoint.
@@ -425,7 +413,10 @@ def list(self, **kwargs):
         req = Request(**self.request_kwargs).get(add_params=kwargs)
 
         if self.custom_return:
-            return response_loader(req, self.custom_return, self.parent_obj.endpoint)
+            for i in req:
+                yield self.custom_return(
+                    i, self.parent_obj.endpoint.api, self.parent_obj.endpoint
+                )
         return req
 
     def create(self, data=None):
@@ -444,7 +435,9 @@ def create(self, data=None):
         data = data or {}
         req = Request(**self.request_kwargs).post(data)
         if self.custom_return:
-            return response_loader(req, self.custom_return, self.parent_obj.endpoint)
+            return self.custom_return(
+                req, self.parent_obj.endpoint.api, self.parent_obj.endpoint
+            )
         return req
 
 

pynetbox/core/endpoint.py

@@ -130,6 +130,7 @@ def __init__(
         base,
         http_session,
         filters=None,
+        limit=None,
         key=None,
         token=None,
         private_key=None,
@@ -158,6 +159,7 @@ def __init__(
         self.http_session = http_session
         self.url = self.base if not key else "{}{}/".format(self.base, key)
         self.threading = threading
+        self.limit = limit
 
     def get_openapi(self):
         """ Gets the OpenAPI Spec """
@@ -305,36 +307,12 @@ def get(self, add_params=None):
             endpoint.
         """
 
-        def req_all():
-            req = self._make_call(add_params=add_params)
-            if isinstance(req, dict) and req.get("results") is not None:
-                ret = req["results"]
-                first_run = True
-                while req["next"]:
-                    # Not worrying about making sure add_params kwargs is
-                    # passed in here because results from detail routes aren't
-                    # paginated, thus far.
-                    if first_run:
-                        req = self._make_call(
-                            add_params={
-                                "limit": req["count"],
-                                "offset": len(req["results"]),
-                            }
-                        )
-                    else:
-                        req = self._make_call(url_override=req["next"])
-                    first_run = False
-                    ret.extend(req["results"])
-                return ret
-            else:
-                return req
-
-        def req_all_threaded(add_params):
-            if add_params is None:
-                # Limit must be 0 to discover the max page size
-                add_params = {"limit": 0}
-            req = self._make_call(add_params=add_params)
-            if isinstance(req, dict) and req.get("results") is not None:
+        if not add_params and self.limit is not None:
+            add_params = {"limit": self.limit}
+        req = self._make_call(add_params=add_params)
+        if isinstance(req, dict) and req.get("results") is not None:
+            self.count = req["count"]
+            if self.threading:
                 ret = req["results"]
                 if req.get("next"):
                     page_size = len(req["results"])
@@ -347,15 +325,30 @@ def req_all_threaded(add_params):
                         ret.extend(req["results"])
                     else:
                         self.concurrent_get(ret, page_size, page_offsets)
-
-                return ret
-            else:
-                return req
-
-        if self.threading:
-            return req_all_threaded(add_params)
-
-        return req_all()
+                for i in ret:
+                    yield i
+            first_run = True
+            for i in req["results"]:
+                yield i
+            while req["next"]:
+                # Not worrying about making sure add_params kwargs is
+                # passed in here because results from detail routes aren't
+                # paginated, thus far.
+                if first_run:
+                    req = self._make_call(
+                        add_params={
+                            "limit": self.limit or req["count"],
+                            "offset": len(req["results"]),
+                        }
+                    )
+                else:
+                    req = self._make_call(url_override=req["next"])
+                first_run = False
+                for i in req["results"]:
+                    yield i
+        else:
+            self.count = len(req)
+            yield req
 
     def put(self, data):
         """Makes PUT request.
@@ -438,4 +431,6 @@ def get_count(self, *args, **kwargs):
         :returns: Int of number of objects query returned.
         """
 
-        return self._make_call(add_params={"limit": 1})["count"]
+        if not hasattr(self, "count"):
+            self.count = self._make_call(add_params={"limit": 1})["count"]
+        return self.count

pynetbox/core/query.py

@@ -18,7 +18,7 @@
 
 import pynetbox.core.app
 from six.moves.urllib.parse import urlsplit
-from pynetbox.core.query import Request
+from pynetbox.core.query import Request, RequestError
 from pynetbox.core.util import Hashabledict
 
 
@@ -70,6 +70,27 @@ class JsonField(object):
     _json_field = True
 
 
+class RecordSet(object):
+    def __init__(self, endpoint, request, **kwargs):
+        self.endpoint = endpoint
+        self.request = request
+        self.response = self.request.get()
+        self._response_cache = []
+
+    def __iter__(self):
+        for i in self._response_cache:
+            yield self.endpoint.return_obj(i, self.endpoint.api, self.endpoint)
+        for i in self.response:
+            yield self.endpoint.return_obj(i, self.endpoint.api, self.endpoint)
+
+    def __len__(self):
+        try:
+            return self.request.count
+        except AttributeError:
+            self._response_cache.append(next(self.response))
+            return self.request.count
+
+
 class Record(object):
     """Create python objects from netbox API responses.
 
@@ -311,7 +332,7 @@ def full_details(self):
                 session_key=self.api.session_key,
                 http_session=self.api.http_session,
             )
-            self._parse_values(req.get())
+            self._parse_values(next(req.get()))
             self.has_details = True
             return True
         return False

pynetbox/core/response.py

@@ -24,21 +24,23 @@
 
 class TraceableRecord(Record):
     def trace(self):
-        req = Request(
-            key=str(self.id) + "/trace",
-            base=self.endpoint.url,
-            token=self.api.token,
-            session_key=self.api.session_key,
-            http_session=self.api.http_session,
-        )
+        req = list(
+            Request(
+                key=str(self.id) + "/trace",
+                base=self.endpoint.url,
+                token=self.api.token,
+                session_key=self.api.session_key,
+                http_session=self.api.http_session,
+            ).get()
+        )[0]
         uri_to_obj_class_map = {
             "dcim/cables": Cables,
             "dcim/front-ports": FrontPorts,
             "dcim/interfaces": Interfaces,
             "dcim/rear-ports": RearPorts,
         }
         ret = []
-        for (termination_a_data, cable_data, termination_b_data) in req.get():
+        for (termination_a_data, cable_data, termination_b_data) in req:
             this_hop_ret = []
             for hop_item_data in (termination_a_data, cable_data, termination_b_data):
                 # if not fully terminated then some items will be None

pynetbox/models/dcim.py

@@ -1,4 +1,3 @@
 black==19.10b0
-Faker==6.5.*
 pytest==6.2.*
 pytest-docker==0.10.*