From 54c2fa33d4aa98bc73007ae84928945425688eef Mon Sep 17 00:00:00 2001
From: Jen Hamon <jhamon@pinecone.io>
Date: Thu, 31 Oct 2024 09:44:51 -0400
Subject: [PATCH] Fixes for mypy

---
 pinecone/grpc/future.py     | 22 ++++++++++++++--------
 pinecone/grpc/index_grpc.py |  2 +-
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/pinecone/grpc/future.py b/pinecone/grpc/future.py
index fa8c4d71..40a1b28a 100644
--- a/pinecone/grpc/future.py
+++ b/pinecone/grpc/future.py
@@ -6,12 +6,15 @@
 
 class PineconeGrpcFuture(ConcurrentFuture):
     def __init__(
-        self, grpc_future: GrpcFuture, timeout: Optional[int] = 10, result_transformer=None
+        self, grpc_future: GrpcFuture, timeout: Optional[int] = None, result_transformer=None
     ):
         super().__init__()
         self._grpc_future = grpc_future
-        self.default_timeout = timeout  # seconds
-        self.result_transformer = result_transformer
+        self._result_transformer = result_transformer
+        if timeout is not None:
+            self._default_timeout = timeout  # seconds
+        else:
+            self._default_timeout = 5  # seconds
 
         # Sync initial state, in case the gRPC future is already done
         self._sync_state(self._grpc_future)
@@ -29,11 +32,11 @@ def _sync_state(self, grpc_future):
 
         if grpc_future.cancelled():
             self.cancel()
-        elif grpc_future.exception(timeout=self.default_timeout):
+        elif grpc_future.exception(timeout=self._default_timeout):
             self.set_exception(grpc_future.exception())
         elif grpc_future.done():
             try:
-                result = grpc_future.result(timeout=self.default_timeout)
+                result = grpc_future.result(timeout=self._default_timeout)
                 self.set_result(result)
             except Exception as e:
                 self.set_exception(e)
@@ -41,7 +44,7 @@ def _sync_state(self, grpc_future):
             self.set_running_or_notify_cancel()
 
     def set_result(self, result):
-        if self.result_transformer:
+        if self._result_transformer:
             result = self.result_transformer(result)
         return super().set_result(result)
 
@@ -66,8 +69,11 @@ def result(self, timeout=None):
         except RpcError as e:
             raise self._wrap_rpc_exception(e) from e
 
-    def _timeout(self, timeout: Optional[int]) -> int:
-        return timeout if timeout is not None else self.default_timeout
+    def _timeout(self, timeout: Optional[int] = None) -> int:
+        if timeout is not None:
+            return timeout
+        else:
+            return self._default_timeout
 
     def _wrap_rpc_exception(self, e):
         if e._state and e._state.debug_error_string:
diff --git a/pinecone/grpc/index_grpc.py b/pinecone/grpc/index_grpc.py
index eff8fafe..317a0fc5 100644
--- a/pinecone/grpc/index_grpc.py
+++ b/pinecone/grpc/index_grpc.py
@@ -287,7 +287,7 @@ def fetch(
         namespace: Optional[str] = None,
         async_req: Optional[bool] = False,
         **kwargs,
-    ) -> FetchResponse:
+    ) -> Union[FetchResponse, PineconeGrpcFuture]:
         """
         The fetch operation looks up and returns vectors, by ID, from a single namespace.
         The returned vectors include the vector data and/or metadata.