Refactoring

Author: jhamon

pinecone/control/db_control.py

@@ -0,0 +1,54 @@
+import logging
+from typing import Optional, TYPE_CHECKING
+
+from pinecone.core.openapi.db_control.api.manage_indexes_api import ManageIndexesApi
+from pinecone.openapi_support.api_client import ApiClient
+
+from pinecone.utils import setup_openapi_client
+from pinecone.core.openapi.db_control import API_VERSION
+
+
+logger = logging.getLogger(__name__)
+""" @private """
+
+if TYPE_CHECKING:
+    from .resources.index import IndexResource
+    from .resources.collection import CollectionResource
+
+
+class DBControl:
+    def __init__(self, config, openapi_config, pool_threads):
+        self.config = config
+        """ @private """
+
+        self.index_api = setup_openapi_client(
+            api_client_klass=ApiClient,
+            api_klass=ManageIndexesApi,
+            config=self.config,
+            openapi_config=self.openapi_config,
+            pool_threads=pool_threads,
+            api_version=API_VERSION,
+        )
+        """ @private """
+
+        self._index_resource: Optional["IndexResource"] = None
+        """ @private """
+
+        self._collection_resource: Optional["CollectionResource"] = None
+        """ @private """
+
+    @property
+    def index(self) -> "IndexResource":
+        if self._index_resource is None:
+            from .resources.index import IndexResource
+
+            self._index_resource = IndexResource(index_api=self.index_api, config=self.config)
+        return self._index_resource
+
+    @property
+    def collection(self) -> "CollectionResource":
+        if self._collection_resource is None:
+            from .resources.collection import CollectionResource
+
+            self._collection_resource = CollectionResource(self.index_api)
+        return self._collection_resource

pinecone/control/db_control_asyncio.py

@@ -0,0 +1,55 @@
+import logging
+from typing import Optional, TYPE_CHECKING
+
+from pinecone.core.openapi.db_control.api.manage_indexes_api import AsyncioManageIndexesApi
+from pinecone.openapi_support import AsyncioApiClient
+
+from pinecone.utils import setup_async_openapi_client
+from pinecone.core.openapi.db_control import API_VERSION
+
+logger = logging.getLogger(__name__)
+""" @private """
+
+
+if TYPE_CHECKING:
+    from .resources_asyncio.index import IndexResourceAsyncio
+    from .resources_asyncio.collection import CollectionResourceAsyncio
+
+
+class DBControlAsyncio:
+    def __init__(self, config, openapi_config, pool_threads):
+        self.config = config
+        """ @private """
+
+        self.index_api = setup_async_openapi_client(
+            api_client_klass=AsyncioApiClient,
+            api_klass=AsyncioManageIndexesApi,
+            config=self.config,
+            openapi_config=self.openapi_config,
+            api_version=API_VERSION,
+        )
+        """ @private """
+
+        self._index_resource: Optional["IndexResourceAsyncio"] = None
+        """ @private """
+
+        self._collection_resource: Optional["CollectionResourceAsyncio"] = None
+        """ @private """
+
+    @property
+    def index(self) -> "IndexResourceAsyncio":
+        if self._index_resource is None:
+            from .resources_asyncio.index import IndexResourceAsyncio
+
+            self._index_resource = IndexResourceAsyncio(
+                index_api=self.index_api, config=self.config
+            )
+        return self._index_resource
+
+    @property
+    def collection(self) -> "CollectionResourceAsyncio":
+        if self._collection_resource is None:
+            from .resources_asyncio.collection import CollectionResourceAsyncio
+
+            self._collection_resource = CollectionResourceAsyncio(self.index_api)
+        return self._collection_resource

pinecone/control/pinecone_interface.py renamed to pinecone/control/legacy_pinecone_interface.py

@@ -1,30 +1,30 @@
 from abc import ABC, abstractmethod
 
-from typing import Optional, Dict, Union
-
-
-from pinecone.models import (
-    ServerlessSpec,
-    PodSpec,
-    IndexList,
-    CollectionList,
-    IndexModel,
-    IndexEmbed,
-)
-from pinecone.enums import (
-    Metric,
-    VectorType,
-    DeletionProtection,
-    PodType,
-    CloudProvider,
-    AwsRegion,
-    GcpRegion,
-    AzureRegion,
-)
-from .types import CreateIndexForModelEmbedTypedDict
-
-
-class PineconeDBControlInterface(ABC):
+from typing import Optional, Dict, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pinecone.models import (
+        ServerlessSpec,
+        PodSpec,
+        IndexList,
+        CollectionList,
+        IndexModel,
+        IndexEmbed,
+    )
+    from pinecone.enums import (
+        Metric,
+        VectorType,
+        DeletionProtection,
+        PodType,
+        CloudProvider,
+        AwsRegion,
+        GcpRegion,
+        AzureRegion,
+    )
+    from .types import CreateIndexForModelEmbedTypedDict
+
+
+class LegacyPineconeDBControlInterface(ABC):
     @abstractmethod
     def __init__(
         self,
@@ -190,14 +190,16 @@ def __init__(
     def create_index(
         self,
         name: str,
-        spec: Union[Dict, ServerlessSpec, PodSpec],
+        spec: Union[Dict, "ServerlessSpec", "PodSpec"],
         dimension: Optional[int],
-        metric: Optional[Union[Metric, str]] = Metric.COSINE,
+        metric: Optional[Union["Metric", str]] = "Metric.COSINE",
         timeout: Optional[int] = None,
-        deletion_protection: Optional[Union[DeletionProtection, str]] = DeletionProtection.DISABLED,
-        vector_type: Optional[Union[VectorType, str]] = VectorType.DENSE,
+        deletion_protection: Optional[
+            Union["DeletionProtection", str]
+        ] = "DeletionProtection.DISABLED",
+        vector_type: Optional[Union["VectorType", str]] = "VectorType.DENSE",
         tags: Optional[Dict[str, str]] = None,
-    ) -> IndexModel:
+    ) -> "IndexModel":
         """Creates a Pinecone index.
 
         :param name: The name of the index to create. Must be unique within your project and
@@ -299,13 +301,15 @@ def create_index(
     def create_index_for_model(
         self,
         name: str,
-        cloud: Union[CloudProvider, str],
-        region: Union[AwsRegion, GcpRegion, AzureRegion, str],
-        embed: Union[IndexEmbed, CreateIndexForModelEmbedTypedDict],
+        cloud: Union["CloudProvider", str],
+        region: Union["AwsRegion", "GcpRegion", "AzureRegion", str],
+        embed: Union["IndexEmbed", "CreateIndexForModelEmbedTypedDict"],
         tags: Optional[Dict[str, str]] = None,
-        deletion_protection: Optional[Union[DeletionProtection, str]] = DeletionProtection.DISABLED,
+        deletion_protection: Optional[
+            Union["DeletionProtection", str]
+        ] = "DeletionProtection.DISABLED",
         timeout: Optional[int] = None,
-    ) -> IndexModel:
+    ) -> "IndexModel":
         """
         :param name: The name of the index to create. Must be unique within your project and
             cannot be changed once created. Allowed characters are lowercase letters, numbers,
@@ -414,7 +418,7 @@ def delete_index(self, name: str, timeout: Optional[int] = None):
         pass
 
     @abstractmethod
-    def list_indexes(self) -> IndexList:
+    def list_indexes(self) -> "IndexList":
         """
         :return: Returns an `IndexList` object, which is iterable and contains a
             list of `IndexModel` objects. The `IndexList` also has a convenience method `names()`
@@ -447,7 +451,7 @@ def list_indexes(self) -> IndexList:
         pass
 
     @abstractmethod
-    def describe_index(self, name: str) -> IndexModel:
+    def describe_index(self, name: str) -> "IndexModel":
         """
         :param name: the name of the index to describe.
         :return: Returns an `IndexModel` object
@@ -534,8 +538,8 @@ def configure_index(
         self,
         name: str,
         replicas: Optional[int] = None,
-        pod_type: Optional[Union[PodType, str]] = None,
-        deletion_protection: Optional[Union[DeletionProtection, str]] = None,
+        pod_type: Optional[Union["PodType", str]] = None,
+        deletion_protection: Optional[Union["DeletionProtection", str]] = None,
         tags: Optional[Dict[str, str]] = None,
     ):
         """
@@ -622,7 +626,7 @@ def configure_index(
         pass
 
     @abstractmethod
-    def create_collection(self, name: str, source: str):
+    def create_collection(self, name: str, source: str) -> None:
         """Create a collection from a pod-based index
 
         :param name: Name of the collection
@@ -631,7 +635,7 @@ def create_collection(self, name: str, source: str):
         pass
 
     @abstractmethod
-    def list_collections(self) -> CollectionList:
+    def list_collections(self) -> "CollectionList":
         """List all collections
 
         ```python