Skip to content

Commit

Permalink
Fixes for mypy
Browse files Browse the repository at this point in the history
  • Loading branch information
jhamon committed Oct 31, 2024
1 parent a75805d commit 54c2fa3
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 9 deletions.
22 changes: 14 additions & 8 deletions pinecone/grpc/future.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -29,19 +32,19 @@ 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)
elif grpc_future.running():
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)

Expand All @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion pinecone/grpc/index_grpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit 54c2fa3

Please sign in to comment.