Add mypy checks along with pre-commit hook

Author: subhashb

.pre-commit-config.yaml

@@ -21,3 +21,9 @@ repos:
     hooks:
       - id: isort
         args: ["--profile", "black", "--filter-files"]
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v0.910
+    hooks:
+      - id: mypy
+        exclude: ^tests/

message_db/client.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Dict, List
+from typing import Any, Dict, List
 from uuid import uuid4
 
 from psycopg2.extensions import connection
@@ -13,7 +13,7 @@ class MessageDB:
     """This class provides a Python interface to all MessageDB commands."""
 
     @classmethod
-    def from_url(cls, url: str, **kwargs: str) -> MessageDB:
+    def from_url(cls, url: str, **kwargs: Any) -> MessageDB:
         """Returns a MessageDB client object configured from the given URL.
 
         The general form of a connection string is:
@@ -51,10 +51,10 @@ def _write(
         connection: connection,
         stream_name: str,
         message_type: str,
-        data: Dict,
-        metadata: Dict = None,
+        data: Dict[str, Any],
+        metadata: Dict[str, Any] = None,
         expected_version: int = None,
-    ):
+    ) -> int:
         try:
             with connection.cursor(cursor_factory=RealDictCursor) as cursor:
                 cursor.execute(
@@ -85,7 +85,7 @@ def write(
         data: Dict,
         metadata: Dict = None,
         expected_version: int = None,
-    ):
+    ) -> int:
         conn = self.connection_pool.get_connection()
 
         try:
@@ -98,32 +98,34 @@ def write(
 
         return position
 
-    def write_batch(self, stream_name, data, expected_version: int = None) -> None:
+    def write_batch(self, stream_name, data, expected_version: int = None) -> int:
         conn = self.connection_pool.get_connection()
 
         try:
             with conn:
                 for record in data:
-                    expected_version = self._write(
+                    position = self._write(
                         conn,
                         stream_name,
                         record[0],
                         record[1],
                         metadata=record[2] if len(record) > 2 else None,
                         expected_version=expected_version,
                     )
+
+                    expected_version = position
         finally:
             self.connection_pool.release(conn)
 
-        return expected_version
+        return position
 
     def read(
         self,
         stream_name: str,
         sql: str = None,
         position: int = 0,
         no_of_messages: int = 1000,
-    ) -> List[Dict]:
+    ) -> List[Dict[str, Any]]:
         conn = self.connection_pool.get_connection()
         cursor = conn.cursor(cursor_factory=RealDictCursor)
 
@@ -149,7 +151,9 @@ def read(
 
         return messages
 
-    def read_stream(self, stream_name, position=0, no_of_messages=1000) -> List[Dict]:
+    def read_stream(
+        self, stream_name: str, position: int = 0, no_of_messages: int = 1000
+    ) -> List[Dict[str, Any]]:
         if "-" not in stream_name:
             raise ValueError(f"{stream_name} is not a stream")
 
@@ -160,8 +164,8 @@ def read_stream(self, stream_name, position=0, no_of_messages=1000) -> List[Dict
         )
 
     def read_category(
-        self, category_name, position=0, no_of_messages=1000
-    ) -> List[Dict]:
+        self, category_name: str, position: int = 0, no_of_messages: int = 1000
+    ) -> List[Dict[str, Any]]:
         if "-" in category_name:
             raise ValueError(f"{category_name} is not a category")
 
@@ -171,7 +175,7 @@ def read_category(
             category_name, sql=sql, position=position, no_of_messages=no_of_messages
         )
 
-    def read_last_message(self, stream_name):
+    def read_last_message(self, stream_name) -> Dict[str, Any]:
         conn = self.connection_pool.get_connection()
         cursor = conn.cursor(cursor_factory=RealDictCursor)
 

message_db/connection.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
+from typing import Any
+
 from psycopg2.extensions import connection
 from psycopg2.pool import SimpleConnectionPool
 
 
 class ConnectionPool:
     @classmethod
     def from_url(
-        cls, *args: str, max_connections: int = 100, **kwargs: str
+        cls, *args: str, max_connections: int = 100, **kwargs: Any
     ) -> ConnectionPool:
         """Return a Connection Pool configured from the given URL.
 
@@ -21,7 +23,7 @@ def from_url(
         """
         return cls(*args, max_connections=max_connections, **kwargs)
 
-    def __init__(self, *args: str, max_connections: int = 100, **kwargs: str) -> None:
+    def __init__(self, *args: str, max_connections: int = 100, **kwargs: Any) -> None:
         if not isinstance(max_connections, int) or max_connections < 0:
             raise ValueError('"max_connections" must be a positive integer')
 

pyproject.toml

@@ -19,6 +19,8 @@ pre-commit = "^2.15.0"
 isort = "^5.10.1"
 autoflake = "^1.4"
 pytest-cov = "^3.0.0"
+mypy = "^0.910"
+types-psycopg2 = "^2.9.1"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]