Refactor(policy): convert list method to Iterator pattern

examples/policy.py

@@ -82,7 +82,6 @@ def main():
     _print_header(f"Listing policies in organization: {args.org}")
 
     list_options = PolicyListOptions(
-        page_number=args.page,
         page_size=args.page_size,
     )
 
@@ -93,14 +92,10 @@ def main():
             PolicyKind.SENTINEL if args.kind == "sentinel" else PolicyKind.OPA
         )
 
-    policy_list = client.policies.list(args.org, list_options)
-
-    print(f"Total policies: {policy_list.total_count}")
-    print(f"Page {policy_list.current_page} of {policy_list.total_pages}")
-    print()
+    policy_iter = client.policies.list(args.org, list_options)
 
     existing_policy = None
-    for policy in policy_list.items:
+    for policy in policy_iter:
         print(
             f"- {policy.id} | {policy.name} | kind={policy.kind} | enforcement={policy.enforcement_level}"
         )

src/pytfe/models/policy.py

@@ -38,7 +38,6 @@ class PolicyListOptions(BaseModel):
 
     search: str | None = Field(None, alias="search[name]")
     kind: PolicyKind | None = Field(None, alias="filter[kind]")
-    page_number: int | None = Field(None, alias="page[number]")
     page_size: int | None = Field(None, alias="page[size]")
 
 

src/pytfe/resources/policy.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Iterator
+from typing import Any
+
 from ..errors import (
     InvalidNameError,
     InvalidOrgError,
@@ -11,7 +14,6 @@
 from ..models.policy import (
     Policy,
     PolicyCreateOptions,
-    PolicyList,
     PolicyListOptions,
     PolicyUpdateOptions,
 )
@@ -22,35 +24,28 @@
 class Policies(_Service):
     def list(
         self, organization: str, options: PolicyListOptions | None = None
-    ) -> PolicyList:
-        """List all the policies of the given organization."""
+    ) -> Iterator[Policy]:
+        """Iterate all the policies of the given organization."""
         if not valid_string_id(organization):
             raise InvalidOrgError()
-        params = (
-            options.model_dump(by_alias=True, exclude_none=True) if options else None
-        )
-        r = self.t.request(
-            "GET",
-            f"/api/v2/organizations/{organization}/policies",
-            params=params,
-        )
-        jd = r.json()
-        items = []
-        meta = jd.get("meta", {})
-        pagination = meta.get("pagination", {})
-        for d in jd.get("data", []):
-            attrs = d.get("attributes", {})
-            attrs["id"] = d.get("id")
-            attrs["organization"] = d.get("relationships", {}).get("organization", {})
-            items.append(Policy.model_validate(attrs))
-        return PolicyList(
-            items=items,
-            current_page=pagination.get("current-page"),
-            total_pages=pagination.get("total-pages"),
-            prev_page=pagination.get("prev-page"),
-            next_page=pagination.get("next-page"),
-            total_count=pagination.get("total-count"),
-        )
+
+        path = f"/api/v2/organizations/{organization}/policies"
+        params: dict[str, Any] = {}
+
+        if options:
+            if getattr(options, "page_size", None):
+                params["page[size]"] = str(options.page_size)
+
+        def _gen() -> Iterator[Policy]:
+            for item in self._list(path, params=params):
+                attrs = item.get("attributes", {})
+                attrs["id"] = item.get("id")
+                attrs["organization"] = item.get("relationships", {}).get(
+                    "organization", {}
+                )
+                yield Policy.model_validate(attrs)
+
+        return _gen()
 
     def create(self, organization: str, options: PolicyCreateOptions) -> Policy:
         """Create a new policy in the given organization."""

tests/units/test_policy.py

@@ -14,7 +14,6 @@
     EnforcementLevel,
     Policy,
     PolicyCreateOptions,
-    PolicyList,
     PolicyUpdateOptions,
 )
 from pytfe.models.policy_set import PolicyKind
@@ -84,22 +83,16 @@ def test_list_policies_success_without_options(
         mock_response.json.return_value = mock_response_data
         mock_transport.request.return_value = mock_response
 
-        result = policies_service.list("org-123")
+        result_iter = policies_service.list("org-123")
+        items = list(result_iter)
 
-        mock_transport.request.assert_called_once_with(
-            "GET", "/api/v2/organizations/org-123/policies", params=None
-        )
+        assert mock_transport.request.called
 
-        assert isinstance(result, PolicyList)
-        assert len(result.items) == 1
-        assert result.items[0].id == "pol-123"
-        assert result.items[0].name == "test-policy"
-        assert result.items[0].kind == PolicyKind.SENTINEL
-        assert (
-            result.items[0].enforcement_level == EnforcementLevel.ENFORCEMENT_ADVISORY
-        )
-        assert result.current_page == 1
-        assert result.total_count == 1
+        assert len(items) == 1
+        assert items[0].id == "pol-123"
+        assert items[0].name == "test-policy"
+        assert items[0].kind == PolicyKind.SENTINEL
+        assert items[0].enforcement_level == EnforcementLevel.ENFORCEMENT_ADVISORY
 
     def test_create_policy_validations(self, policies_service):
         """Test create method validations."""