Add static type checking with mypy (pinecone-io#240)

## Problem

Some parts of this code were including typehints, but we haven't been
running any checks to verify the correctness of those annotations.
Static typechecking helps to improve overall quality and catch mistakes.

## Solution

- Add a dev dependency on `mypy` ([site](https://mypy-lang.org/)) and
packages with types for third-party dependencies.
- Make adjustments throughout the code to satify the typechecker
- Add steps in CI to run mypy checks automatically on every change.
- For now, exclude generated code in `pinecone/core` because it has a
bunch of issues but is not easy to change.

Summary of code changes:
- Everywhere that has a default value of `None` needs to be wrapped in
the `Optional` type from the `typings` package. Correct: `def
my_method(foo: Optional[str] = None)`. Incorrect: `def my_method(foo:
str = None)`
- Refactored config-related stuff a bit because having config classes
with inner `_config` attributes seemed like a smell and also made typing
config-related stuff more challening. Ended up converting what was
previously the `Config` class into a `ConfigBuilder` and the actual
config object is now just a `NamedTuple` which is a container for
immutable data.
- Mypy doesn't like when you overwrite variables in a way that changes
the type. So some places required me to introduce an additional
intermediate variable to satisfy the type checker.
- Mypy got mad about having a `Pinecone` class in both `pinecone` and
`pinecone.grpc`. So I renamed the GRPC one to `PineconeGRPC`.

## Type of Change

- [x] Infrastructure change (CI configs, etc)

## Test Plan

See new step running in CI.

Author: jhamon

.github/actions/setup-poetry/action.yml

@@ -5,6 +5,7 @@ inputs:
     description: 'Install gRPC dependencies'
     required: true
     default: 'false'
+
 runs:
   using: 'composite'
   steps:

.github/workflows/testing.yaml

@@ -24,6 +24,13 @@ jobs:
 
       - name: Run unit tests
         run: poetry run pytest --cov=pinecone --timeout=120 tests/unit
+      
+      - name: mypy check
+        run: |
+          # Still lots of errors when running on the whole package (especially
+          # in the generated core module), but we can check these subpackages
+          # so we don't add new regressions.
+          poetry run mypy pinecone --exclude pinecone/core --exclude pinecone/grpc
 
   units-grpc:
     name: Run tests (GRPC)
@@ -48,6 +55,14 @@ jobs:
       - name: Run unit tests (GRPC)
         run: poetry run pytest --cov=pinecone --timeout=120 tests/unit_grpc
 
+      - name: mypy check
+        run: |
+          # Still lots of errors when running on the whole package (especially
+          # in the generated core module), but we can check these subpackages
+          # so we don't add new regressions.
+          poetry run mypy pinecone --exclude pinecone/core
+
+
   package:
     name: Check packaging
     runs-on: ubuntu-latest

Makefile

@@ -17,6 +17,9 @@ test-grpc-unit:
 	@echo "Running tests..."
 	poetry run pytest --cov=pinecone --timeout=120 tests/unit_grpc
 
+make type-check:
+	poetry run mypy pinecone --exclude pinecone/core
+
 version:
 	poetry version
 

pinecone/config/__init__.py

@@ -1,2 +1,2 @@
-from .config import Config
+from .config import ConfigBuilder, Config
 from .pinecone_config import PineconeConfig

pinecone/config/config.py

@@ -1,20 +1,19 @@
-from typing import NamedTuple
+from typing import NamedTuple, Optional
 import os
 
-from pinecone.utils import check_kwargs
 from pinecone.exceptions import PineconeConfigurationError
 from pinecone.core.client.exceptions import ApiKeyError
 from pinecone.config.openapi import OpenApiConfigFactory
 from pinecone.core.client.configuration import Configuration as OpenApiConfiguration
 
 
-class ConfigBase(NamedTuple):
+class Config(NamedTuple):
     api_key: str = ""
     host: str = ""
-    openapi_config: OpenApiConfiguration = None
+    openapi_config: Optional[OpenApiConfiguration] = None
 
 
-class Config:
+class ConfigBuilder:
     """
 
     Configurations are resolved in the following order:
@@ -31,39 +30,25 @@ class Config:
     :param openapi_config: Optional. Set OpenAPI client configuration.
     """
 
-    def __init__(
-        self,
-        api_key: str = None,
-        host: str = None,
-        openapi_config: OpenApiConfiguration = None,
+    @staticmethod
+    def build(
+        api_key: Optional[str] = None,
+        host: Optional[str] = None,
+        openapi_config: Optional[OpenApiConfiguration] = None,
         **kwargs,
-    ):
+    ) -> Config:
         api_key = api_key or kwargs.pop("api_key", None) or os.getenv("PINECONE_API_KEY")
         host = host or kwargs.pop("host", None)
+
+        if not api_key:
+            raise PineconeConfigurationError("You haven't specified an Api-Key.")
+        if not host:
+            raise PineconeConfigurationError("You haven't specified a host.")
+
         openapi_config = (
             openapi_config
             or kwargs.pop("openapi_config", None)
             or OpenApiConfigFactory.build(api_key=api_key, host=host)
         )
 
-        check_kwargs(self.__init__, kwargs)
-        self._config = ConfigBase(api_key, host, openapi_config)
-        self.validate()
-
-    def validate(self):
-        if not self._config.api_key:
-            raise PineconeConfigurationError("You haven't specified an Api-Key.")
-        if not self._config.host:
-            raise PineconeConfigurationError("You haven't specified a host.")
-
-    @property
-    def API_KEY(self):
-        return self._config.api_key
-
-    @property
-    def HOST(self):
-        return self._config.host
-
-    @property
-    def OPENAPI_CONFIG(self):
-        return self._config.openapi_config
+        return Config(api_key, host, openapi_config)

pinecone/config/openapi.py

@@ -1,5 +1,5 @@
 import sys
-from typing import List
+from typing import List, Optional
 
 import certifi
 import requests
@@ -16,7 +16,7 @@
 
 class OpenApiConfigFactory:
     @classmethod
-    def build(cls, api_key: str, host: str = None, **kwargs):
+    def build(cls, api_key: str, host: Optional[str] = None, **kwargs):
         openapi_config = OpenApiConfiguration()
         openapi_config.host = host
         openapi_config.ssl_ca_cert = certifi.where()
@@ -26,6 +26,7 @@ def build(cls, api_key: str, host: str = None, **kwargs):
 
     @classmethod
     def _get_socket_options(
+        self,
         do_keep_alive: bool = True,
         keep_alive_idle_sec: int = TCP_KEEPIDLE,
         keep_alive_interval_sec: int = TCP_KEEPINTVL,

pinecone/config/pinecone_config.py

@@ -1,10 +1,12 @@
+from typing import Optional
 import os
-from .config import Config
+from .config import ConfigBuilder, Config
 
 DEFAULT_CONTROLLER_HOST = "https://api.pinecone.io"
 
 
-class PineconeConfig(Config):
-    def __init__(self, api_key: str = None, host: str = None, **kwargs):
+class PineconeConfig():
+    @staticmethod
+    def build(api_key: Optional[str] = None, host: Optional[str] = None, **kwargs) -> Config:
         host = host or kwargs.get("host") or os.getenv("PINECONE_CONTROLLER_HOST") or DEFAULT_CONTROLLER_HOST
-        super().__init__(api_key=api_key, host=host, **kwargs)
+        return ConfigBuilder.build(api_key=api_key, host=host, **kwargs)

pinecone/control/index_host_store.py

@@ -1,9 +1,10 @@
+from typing import Dict
 from pinecone.config import Config
 from pinecone.core.client.api.index_operations_api import IndexOperationsApi
 
 
 class SingletonMeta(type):
-    _instances = {}
+    _instances: Dict[str, str] = {}
 
     def __call__(cls, *args, **kwargs):
         if cls not in cls._instances:
@@ -17,7 +18,7 @@ def __init__(self):
         self._indexHosts = {}
 
     def _key(self, config: Config, index_name: str) -> str:
-        return ":".join([config.API_KEY, index_name])
+        return ":".join([config.api_key, index_name])
 
     def delete_host(self, config: Config, index_name: str):
         key = self._key(config, index_name)

pinecone/control/pinecone.py

@@ -1,5 +1,5 @@
 import time
-from typing import Optional
+from typing import Optional, Dict, Any, cast
 
 from .index_host_store import IndexHostStore
 
@@ -18,21 +18,25 @@
 class Pinecone:
     def __init__(
         self,
-        api_key: str = None,
-        host: str = None,
-        config: Config = None,
-        index_api: IndexOperationsApi = None,
+        api_key: Optional[str] = None,
+        host: Optional[str] = None,
+        config: Optional[Config] = None,
+        index_api: Optional[IndexOperationsApi] = None,
         **kwargs,
     ):
         if config or kwargs.get("config"):
-            self.config = config or kwargs.get("config")
+            configKwarg = config or kwargs.get("config")
+            if not isinstance(configKwarg, Config):
+                raise TypeError("config must be of type pinecone.config.Config")
+            else:
+                self.config = configKwarg
         else:
-            self.config = PineconeConfig(api_key=api_key, host=host, **kwargs)
-
+            self.config = PineconeConfig.build(api_key=api_key, host=host, **kwargs)
+        
         if index_api:
             self.index_api = index_api
         else:
-            api_client = ApiClient(configuration=self.config.OPENAPI_CONFIG)
+            api_client = ApiClient(configuration=self.config.openapi_config)
             api_client.user_agent = get_user_agent()
             self.index_api = IndexOperationsApi(api_client)
 
@@ -45,15 +49,15 @@ def create_index(
         cloud: str,
         region: str,
         capacity_mode: str,
-        timeout: int = None,
+        timeout: Optional[int] = None,
         index_type: str = "approximated",
         metric: str = "cosine",
         replicas: int = 1,
         shards: int = 1,
         pods: int = 1,
         pod_type: str = "p1",
-        index_config: dict = None,
-        metadata_config: dict = None,
+        index_config: Optional[dict] = None,
+        metadata_config: Optional[dict] = None,
         source_collection: str = "",
     ):
         """Creates a Pinecone index.
@@ -138,7 +142,7 @@ def is_ready():
                 )
             )
 
-    def delete_index(self, name: str, timeout: int = None):
+    def delete_index(self, name: str, timeout: Optional[int] = None):
         """Deletes a Pinecone index.
 
         :param name: the name of the index.
@@ -210,7 +214,7 @@ def configure_index(self, name: str, replicas: Optional[int] = None, pod_type: O
         :param: pod_type: the new pod_type for the index.
         """
         api_instance = self.index_api
-        config_args = {}
+        config_args: Dict[str, Any] = {}
         if pod_type != "":
             config_args.update(pod_type=pod_type)
         if replicas:
@@ -267,4 +271,4 @@ def _get_status(self, name: str):
 
     def Index(self, name: str):
         index_host = self.index_host_store.get_host(self.index_api, self.config, name)
-        return Index(api_key=self.config.API_KEY, host=index_host)
+        return Index(api_key=self.config.api_key, host=index_host)

pinecone/data/index.py

@@ -3,7 +3,7 @@
 from collections.abc import Iterable
 from typing import Union, List, Tuple, Optional, Dict, Any
 
-from pinecone.config import Config
+from pinecone.config import ConfigBuilder
 
 from pinecone.core.client.models import SparseValues
 from pinecone.core.client import ApiClient
@@ -76,13 +76,13 @@ class Index():
     """
 
     def __init__(self, api_key: str, host: str, pool_threads=1, **kwargs):
-        api_key = api_key or kwargs.get("api_key")
-        host = host or kwargs.get('host')
+        api_key = api_key or kwargs.get("api_key", None)
+        host = host or kwargs.get('host', None)
         pool_threads = pool_threads or kwargs.get("pool_threads")
 
-        self._config = Config(api_key=api_key, host=host, **kwargs)
+        self._config = ConfigBuilder.build(api_key=api_key, host=host, **kwargs)
 
-        api_client = ApiClient(configuration=self._config.OPENAPI_CONFIG, pool_threads=pool_threads)
+        api_client = ApiClient(configuration=self._config.openapi_config, pool_threads=pool_threads)
         api_client.user_agent = get_user_agent()
         self._api_client = api_client
         self._vector_api = VectorOperationsApi(api_client=api_client)
@@ -180,7 +180,7 @@ def upsert(
         return UpsertResponse(upserted_count=total_upserted)
 
     def _upsert_batch(
-        self, vectors: List[Vector], namespace: Optional[str], _check_type: bool, **kwargs
+        self, vectors: Union[List[Vector], List[tuple], List[dict]], namespace: Optional[str], _check_type: bool, **kwargs
     ) -> UpsertResponse:
         args_dict = self._parse_non_empty_args([("namespace", namespace)])
         vec_builder = lambda v: VectorFactory.build(v, check_type=_check_type)
@@ -202,7 +202,7 @@ def _iter_dataframe(df, batch_size):
             yield batch
 
     def upsert_from_dataframe(
-        self, df, namespace: str = None, batch_size: int = 500, show_progress: bool = True
+        self, df, namespace: Optional[str] = None, batch_size: int = 500, show_progress: bool = True
     ) -> UpsertResponse:
         """Upserts a dataframe into the index.
 
@@ -416,7 +416,8 @@ def _query_transform(item):
             ),
             **{k: v for k, v in kwargs.items() if k in _OPENAPI_ENDPOINT_PARAMS},
         )
-        return parse_query_response(response, vector is not None or id)
+        unary_query = True if vector is not None or id else False
+        return parse_query_response(response, unary_query)
 
     @validate_and_convert_errors
     def update(

pinecone/data/vector_factory.py

@@ -91,7 +91,7 @@ def _dict_to_vector(item, check_type: bool) -> Vector:
         try:
             return Vector(**item, _check_type=check_type)
         except TypeError as e:
-            if not isinstance(item["values"], Iterable) or not isinstance(item["values"][0], numbers.Real):
+            if not isinstance(item["values"], Iterable) or not isinstance(item["values"].__iter__().__next__(), numbers.Real):
                 raise TypeError(f"Column `values` is expected to be a list of floats")
             raise e
 

pinecone/grpc/__init__.py

@@ -1,2 +1,2 @@
 from .index_grpc import GRPCIndex
-from .pinecone import Pinecone
+from .pinecone import PineconeGRPC

pinecone/grpc/base.py

@@ -1,11 +1,11 @@
 import logging
 from abc import ABC, abstractmethod
 from functools import wraps
-from typing import Dict
+from typing import Dict, Optional
 
 import certifi
 import grpc
-from grpc._channel import _InactiveRpcError
+from grpc._channel import _InactiveRpcError, Channel
 import json
 
 from .retry import RetryConfig
@@ -29,15 +29,15 @@ def __init__(
         self,
         index_name: str,
         config: Config,
-        channel=None,
-        grpc_config: GRPCClientConfig = None,
-        _endpoint_override: str = None,
+        channel: Optional[Channel] =None,
+        grpc_config: Optional[GRPCClientConfig] = None,
+        _endpoint_override: Optional[str] = None,
     ):
         self.name = index_name
 
         self.grpc_client_config = grpc_config or GRPCClientConfig()
         self.retry_config = self.grpc_client_config.retry_config or RetryConfig()
-        self.fixed_metadata = {"api-key": config.API_KEY, "service-name": index_name, "client-version": CLIENT_VERSION}
+        self.fixed_metadata = {"api-key": config.api_key, "service-name": index_name, "client-version": CLIENT_VERSION}
         self._endpoint_override = _endpoint_override
 
         self.method_config = json.dumps(

pinecone/grpc/config.py

@@ -25,7 +25,7 @@ class GRPCClientConfig(NamedTuple):
     conn_timeout: int = 1
     reuse_channel: bool = True
     retry_config: Optional[RetryConfig] = None
-    grpc_channel_options: Dict[str, str] = None
+    grpc_channel_options: Optional[Dict[str, str]] = None
 
     @classmethod
     def _from_dict(cls, kwargs: dict):