chore: add static type checking with flake8-annotations and mypy (#74)

* add type checking to lint session

* add missing type annotations

* remove unused utils.connect function

* add static type checking with mypy

Author: shubha-rajan

.flake8

@@ -1,5 +1,5 @@
 [flake8]
-ignore = E203, E266, E501, W503
+ignore = E203, E266, E501, W503, ANN101
 exclude =
   # Exclude generated code.
   **/proto/**

.mypy.ini

@@ -0,0 +1,10 @@
+[mypy]
+python_version = 3.6
+warn_unused_configs = True
+plugins = sqlmypy
+
+[mypy-google.auth.*]
+ignore_missing_imports = True
+
+[mypy-pg8000]
+ignore_missing_imports = True

google/__init__.py

@@ -12,11 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Iterable
+
 try:
     import pkg_resources
 
     pkg_resources.declare_namespace(__name__)
 except ImportError:
     import pkgutil
 
-    __path__ = pkgutil.extend_path(__path__, __name__)
+    __path__: Iterable[str] = pkgutil.extend_path(__path__, __name__)

google/cloud/__init__.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Iterable
 
 try:
     import pkg_resources
@@ -19,4 +20,4 @@
 except ImportError:
     import pkgutil
 
-    __path__ = pkgutil.extend_path(__path__, __name__)
+    __path__: Iterable[str] = pkgutil.extend_path(__path__, __name__)

google/cloud/sql/connector/__init__.py

@@ -14,9 +14,12 @@
 limitations under the License.
 """
 
+from typing import Iterable
+
 from .connector import connect
 from .instance_connection_manager import CloudSQLConnectionError
 
+
 __ALL__ = [connect, CloudSQLConnectionError]
 
 try:
@@ -26,4 +29,4 @@
 except ImportError:
     import pkgutil
 
-    __path__ = pkgutil.extend_path(__path__, __name__)
+    __path__: Iterable[str] = pkgutil.extend_path(__path__, __name__)

google/cloud/sql/connector/connector.py

@@ -21,15 +21,15 @@
 from google.cloud.sql.connector.utils import generate_keys
 
 from threading import Thread
-from typing import Optional
+from typing import Any, Dict, Optional
 
 
 # This thread is used to background processing
 _thread: Optional[Thread] = None
 _loop: Optional[asyncio.AbstractEventLoop] = None
 _keys: Optional[concurrent.futures.Future] = None
 
-_instances = {}
+_instances: Dict[str, InstanceConnectionManager] = {}
 
 
 def _get_loop() -> asyncio.AbstractEventLoop:
@@ -41,14 +41,14 @@ def _get_loop() -> asyncio.AbstractEventLoop:
     return _loop
 
 
-def _get_keys() -> concurrent.futures.Future:
+def _get_keys(loop: asyncio.AbstractEventLoop) -> concurrent.futures.Future:
     global _keys
     if _keys is None:
-        _keys = asyncio.run_coroutine_threadsafe(generate_keys(), _loop)
+        _keys = asyncio.run_coroutine_threadsafe(generate_keys(), loop)
     return _keys
 
 
-def connect(instance_connection_string, driver: str, **kwargs):
+def connect(instance_connection_string: str, driver: str, **kwargs: Any) -> Any:
     """Prepares and returns a database connection object and starts a
     background thread to refresh the certificates and metadata.
 
@@ -81,7 +81,7 @@ def connect(instance_connection_string, driver: str, **kwargs):
     if instance_connection_string in _instances:
         icm = _instances[instance_connection_string]
     else:
-        keys = _get_keys()
+        keys = _get_keys(loop)
         icm = InstanceConnectionManager(instance_connection_string, driver, keys, loop)
         _instances[instance_connection_string] = icm
 

google/cloud/sql/connector/instance_connection_manager.py

@@ -28,11 +28,22 @@
 import ssl
 import socket
 from tempfile import NamedTemporaryFile
-from typing import Any, Awaitable
+from typing import (
+    Any,
+    Awaitable,
+    Coroutine,
+    IO,
+    Optional,
+    TYPE_CHECKING,
+    Union,
+)
 
 from functools import partial
 import logging
 
+if TYPE_CHECKING:
+    import pymysql
+    import pg8000
 logger = logging.getLogger(name=__name__)
 
 APPLICATION_NAME = "cloud-sql-python-connector"
@@ -47,25 +58,25 @@ class ConnectionSSLContext(ssl.SSLContext):
     required for compatibility with pg8000 driver.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.request_ssl = False
         super(ConnectionSSLContext, self).__init__(*args, **kwargs)
 
 
 class InstanceMetadata:
     ip_address: str
-    _ca_fileobject: NamedTemporaryFile
-    _cert_fileobject: NamedTemporaryFile
-    _key_fileobject: NamedTemporaryFile
+    _ca_fileobject: IO
+    _cert_fileobject: IO
+    _key_fileobject: IO
     context: ssl.SSLContext
 
     def __init__(
         self,
         ephemeral_cert: str,
         ip_address: str,
-        private_key: str,
+        private_key: bytes,
         server_ca_cert: str,
-    ):
+    ) -> None:
         self.ip_address = ip_address
 
         self._ca_fileobject = NamedTemporaryFile(suffix=".pem")
@@ -96,8 +107,8 @@ class CloudSQLConnectionError(Exception):
     correctly.
     """
 
-    def __init__(self, *args, **kwargs) -> None:
-        super(CloudSQLConnectionError, self).__init__(self, *args, **kwargs)
+    def __init__(self, *args: Any) -> None:
+        super(CloudSQLConnectionError, self).__init__(self, *args)
 
 
 class InstanceConnectionManager:
@@ -124,9 +135,9 @@ class InstanceConnectionManager:
     # while developing on Windows.
     # Link to Github issue:
     # https://github.com/GoogleCloudPlatform/cloud-sql-python-connector/issues/22
-    _loop: asyncio.AbstractEventLoop = None
+    _loop: asyncio.AbstractEventLoop
 
-    __client_session: aiohttp.ClientSession = None
+    __client_session: Optional[aiohttp.ClientSession] = None
 
     @property
     def _client_session(self) -> aiohttp.ClientSession:
@@ -140,16 +151,17 @@ def _client_session(self) -> aiohttp.ClientSession:
             )
         return self.__client_session
 
-    _credentials: Credentials = None
+    _credentials: Optional[Credentials] = None
+    _keys: Awaitable
 
-    _instance_connection_string: str = None
-    _user_agent_string: str = None
-    _instance: str = None
-    _project: str = None
-    _region: str = None
+    _instance_connection_string: str
+    _user_agent_string: str
+    _instance: str
+    _project: str
+    _region: str
 
-    _current: asyncio.Task = None
-    _next: asyncio.Task = None
+    _current: Union[Coroutine, asyncio.Task]
+    _next: Union[Coroutine, asyncio.Task]
 
     def __init__(
         self,
@@ -174,7 +186,7 @@ def __init__(
 
         self._user_agent_string = f"{APPLICATION_NAME}/{version}+{driver_name}"
         self._loop = loop
-        self._keys: Awaitable = asyncio.wrap_future(keys, loop=self._loop)
+        self._keys = asyncio.wrap_future(keys, loop=self._loop)
         self._auth_init()
 
         logger.debug("Updating instance data")
@@ -183,17 +195,17 @@ def __init__(
         self._next = self._current
         asyncio.run_coroutine_threadsafe(self._current, self._loop)
 
-    def __del__(self):
+    def __del__(self) -> None:
         """Deconstructor to make sure ClientSession is closed and tasks have
         finished to have a graceful exit.
         """
         logger.debug("Entering deconstructor")
 
         async def _deconstruct() -> None:
-            if self._current is not None:
+            if isinstance(self._current, asyncio.Task):
                 logger.debug("Waiting for _current to be cancelled")
                 self._current.cancel()
-            if self._next is not None:
+            if isinstance(self._next, asyncio.Task):
                 logger.debug("Waiting for _next to be cancelled")
                 self._next.cancel()
             if not self._client_session.closed:
@@ -293,9 +305,9 @@ async def _schedule_refresh(self, delay: int) -> asyncio.Task:
             logger.debug("Schedule refresh task cancelled.")
             raise e
 
-        return self._perform_refresh()
+        return await self._perform_refresh()
 
-    def connect(self, driver: str, timeout: int, **kwargs):
+    def connect(self, driver: str, timeout: int, **kwargs: Any) -> Any:
         """A method that returns a DB-API connection to the database.
 
         :type driver: str
@@ -308,7 +320,7 @@ def connect(self, driver: str, timeout: int, **kwargs):
         :returns: A DB-API connection to the primary IP of the database.
         """
 
-        connect_future = asyncio.run_coroutine_threadsafe(
+        connect_future: concurrent.futures.Future = asyncio.run_coroutine_threadsafe(
             self._connect(driver, **kwargs), self._loop
         )
 
@@ -320,7 +332,7 @@ def connect(self, driver: str, timeout: int, **kwargs):
         else:
             return connection
 
-    async def _connect(self, driver: str, **kwargs) -> Any:
+    async def _connect(self, driver: str, **kwargs: Any) -> Any:
         """A method that returns a DB-API connection to the database.
 
         :type driver: str
@@ -354,7 +366,9 @@ async def _connect(self, driver: str, **kwargs) -> Any:
 
         return await self._loop.run_in_executor(None, connect_partial)
 
-    def _connect_with_pymysql(self, ip_address: str, ctx: ssl.SSLContext, **kwargs):
+    def _connect_with_pymysql(
+        self, ip_address: str, ctx: ssl.SSLContext, **kwargs: Any
+    ) -> "pymysql.connections.Connection":
         """Helper function to create a pymysql DB-API connection object.
 
         :type ip_address: str
@@ -385,7 +399,9 @@ def _connect_with_pymysql(self, ip_address: str, ctx: ssl.SSLContext, **kwargs):
         conn.connect(sock)
         return conn
 
-    def _connect_with_pg8000(self, ip_address: str, ctx: ssl.SSLContext, **kwargs):
+    def _connect_with_pg8000(
+        self, ip_address: str, ctx: ssl.SSLContext, **kwargs: Any
+    ) -> "pg8000.dbapi.Connection":
         """Helper function to create a pg8000 DB-API connection object.
 
         :type ip_address: str
@@ -409,7 +425,7 @@ def _connect_with_pg8000(self, ip_address: str, ctx: ssl.SSLContext, **kwargs):
         user = kwargs.pop("user")
         db = kwargs.pop("db")
         passwd = kwargs.pop("password")
-        ctx.request_ssl = False
+        setattr(ctx, "request_ssl", False)
         return pg8000.dbapi.connect(
             user,
             database=db,

google/cloud/sql/connector/refresh_utils.py

@@ -20,7 +20,7 @@
 from google.auth.credentials import Credentials
 import google.auth.transport.requests
 import json
-from typing import Dict, Union
+from typing import Any, Dict
 
 import logging
 
@@ -34,7 +34,7 @@ async def _get_metadata(
     credentials: Credentials,
     project: str,
     instance: str,
-) -> Dict[str, Union[Dict, str]]:
+) -> Dict[str, Any]:
     """Requests metadata from the Cloud SQL Instance
     and returns a dictionary containing the IP addresses and certificate
     authority of the Cloud SQL Instance.

google/cloud/sql/connector/utils.py

@@ -14,13 +14,14 @@
 limitations under the License.
 """
 
-import pymysql.cursors
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
 
+from typing import Tuple
 
-async def generate_keys():
+
+async def generate_keys() -> Tuple[bytes, str]:
     """A helper function to generate the private and public keys.
 
     backend - The value specified is default_backend(). This is because the
@@ -57,27 +58,7 @@ async def generate_keys():
     return priv_key, pub_key
 
 
-def connect(host, user, password, db_name):
-    """
-    Connect method to be used as a custom creator in the SQLAlchemy engine
-    creation.
-    """
-    return pymysql.connect(
-        host=host,
-        user=user,
-        password=password,
-        db=db_name,
-        ssl={
-            "ssl": {
-                "ca": "./ca.pem",
-                "cert": "./cert.pem",
-                "key": "./priv.pem",
-            }  # noqa: E501
-        },
-    )
-
-
-def write_to_file(serverCaCert, ephemeralCert, priv_key):
+def write_to_file(serverCaCert: str, ephemeralCert: str, priv_key: bytes) -> None:
     """
     Helper function to write the serverCaCert, ephemeral certificate and
     private key to .pem files

noxfile.py

@@ -26,17 +26,17 @@
 if os.path.exists("samples"):
     BLACK_PATHS.append("samples")
 
-
 @nox.session(python="3.7")
 def lint(session):
     """Run linters.
     Returns a failure if the linters find linting errors or sufficiently
     serious code quality issues.
     """
-    session.install("flake8", "black")
+    session.install("flake8", "flake8-annotations", "black", "mypy", "sqlalchemy-stubs")
     session.install("-r", "requirements.txt")
     session.run("black", "--check", *BLACK_PATHS)
     session.run("flake8", "google", "tests")
+    session.run("mypy", "google", "tests")
 
 
 @nox.session(python="3.6")