List typing fixes (#287)

* Update ini to pick up more files * Remove runtime type checking for list types * Change approach to specifying list params * Fix after rebase * Formatting * Remove som unused imports and switch typed dict import to typing extensions
workos · Jul 23, 2024 · 3138264 · 3138264
1 parent 2a7ac32
commit 3138264
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 67 deletions.
diff --git a/mypy.ini b/mypy.ini
@@ -1,2 +1,2 @@
 [mypy]
-files=./workos/resources/organizations.py,./workos/resources/directory_sync.py
+files=./workos/**/*/organizations.py,./workos/**/*/directory_sync.py
diff --git a/tests/utils/list_resource.py b/tests/utils/list_resource.py
@@ -1,5 +1,4 @@
-from typing import Dict, List, Optional, TypeVar, TypedDict, Union
-from workos.resources.list import ListPage, WorkOsListResource
+from typing import Dict, List, Optional
 
 
 def list_data_to_dicts(list_data: List):

diff --git a/workos/directory_sync.py b/workos/directory_sync.py
@@ -1,5 +1,4 @@
 from typing import Optional, Protocol
-
 import workos
 from workos.utils.pagination_order import PaginationOrder
 from workos.utils.request import (
@@ -20,6 +19,25 @@
 RESPONSE_LIMIT = 10
 
 
+class DirectoryListFilters(ListArgs, total=False):
+    search: Optional[str]
+    organization_id: Optional[str]
+    domain: Optional[str]
+
+
+class DirectoryUserListFilters(
+    ListArgs,
+    total=False,
+):
+    group: Optional[str]
+    directory: Optional[str]
+
+
+class DirectoryGroupListFilters(ListArgs, total=False):
+    user: Optional[str]
+    directory: Optional[str]
+
+
 class DirectorySyncModule(Protocol):
     def list_users(
         self,
@@ -29,7 +47,7 @@ def list_users(
         before: Optional[str] = None,
         after: Optional[str] = None,
         order: PaginationOrder = "desc",
-    ) -> WorkOsListResource[DirectoryUser]:
+    ) -> WorkOsListResource[DirectoryUser, DirectoryUserListFilters]:
         ...
 
     def list_groups(
@@ -40,7 +58,7 @@ def list_groups(
         before: Optional[str] = None,
         after: Optional[str] = None,
         order: PaginationOrder = "desc",
-    ) -> WorkOsListResource[DirectoryGroup]:
+    ) -> WorkOsListResource[DirectoryGroup, DirectoryGroupListFilters]:
         ...
 
     def get_user(self, user: str) -> DirectoryUser:
@@ -61,7 +79,7 @@ def list_directories(
         after: Optional[str] = None,
         organization: Optional[str] = None,
         order: PaginationOrder = "desc",
-    ) -> WorkOsListResource[Directory]:
+    ) -> WorkOsListResource[Directory, DirectoryListFilters]:
         ...
 
     def delete_directory(self, directory: str) -> None:
@@ -89,7 +107,7 @@ def list_users(
         before: Optional[str] = None,
         after: Optional[str] = None,
         order: PaginationOrder = "desc",
-    ) -> WorkOsListResource[DirectoryUser]:
+    ) -> WorkOsListResource[DirectoryUser, DirectoryUserListFilters]:
         """Gets a list of provisioned Users for a Directory.
 
         Note, either 'directory' or 'group' must be provided.
@@ -106,30 +124,29 @@ def list_users(
             dict: Directory Users response from WorkOS.
         """
 
-        params = {
+        list_params: DirectoryUserListFilters = {
             "limit": limit,
             "before": before,
             "after": after,
             "order": order,
         }
 
         if group is not None:
-            params["group"] = group
+            list_params["group"] = group
         if directory is not None:
-            params["directory"] = directory
+            list_params["directory"] = directory
 
         response = self.request_helper.request(
             "directory_users",
             method=REQUEST_METHOD_GET,
-            params=params,
+            params=list_params,
             token=workos.api_key,
         )
 
         return WorkOsListResource(
             list_method=self.list_users,
-            # TODO: Should we even bother with this validation?
-            list_args=ListArgs.model_validate(params),
-            **ListPage[DirectoryUser](**response).model_dump()
+            list_args=list_params,
+            **ListPage[DirectoryUser](**response).model_dump(),
         )
 
     def list_groups(
@@ -140,7 +157,7 @@ def list_groups(
         before: Optional[str] = None,
         after: Optional[str] = None,
         order: PaginationOrder = "desc",
-    ) -> WorkOsListResource[DirectoryGroup]:
+    ) -> WorkOsListResource[DirectoryGroup, DirectoryGroupListFilters]:
         """Gets a list of provisioned Groups for a Directory .
 
         Note, either 'directory' or 'user' must be provided.
@@ -156,29 +173,29 @@ def list_groups(
         Returns:
             dict: Directory Groups response from WorkOS.
         """
-        params = {
+        list_params: DirectoryGroupListFilters = {
             "limit": limit,
             "before": before,
             "after": after,
             "order": order,
         }
+
         if user is not None:
-            params["user"] = user
+            list_params["user"] = user
         if directory is not None:
-            params["directory"] = directory
+            list_params["directory"] = directory
 
         response = self.request_helper.request(
             "directory_groups",
             method=REQUEST_METHOD_GET,
-            params=params,
+            params=list_params,
             token=workos.api_key,
         )
 
         return WorkOsListResource(
             list_method=self.list_groups,
-            # TODO: Should we even bother with this validation?
-            list_args=ListArgs.model_validate(params),
-            **ListPage[DirectoryGroup](**response).model_dump()
+            list_args=list_params,
+            **ListPage[DirectoryGroup](**response).model_dump(),
         )
 
     def get_user(self, user: str):
@@ -242,7 +259,7 @@ def list_directories(
         after: Optional[str] = None,
         organization: Optional[str] = None,
         order: PaginationOrder = "desc",
-    ) -> WorkOsListResource[Directory]:
+    ) -> WorkOsListResource[Directory, DirectoryListFilters]:
         """Gets details for existing Directories.
 
         Args:
@@ -258,27 +275,26 @@ def list_directories(
             dict: Directories response from WorkOS.
         """
 
-        params = {
-            "domain": domain,
-            "organization": organization,
-            "search": search,
+        list_params: DirectoryListFilters = {
             "limit": limit,
             "before": before,
             "after": after,
             "order": order,
+            "domain": domain,
+            "organization_id": organization,
+            "search": search,
         }
 
         response = self.request_helper.request(
             "directories",
             method=REQUEST_METHOD_GET,
-            params=params,
+            params=list_params,
             token=workos.api_key,
         )
         return WorkOsListResource(
             list_method=self.list_directories,
-            # TODO: Should we even bother with this validation?
-            list_args=ListArgs.model_validate(params),
-            **ListPage[Directory](**response).model_dump()
+            list_args=list_params,
+            **ListPage[Directory](**response).model_dump(),
         )
 
     def delete_directory(self, directory: str):

diff --git a/workos/organizations.py b/workos/organizations.py
@@ -1,5 +1,4 @@
 from typing import List, Optional, Protocol
-
 import workos
 from workos.utils.pagination_order import PaginationOrder
 from workos.utils.request import (
@@ -20,6 +19,10 @@
 RESPONSE_LIMIT = 10
 
 
+class OrganizationListFilters(ListArgs, total=False):
+    domains: Optional[List[str]]
+
+
 class OrganizationsModule(Protocol):
     def list_organizations(
         self,
@@ -28,7 +31,7 @@ def list_organizations(
         before: Optional[str] = None,
         after: Optional[str] = None,
         order: PaginationOrder = "desc",
-    ) -> WorkOsListResource[Organization]:
+    ) -> WorkOsListResource[Organization, OrganizationListFilters]:
         ...
 
     def get_organization(self, organization: str) -> Organization:
@@ -75,7 +78,7 @@ def list_organizations(
         before: Optional[str] = None,
         after: Optional[str] = None,
         order: PaginationOrder = "desc",
-    ) -> WorkOsListResource[Organization]:
+    ) -> WorkOsListResource[Organization, OrganizationListFilters]:
         """Retrieve a list of organizations that have connections configured within your WorkOS dashboard.
 
         Kwargs:
@@ -89,25 +92,24 @@ def list_organizations(
             dict: Organizations response from WorkOS.
         """
 
-        params = {
-            "domains": domains,
+        list_params: OrganizationListFilters = {
             "limit": limit,
             "before": before,
             "after": after,
             "order": order,
+            "domains": domains,
         }
 
         response = self.request_helper.request(
             ORGANIZATIONS_PATH,
             method=REQUEST_METHOD_GET,
-            params=params,
+            params=list_params,
             token=workos.api_key,
         )
 
-        return WorkOsListResource[Organization](
+        return WorkOsListResource[Organization, OrganizationListFilters](
             list_method=self.list_organizations,
-            # TODO: Should we even bother with this validation?
-            list_args=ListArgs.model_validate(params),
+            list_args=list_params,
             **ListPage[Organization](**response).model_dump()
         )
 

diff --git a/workos/resources/list.py b/workos/resources/list.py
@@ -1,25 +1,22 @@
-from abc import abstractmethod
 from typing import (
     List,
-    Any,
     Literal,
     TypeVar,
     Generic,
     Callable,
     Iterator,
     Optional,
 )
-
+from typing_extensions import TypedDict
 from workos.resources.base import WorkOSBaseResource
 from workos.resources.directory_sync import Directory, DirectoryGroup, DirectoryUser
 from workos.resources.organizations import Organization
-
-from pydantic import BaseModel, Extra, Field
-
-# TODO: THIS OLD RESOURCE GOES AWAY
+from pydantic import BaseModel, Field
+from workos.resources.workos_model import WorkOSModel
 
 
 class WorkOSListResource(WorkOSBaseResource):
+    # TODO: THIS OLD RESOURCE GOES AWAY
     """Representation of a WorkOS List Resource as returned through the API.
 
     Attributes:
@@ -124,44 +121,49 @@ class ListMetadata(BaseModel):
     before: Optional[str] = None
 
 
-class ListPage(BaseModel, Generic[ListableResource]):
+class ListPage(WorkOSModel, Generic[ListableResource]):
     object: Literal["list"]
     data: List[ListableResource]
     list_metadata: ListMetadata
 
 
-class ListArgs(BaseModel, extra="allow"):
-    limit: Optional[int] = 10
-    before: Optional[str] = None
-    after: Optional[str] = None
-    order: Literal["asc", "desc"] = "desc"
+class ListArgs(TypedDict):
+    limit: int
+    before: Optional[str]
+    after: Optional[str]
+    order: Literal["asc", "desc"]
+
 
-    class Config:
-        extra = "allow"
+ListAndFilterParams = TypeVar("ListAndFilterParams", bound=ListArgs)
 
 
-class WorkOsListResource(BaseModel, Generic[ListableResource]):
+class WorkOsListResource(
+    WorkOSModel,
+    Generic[ListableResource, ListAndFilterParams],
+):
     object: Literal["list"]
     data: List[ListableResource]
     list_metadata: ListMetadata
 
-    # These fields end up exposed in the types. Does we care?
     list_method: Callable = Field(exclude=True)
-    list_args: ListArgs = Field(exclude=True)
+    list_args: ListAndFilterParams = Field(exclude=True)
 
     def auto_paging_iter(self) -> Iterator[ListableResource]:
-        next_page: WorkOsListResource[ListableResource]
+        next_page: WorkOsListResource[ListableResource, ListAndFilterParams]
 
         after = self.list_metadata.after
-        order = self.list_args.order
-
-        fixed_pagination_params = {"order": order, "limit": self.list_args.limit}
-        filter_params = self.list_args.model_dump(
-            exclude={"after", "before", "order", "limit"}
-        )
 
+        fixed_pagination_params = {
+            "order": self.list_args["order"],
+            "limit": self.list_args["limit"],
+        }
+        # Omit common list parameters
+        filter_params = {
+            k: v
+            for k, v in self.list_args.items()
+            if k not in {"order", "limit", "before", "after"}
+        }
         index: int = 0
-
         while True:
             if index >= len(self.data):
                 if after is not None: