finalize list() and extend PluginAware in NamespaceResource

Author: rohanshah18

pinecone/db_data/index.py

@@ -1,3 +1,4 @@
+from pinecone.core.openapi.db_data.model.list_namespaces_response import ListNamespacesResponse
 from pinecone.utils.tqdm import tqdm
 import warnings
 import logging
@@ -8,7 +9,6 @@
 
 from pinecone.openapi_support import ApiClient
 from pinecone.core.openapi.db_data.api.vector_operations_api import VectorOperationsApi
-from pinecone.core.openapi.db_data.api.namespace_operations_api import NamespaceOperationsApi
 from pinecone.core.openapi.db_data import API_VERSION
 from pinecone.core.openapi.db_data.models import (
     QueryResponse,
@@ -630,10 +630,11 @@ def delete_namespace(self, namespace: str) -> Dict[str, Any]:
         return self.namespace.delete(namespace=namespace)
 
     @validate_and_convert_errors
-    def list_namespaces(self,
-        limit: Optional[int] = None,
-        pagination_token: Optional[str] = None,
-    ) -> Iterator["NamespaceDescription"]:
-        return self.namespace.list(limit=limit, pagination_token=pagination_token)
+    def list_namespaces(self, **kwargs) -> Iterator[ListNamespacesResponse]:
+        return self.namespace.list(**kwargs)
 
-    
+    @validate_and_convert_errors
+    def list_namespaces_paginated(
+        self, limit: Optional[int] = None, pagination_token: Optional[str] = None
+    ) -> ListNamespacesResponse:
+        return self.namespace.list_paginated(limit=limit, pagination_token=pagination_token)

pinecone/db_data/index_asyncio.py

@@ -669,27 +669,19 @@ async def cancel_import(self, id: str):
         return await self.bulk_import.cancel(id=id)
 
     @validate_and_convert_errors
-    async def describe_namespace(
-        self,
-        namespace: str,
-        **kwargs
-    ) -> "NamespaceDescription":
+    async def describe_namespace(self, namespace: str) -> "NamespaceDescription":
         return await self.namespace.describe(namespace=namespace)
 
     @validate_and_convert_errors
-    async def delete_namespace(
-        self,
-        namespace: str,
-        **kwargs
-    ) -> Dict[str, Any]:
+    async def delete_namespace(self, namespace: str) -> Dict[str, Any]:
         return await self.namespace.delete(namespace=namespace)
 
     @validate_and_convert_errors
-    async def list_namespaces(
-        self,
-        limit: Optional[int] = None,
-        pagination_token: Optional[str] = None,
-        **kwargs
-    ) -> AsyncIterator["NamespaceDescription"]:
-        async for namespace in self.namespace.list(limit=limit, pagination_token=pagination_token, **kwargs):
-            yield namespace
+    async def list_namespaces(self, **kwargs) -> AsyncIterator[ListNamespacesResponse]:
+        return await self.namespace.list(**kwargs)
+
+    @validate_and_convert_errors
+    async def list_namespaces_paginated(
+        self, limit: Optional[int] = None, pagination_token: Optional[str] = None
+    ) -> ListNamespacesResponse:
+        return await self.namespace.list_paginated(limit=limit, pagination_token=pagination_token)

pinecone/db_data/index_asyncio_interface.py

@@ -11,6 +11,7 @@
     SparseValues,
     SearchRecordsResponse,
     NamespaceDescription,
+    ListNamespacesResponse,
 )
 from .query_results_aggregator import QueryNamespacesResults
 from .types import (
@@ -810,14 +811,12 @@ async def search_records(
     @abstractmethod
     async def describe_namespace(
         self,
-        namespace: str,
-        **kwargs
+        namespace: str
     ) -> NamespaceDescription:
         """Describe a namespace within an index, showing the vector count within the namespace.
 
         Args:
             namespace (str): The namespace to describe
-            **kwargs: Additional arguments to pass to the API call
 
         Returns:
             NamespaceDescription: Information about the namespace including vector count
@@ -827,14 +826,12 @@ async def describe_namespace(
     @abstractmethod
     async def delete_namespace(
         self,
-        namespace: str,
-        **kwargs
+        namespace: str
     ) -> Dict[str, Any]:
         """Delete a namespace from an index.
 
         Args:
             namespace (str): The namespace to delete
-            **kwargs: Additional arguments to pass to the API call
 
         Returns:
             Dict[str, Any]: Response from the delete operation
@@ -843,19 +840,44 @@ async def delete_namespace(
 
     @abstractmethod
     async def list_namespaces(
-        self,
-        limit: Optional[int] = None,
-        pagination_token: Optional[str] = None,
-        **kwargs
-    ) -> AsyncIterator[NamespaceDescription]:
-        """Get a list of all namespaces within an index.
+        self, **kwargs
+    ) -> AsyncIterator[ListNamespacesResponse]:
+        """List all namespaces in an index. This method automatically handles pagination to return all results.
 
         Args:
-            limit (Optional[int]): Max number of namespaces to return per page
-            pagination_token (Optional[str]): Token to continue a previous listing operation
-            **kwargs: Additional arguments to pass to the API call
+            limit (Optional[int]): The maximum number of namespaces to return. If unspecified, the server will use a default value. [optional]
 
         Returns:
-            AsyncIterator[NamespaceDescription]: Async generator yielding namespace descriptions
+            `ListNamespacesResponse`: Object containing the list of namespaces.
+
+        Examples:
+            >>> async for namespace in index.list_namespaces(limit=5):
+            ...     print(f"Namespace: {namespace.name}, Vector count: {namespace.vector_count}")
+            Namespace: namespace1, Vector count: 1000
+            Namespace: namespace2, Vector count: 2000
         """
         pass
+
+    @abstractmethod
+    async def list_namespaces_paginated(
+        self, limit: Optional[int] = None, pagination_token: Optional[str] = None
+    ) -> ListNamespacesResponse:
+        """List all namespaces in an index with pagination support. The response includes pagination information if there are more results available.
+
+        Consider using the `list_namespaces` method to avoid having to handle pagination tokens manually.
+
+        Args:
+            limit (Optional[int]): The maximum number of namespaces to return. If unspecified, the server will use a default value. [optional]
+            pagination_token (Optional[str]): A token needed to fetch the next page of results. This token is returned
+                in the response if additional results are available. [optional]
+
+        Returns:
+            `ListNamespacesResponse`: Object containing the list of namespaces and pagination information.
+
+        Examples:
+            >>> results = await index.list_namespaces_paginated(limit=5)
+            >>> results.pagination.next
+            eyJza2lwX3Bhc3QiOiI5OTMiLCJwcmVmaXgiOiI5OSJ9
+            >>> next_results = await index.list_namespaces_paginated(limit=5, pagination_token=results.pagination.next)
+        """
+        pass

pinecone/db_data/interfaces.py

@@ -11,6 +11,7 @@
     SparseValues,
     SearchRecordsResponse,
     NamespaceDescription,
+    ListNamespacesResponse,
 )
 from .query_results_aggregator import QueryNamespacesResults
 from multiprocessing.pool import ApplyResult
@@ -682,9 +683,7 @@ def list(self, **kwargs):
         pass
 
     @abstractmethod
-    def describe_namespace(
-            self, namespace: str, **kwargs
-    ) -> NamespaceDescription:
+    def describe_namespace(self, namespace: str) -> NamespaceDescription:
         """Describe a namespace within an index, showing the vector count within the namespace.
 
         Args:
@@ -697,9 +696,7 @@ def describe_namespace(
         pass
 
     @abstractmethod
-    def delete_namespace(
-        self, namespace: str, **kwargs
-    ) -> Dict[str, Any]:
+    def delete_namespace(self, namespace: str) -> Dict[str, Any]:
         """Delete a namespace from an index.
 
         Args:
@@ -712,20 +709,44 @@ def delete_namespace(
         pass
 
     @abstractmethod
-    def list_namespaces(
-        self,
-        limit: Optional[int] = None,
-        pagination_token: Optional[str] = None,
-        **kwargs
-    ) -> Iterator[NamespaceDescription]:
-        """Get a list of all namespaces within an index.
+    def list_namespaces(self, **kwargs) -> Iterator[ListNamespacesResponse]:
+        """List all namespaces in an index. This method automatically handles pagination to return all results.
 
         Args:
-            limit (Optional[int]): Max number of namespaces to return per page
-            pagination_token (Optional[str]): Token to continue a previous listing operation
-            **kwargs: Additional arguments to pass to the API call
+            limit (Optional[int]): The maximum number of namespaces to return. If unspecified, the server will use a default value. [optional]
 
         Returns:
-            Iterator[NamespaceDescription]: Generator yielding namespace descriptions
+            `ListNamespacesResponse`: Object containing the list of namespaces.
+
+        Examples:
+            >>> results = list(index.list_namespaces(limit=5))
+            >>> for namespace in results:
+            ...     print(f"Namespace: {namespace.name}, Vector count: {namespace.vector_count}")
+            Namespace: namespace1, Vector count: 1000
+            Namespace: namespace2, Vector count: 2000
         """
         pass
+
+    @abstractmethod
+    def list_namespaces_paginated(
+        self, limit: Optional[int] = None, pagination_token: Optional[str] = None
+    ) -> ListNamespacesResponse:
+        """List all namespaces in an index with pagination support. The response includes pagination information if there are more results available.
+
+        Consider using the `list_namespaces` method to avoid having to handle pagination tokens manually.
+
+        Args:
+            limit (Optional[int]): The maximum number of namespaces to return. If unspecified, the server will use a default value. [optional]
+            pagination_token (Optional[str]): A token needed to fetch the next page of results. This token is returned
+                in the response if additional results are available. [optional]
+
+        Returns:
+            `ListNamespacesResponse`: Object containing the list of namespaces and pagination information.
+
+        Examples:
+            >>> results = index.list_namespaces_paginated(limit=5)
+            >>> results.pagination.next
+            eyJza2lwX3Bhc3QiOiI5OTMiLCJwcmVmaXgiOiI5OSJ9
+            >>> next_results = index.list_namespaces_paginated(limit=5, pagination_token=results.pagination.next)
+        """
+        pass

pinecone/db_data/resources/asyncio/namespace_asyncio.py

@@ -15,7 +15,7 @@
 
 
 class NamespaceResourceAsyncio:
-    def __init__(self, api_client, **kwargs) -> None:
+    def __init__(self, api_client) -> None:
         self.__namespace_operations_api = AsyncioNamespaceOperationsApi(api_client)
 
     async def describe(self, namespace: str) -> NamespaceDescription:
@@ -41,7 +41,7 @@ async def delete(self, namespace: str):
         args = NamespaceRequestFactory.delete_namespace_args(namespace=namespace)
         return await self.__namespace_operations_api.delete_namespace(**args)
 
-    async def list(self, **kwargs) -> AsyncIterator[NamespaceDescription]:
+    async def list(self, **kwargs) -> AsyncIterator[ListNamespacesResponse]:
         """
         Args:
             limit (Optional[int]): The maximum number of namespaces to fetch in each network call. If unspecified, the server will use a default value. [optional]
@@ -50,22 +50,55 @@ async def list(self, **kwargs) -> AsyncIterator[NamespaceDescription]:
 
         Returns:
             Returns an async generator that yields each namespace. It automatically handles pagination tokens on your behalf so you can
-            easily iterate over all results.
+            easily iterate over all results. The `list` method accepts all of the same arguments as list_paginated
 
         ```python
         async for namespace in index.list_namespaces():
             print(namespace)
         ```
+
+        You can convert the generator into a list by using an async list comprehension:
+
+        ```python
+        namespaces = [namespace async for namespace in index.list_namespaces()]
+        ```
+
+        You should be cautious with this approach because it will fetch all namespaces at once, which could be a large number
+        of network calls and a lot of memory to hold the results.
         """
         done = False
         while not done:
-            args_dict = NamespaceRequestFactory.list_namespaces_args(**kwargs)
-            results = await self.__namespace_operations_api.list_namespaces_operation(**args_dict)
-            if len(results.namespaces) > 0:
+            results = await self.list_paginated(**kwargs)
+            if results.namespaces is not None and len(results.namespaces) > 0:
                 for namespace in results.namespaces:
                     yield namespace
 
-            if results.pagination:
+            if results.pagination and results.pagination.next:
                 kwargs.update({"pagination_token": results.pagination.next})
             else:
-                done = True 
+                done = True
+
+    async def list_paginated(
+        self, limit: Optional[int] = None, pagination_token: Optional[str] = None
+    ) -> ListNamespacesResponse:
+        """
+        Args:
+            limit (Optional[int]): The maximum number of namespaces to return. If unspecified, the server will use a default value. [optional]
+            pagination_token (Optional[str]): A token needed to fetch the next page of results. This token is returned
+                in the response if additional results are available. [optional]
+
+        Returns:
+            `ListNamespacesResponse`: Object containing the list of namespaces and pagination information.
+
+        List all namespaces in an index with pagination support. The response includes pagination information if there are more results available.
+
+        Consider using the `list` method to avoid having to handle pagination tokens manually.
+
+        Examples:
+            >>> results = await index.list_paginated(limit=5)
+            >>> results.pagination.next
+            eyJza2lwX3Bhc3QiOiI5OTMiLCJwcmVmaXgiOiI5OSJ9
+            >>> next_results = await index.list_paginated(limit=5, pagination_token=results.pagination.next)
+        """
+        args = NamespaceRequestFactory.list_namespaces_args(limit=limit, pagination_token=pagination_token)
+        return await self.__namespace_operations_api.list_namespaces_operation(**args)