add type annotations and update mypy configuration

Author: Eugene Shershen

psqlpy_sqlalchemy/connection.py

@@ -1,6 +1,6 @@
 import typing as t
 from collections import deque
-from typing import Any, Optional, Tuple
+from typing import Any, Optional, Tuple, Union
 
 import psqlpy
 from psqlpy import row_factories
@@ -26,19 +26,23 @@ def create_server_side_cursor(self) -> "DBAPICursor":
 
 class AsyncAdapt_psqlpy_cursor(AsyncAdapt_dbapi_cursor):
     __slots__ = (
+        "_adapt_connection",
         "_arraysize",
+        "_connection",
+        "_cursor",
         "_description",
         "_invalidate_schema_cache_asof",
         "_rowcount",
+        "_rows",
     )
 
     _adapt_connection: "AsyncAdapt_psqlpy_connection"
     _connection: psqlpy.Connection
 
-    def __init__(self, adapt_connection: AsyncAdapt_dbapi_connection):
+    def __init__(self, adapt_connection: AsyncAdapt_dbapi_connection) -> None:
         self._adapt_connection = adapt_connection
         self._connection = adapt_connection._connection
-        self._rows = deque()
+        self._rows: deque[t.Any] = deque()
         self._description: t.Optional[t.List[t.Tuple[t.Any, ...]]] = None
         self._arraysize = 1
         self._rowcount = -1
@@ -155,7 +159,7 @@ def _process_parameters(
 
         import uuid
 
-        def process_value(value):
+        def process_value(value: Any) -> Any:
             """Process a single parameter value."""
             if value is None:
                 return None
@@ -341,10 +345,12 @@ def execute(
     ) -> None:
         await_only(self._prepare_execute(operation, parameters))
 
-    def executemany(self, operation, seq_of_parameters) -> None:
+    def executemany(
+        self, operation: t.Any, seq_of_parameters: t.Sequence[t.Any]
+    ) -> None:
         return await_only(self._executemany(operation, seq_of_parameters))
 
-    def setinputsizes(self, *inputsizes):
+    def setinputsizes(self, *inputsizes: t.Any) -> None:
         raise NotImplementedError
 
 
@@ -356,7 +362,9 @@ class AsyncAdapt_psqlpy_ss_cursor(
 
     _cursor: psqlpy.Cursor
 
-    def __init__(self, adapt_connection):
+    def __init__(
+        self, adapt_connection: "AsyncAdapt_psqlpy_connection"
+    ) -> None:
         self._adapt_connection = adapt_connection
         self._connection = adapt_connection._connection
         self.await_ = adapt_connection.await_
@@ -380,7 +388,7 @@ def _convert_result(
             # Return empty tuple on conversion error
             return tuple()
 
-    def close(self):
+    def close(self) -> None:
         """Enhanced close with proper state management"""
         if self._cursor is not None and not self._closed:
             try:
@@ -392,7 +400,7 @@ def close(self):
                 self._cursor = None
                 self._closed = True
 
-    def fetchone(self):
+    def fetchone(self) -> Optional[Tuple[Any, ...]]:
         """Fetch one row with enhanced error handling"""
         if self._closed or self._cursor is None:
             return None
@@ -404,7 +412,7 @@ def fetchone(self):
         except Exception:
             return None
 
-    def fetchmany(self, size=None):
+    def fetchmany(self, size: Optional[int] = None) -> t.List[Tuple[Any, ...]]:
         """Fetch many rows with enhanced error handling"""
         if self._closed or self._cursor is None:
             return []
@@ -417,7 +425,7 @@ def fetchmany(self, size=None):
         except Exception:
             return []
 
-    def fetchall(self):
+    def fetchall(self) -> t.List[Tuple[Any, ...]]:
         """Fetch all rows with enhanced error handling"""
         if self._closed or self._cursor is None:
             return []
@@ -428,7 +436,7 @@ def fetchall(self):
         except Exception:
             return []
 
-    def __iter__(self):
+    def __iter__(self) -> t.Iterator[Tuple[Any, ...]]:
         if self._closed or self._cursor is None:
             return
 
@@ -464,15 +472,15 @@ class AsyncAdapt_psqlpy_connection(AsyncAdapt_dbapi_connection):
         "readonly",
     )
 
-    def __init__(self, dbapi, connection):
+    def __init__(self, dbapi: t.Any, connection: psqlpy.Connection) -> None:
         super().__init__(dbapi, connection)
         self.isolation_level = self._isolation_setting = None
         self.readonly = False
         self.deferrable = False
         self._transaction = None
         self._started = False
         self._connection_valid = True
-        self._last_ping_time = 0
+        self._last_ping_time = 0.0
         self._performance_stats = {
             "queries_executed": 0,
             "transactions_committed": 0,
@@ -496,7 +504,7 @@ async def _start_transaction(self) -> None:
             self._started = False
             raise
 
-    def set_isolation_level(self, level):
+    def set_isolation_level(self, level: t.Any) -> None:
         self.isolation_level = self._isolation_setting = level
 
     def rollback(self) -> None:
@@ -561,7 +569,7 @@ def ping(self) -> bool:
             self._performance_stats["connection_errors"] += 1
             return False
 
-    def get_performance_stats(self) -> dict:
+    def get_performance_stats(self) -> t.Dict[str, int]:
         """Get connection performance statistics"""
         return self._performance_stats.copy()
 
@@ -574,11 +582,13 @@ def reset_performance_stats(self) -> None:
             "connection_errors": 0,
         }
 
-    def close(self):
+    def close(self) -> None:
         self.rollback()
         self._connection.close()
 
-    def cursor(self, server_side=False):
+    def cursor(
+        self, server_side: bool = False
+    ) -> Union[AsyncAdapt_psqlpy_cursor, AsyncAdapt_psqlpy_ss_cursor]:
         if server_side:
             return self._ss_cursor_cls(self)
         return self._cursor_cls(self)

psqlpy_sqlalchemy/dbapi.py

@@ -1,10 +1,12 @@
+import typing as t
+
 from sqlalchemy.util.concurrency import await_only
 
 from .connection import AsyncAdapt_psqlpy_connection
 
 
 class PSQLPyAdaptDBAPI:
-    def __init__(self, psqlpy) -> None:
+    def __init__(self, psqlpy: t.Any) -> None:
         self.psqlpy = psqlpy
         self.paramstyle = "numeric_dollar"
 
@@ -29,7 +31,9 @@ def __init__(self, psqlpy) -> None:
             if k != "connect":
                 self.__dict__[k] = v
 
-    def connect(self, *arg, **kw):
+    def connect(
+        self, *arg: t.Any, **kw: t.Any
+    ) -> AsyncAdapt_psqlpy_connection:
         creator_fn = kw.pop("async_creator_fn", self.psqlpy.connect)
 
         # Handle server_settings parameter that SQLAlchemy might pass
@@ -88,7 +92,7 @@ class PsqlpyDBAPI:
         "numeric_dollar"  # PostgreSQL uses $1, $2, etc. style parameters
     )
 
-    def __init__(self):
+    def __init__(self) -> None:
         # Initialize with psqlpy module
         import psqlpy
 
@@ -108,44 +112,52 @@ def __init__(self):
         self.NotSupportedError = _error_class
 
     # Type constructors
-    def Date(self, year, month, day):
+    def Date(self, year: int, month: int, day: int) -> t.Any:
         """Construct a date value"""
         import datetime
 
         return datetime.date(year, month, day)
 
-    def Time(self, hour, minute, second):
+    def Time(self, hour: int, minute: int, second: int) -> t.Any:
         """Construct a time value"""
         import datetime
 
         return datetime.time(hour, minute, second)
 
-    def Timestamp(self, year, month, day, hour, minute, second):
+    def Timestamp(
+        self,
+        year: int,
+        month: int,
+        day: int,
+        hour: int,
+        minute: int,
+        second: int,
+    ) -> t.Any:
         """Construct a timestamp value"""
         import datetime
 
         return datetime.datetime(year, month, day, hour, minute, second)
 
-    def DateFromTicks(self, ticks):
+    def DateFromTicks(self, ticks: float) -> t.Any:
         """Construct a date from ticks"""
         import datetime
 
         return datetime.date.fromtimestamp(ticks)
 
-    def TimeFromTicks(self, ticks):
+    def TimeFromTicks(self, ticks: float) -> t.Any:
         """Construct a time from ticks"""
         import datetime
 
         dt = datetime.datetime.fromtimestamp(ticks)
         return dt.time()
 
-    def TimestampFromTicks(self, ticks):
+    def TimestampFromTicks(self, ticks: float) -> t.Any:
         """Construct a timestamp from ticks"""
         import datetime
 
         return datetime.datetime.fromtimestamp(ticks)
 
-    def Binary(self, string):
+    def Binary(self, string: t.Union[str, bytes]) -> bytes:
         """Construct a binary value"""
         if isinstance(string, str):
             return string.encode("utf-8")
@@ -158,6 +170,8 @@ def Binary(self, string):
     DATETIME = object  # datetime objects
     ROWID = int
 
-    def connect(self, *args, **kwargs):
+    def connect(
+        self, *args: t.Any, **kwargs: t.Any
+    ) -> AsyncAdapt_psqlpy_connection:
         """Create a connection - delegates to the adapted DBAPI"""
         return self._adapt_dbapi.connect(*args, **kwargs)

psqlpy_sqlalchemy/dialect.py

@@ -24,7 +24,13 @@ class CompatibleNullPool(NullPool):
     but are commonly passed by frameworks like FastAPI with fastapi_async_sqlalchemy.
     """
 
-    def __init__(self, creator, pool_size=None, max_overflow=None, **kw):
+    def __init__(
+        self,
+        creator: t.Any,
+        pool_size: t.Optional[int] = None,
+        max_overflow: t.Optional[int] = None,
+        **kw: t.Any,
+    ) -> None:
         # Filter out pool sizing arguments that NullPool doesn't accept
         filtered_kw = {
             k: v
@@ -140,43 +146,43 @@ class _PGJSONB(sqltypes.JSON):
     class Comparator(sqltypes.JSON.Comparator):
         """Enhanced comparator with JSONB-specific operators"""
 
-        def contains(self, other):
+        def contains(self, other: t.Any) -> t.Any:
             """JSONB containment operator @>"""
             return self.operate(operators.custom_op("@>"), other)
 
-        def contained_by(self, other):
+        def contained_by(self, other: t.Any) -> t.Any:
             """JSONB contained by operator <@"""
             return self.operate(operators.custom_op("<@"), other)
 
-        def has_key(self, key):
+        def has_key(self, key: t.Any) -> t.Any:
             """JSONB has key operator ?"""
             return self.operate(operators.custom_op("?"), key)
 
-        def has_any_key(self, keys):
+        def has_any_key(self, keys: t.Any) -> t.Any:
             """JSONB has any key operator ?|"""
             return self.operate(operators.custom_op("?|"), keys)
 
-        def has_all_keys(self, keys):
+        def has_all_keys(self, keys: t.Any) -> t.Any:
             """JSONB has all keys operator ?&"""
             return self.operate(operators.custom_op("?&"), keys)
 
-        def path_exists(self, path):
+        def path_exists(self, path: t.Any) -> t.Any:
             """JSONB path exists operator @?"""
             return self.operate(operators.custom_op("@?"), path)
 
-        def path_match(self, path):
+        def path_match(self, path: t.Any) -> t.Any:
             """JSONB path match operator @@"""
             return self.operate(operators.custom_op("@@"), path)
 
-        def concat(self, other):
+        def concat(self, other: t.Any) -> t.Any:
             """JSONB concatenation operator ||"""
             return self.operate(operators.custom_op("||"), other)
 
-        def delete_key(self, key):
+        def delete_key(self, key: t.Any) -> t.Any:
             """JSONB delete key operator -"""
             return self.operate(operators.custom_op("-"), key)
 
-        def delete_path(self, path):
+        def delete_path(self, path: t.Any) -> t.Any:
             """JSONB delete path operator #-"""
             return self.operate(operators.custom_op("#-"), path)
 
@@ -222,10 +228,12 @@ class _PGNullType(sqltypes.NullType):
 class _PGUUID(UUID):
     """PostgreSQL UUID type with proper parameter binding for psqlpy."""
 
-    def bind_processor(self, dialect):
+    def bind_processor(
+        self, dialect: t.Any
+    ) -> t.Optional[t.Callable[[t.Any], t.Any]]:
         """Process UUID parameters for psqlpy compatibility."""
 
-        def process(value):
+        def process(value: t.Any) -> t.Optional[bytes]:
             if value is None:
                 return None
             if isinstance(value, uuid.UUID):
@@ -332,23 +340,23 @@ def create_connect_args(
     def set_isolation_level(
         self,
         dbapi_connection: AsyncAdapt_psqlpy_connection,
-        level,
-    ):
+        level: t.Any,
+    ) -> None:
         dbapi_connection.set_isolation_level(self._isolation_lookup[level])
 
-    def set_readonly(self, connection, value):
+    def set_readonly(self, connection: t.Any, value: t.Any) -> None:
         if value is True:
             connection.readonly = psqlpy.ReadVariant.ReadOnly
         else:
             connection.readonly = psqlpy.ReadVariant.ReadWrite
 
-    def get_readonly(self, connection):
+    def get_readonly(self, connection: t.Any) -> t.Any:
         return connection.readonly
 
-    def set_deferrable(self, connection, value):
+    def set_deferrable(self, connection: t.Any, value: t.Any) -> None:
         connection.deferrable = value
 
-    def get_deferrable(self, connection):
+    def get_deferrable(self, connection: t.Any) -> t.Any:
         return connection.deferrable