Skip to content

Commit

Permalink
Merge pull request #358 from jqueuniet/add_next_method
Browse files Browse the repository at this point in the history
Add a __next__ method to complete the iterator protocol
  • Loading branch information
Zach Moody authored Apr 7, 2021
2 parents 9f3f4a9 + 90ee2f3 commit 30b813c
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 4 deletions.
4 changes: 2 additions & 2 deletions pynetbox/core/endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@ def get(self, *args, **kwargs):
key = None

if not key:
return next(iter(self.filter(**kwargs)), None)
return next(self.filter(**kwargs), None)

req = Request(
key=key,
Expand All @@ -143,7 +143,7 @@ def get(self, *args, **kwargs):
http_session=self.api.http_session,
)
try:
return next(iter(RecordSet(self, req)), None)
return next(RecordSet(self, req), None)
except RequestError as e:
if e.req.status_code == 404:
return None
Expand Down
10 changes: 8 additions & 2 deletions pynetbox/core/response.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,11 +78,17 @@ def __init__(self, endpoint, request, **kwargs):
self._response_cache = []

def __iter__(self):
for i in self._response_cache:
yield self.endpoint.return_obj(i, self.endpoint.api, self.endpoint)
if self._response_cache:
yield self.endpoint.return_obj(
self._response_cache.pop(), self.endpoint.api, self.endpoint
)
for i in self.response:
yield self.endpoint.return_obj(i, self.endpoint.api, self.endpoint)

def __next__(self):
for i in self:
return i

def __len__(self):
try:
return self.request.count
Expand Down

0 comments on commit 30b813c

Please sign in to comment.