cleanup tying, pagination to list

Author: arthanson

netbox_python/baseapi.py

@@ -1,54 +1,69 @@
+from typing import Any, Callable, Dict, List
+
+from .rest import Result
+
+
 class baseapi:
     def __init__(self, client):
         self.client = client
         super().__init__()
 
 
 class CreateableAPIResource:
-    def create(self, *args, **kwargs):
+    def create(self, *args, **kwargs) -> Result:
         return self.client.post(self.path, json=args[0] if args else kwargs)
 
 
 class DeletableAPIResource:
-    def delete(self, obj):
-        if obj and isinstance(obj, list):
+    def delete(self, obj: List[Any] | str | int) -> Result:
+        if isinstance(obj, list):
             return self.client.delete(self.path, json=obj)
+
         return self.client.delete(f"{self.path}{obj}/")
 
 
 class ListableAPIResource:
-    def list(self, **kwargs):
+    def paginate(self, result: Result) -> Result:
+        next_token = result.pagination["next"]
+        yield result
+        while next_token:
+            result = self.client.get(
+                self.path, url_override=next_token, params=result.params
+            )
+            yield result
+            next_token = result.pagination["next"]
+
+    def list(self, **kwargs) -> Result:
         return self.client.get(self.path, params=kwargs)
 
     def all(self, **kwargs):
-        ret = self.client.get(self.path, params=kwargs)
-        pagination = ret.pagination
-        while pagination["next"]:
-            partial = self.client.get(
-                self.path, url_override=pagination["next"], params=kwargs
-            )
-            ret.data.append(partial.data)
-            pagination = partial.pagination
+        result = None
+        for page in self.paginate(self.client.get(self.path, params=kwargs)):
+            if not result:
+                result = page
+            else:
+                result.data.append(page.data)
 
-        ret.pagination["next"] = None
-        ret.pagination["previous"] = None
-        return ret
+        result.pagination["next"] = None
+        result.pagination["previous"] = None
+        return result
 
 
 class RetrievableAPIResource:
-    def get(self, id):
+    def get(self, id: str | int) -> Result:
         return self.client.get(f"{self.path}{id}/")
 
 
 class RetrievableRootAPIResource:
-    def get(self, **kwargs):
+    def get(self, **kwargs) -> Result:
         return self.client.get(self.path, params=kwargs)
 
 
 class UpdateableAPIResource:
-    def update(self, obj, **kwargs):
-        if obj and isinstance(obj, list):
+    def update(self, obj: List[Any] | str | int, **kwargs) -> Result:
+        if isinstance(obj, list):
             return self.client.patch(self.path, json=obj)
+
         return self.client.patch(f"{self.path}{obj}/", json=kwargs)
 
 

netbox_python/rest.py

@@ -7,7 +7,7 @@
 
 from netbox_python.exceptions import NetBoxException
 
-JSONType = Union[None, bool, int, float, str, List[Any], Dict[str, Any]]
+# JSONType = Union[None, bool, int, float, str, List[Any], Dict[str, Any]]
 
 
 class Result:
@@ -17,6 +17,7 @@ def __init__(
         headers: CaseInsensitiveDict,
         message: str = "",
         pagination: dict = None,
+        params: dict = None,
         data: List[Dict] = None,
     ):
         """
@@ -30,6 +31,7 @@ def __init__(
         self.message = str(message)
         self.data = data if data else []
         self.pagination = pagination
+        self.params = params
 
 
 class RestClient:
@@ -49,7 +51,7 @@ def __exit__(self, *_):
     def close(self):
         return self._session.close()
 
-    def request(self, method: str, path: str, **kwargs) -> JSONType:
+    def request(self, method: str, path: str, **kwargs) -> Result:
         url = kwargs.pop("url_override", None)
         if not url:
             url = f"{self.base_url}/{path}"
@@ -94,6 +96,7 @@ def request(self, method: str, path: str, **kwargs) -> JSONType:
                 headers=response.headers,
                 message=response.reason,
                 pagination=pagination,
+                params=kwargs.get("params", None),
                 data=data_out,
             )
         # self._logger.error(msg=log_line)