refactor test cases and type annotations; update version to 0.1.0a10

Author: Eugene Shershen

psqlpy_sqlalchemy/__init__.py

@@ -2,5 +2,5 @@
 
 PsqlpyDialect = PSQLPyAsyncDialect
 
-__version__ = "0.1.0a9"
+__version__ = "0.1.0a10"
 __all__ = ["PsqlpyDialect", "PSQLPyAsyncDialect"]

psqlpy_sqlalchemy/connection.py

@@ -1,3 +1,4 @@
+import contextlib
 import typing as t
 from collections import deque
 from typing import Any, Optional, Tuple, Union
@@ -39,7 +40,9 @@ class AsyncAdapt_psqlpy_cursor(AsyncAdapt_dbapi_cursor):
     _adapt_connection: "AsyncAdapt_psqlpy_connection"
     _connection: psqlpy.Connection
 
-    def __init__(self, adapt_connection: AsyncAdapt_dbapi_connection) -> None:
+    def __init__(
+        self, adapt_connection: "AsyncAdapt_psqlpy_connection"
+    ) -> None:
         self._adapt_connection = adapt_connection
         self._connection = adapt_connection._connection
         self._rows: deque[t.Any] = deque()
@@ -177,10 +180,9 @@ def process_value(value: Any) -> Any:
             return {
                 key: process_value(value) for key, value in parameters.items()
             }
-        elif isinstance(parameters, (list, tuple)):
+        if isinstance(parameters, (list, tuple)):
             return [process_value(value) for value in parameters]
-        else:
-            return process_value(parameters)
+        return process_value(parameters)
 
     def _convert_named_params_with_casting(
         self,
@@ -330,9 +332,21 @@ async def _executemany(
             self._process_parameters(params) for params in seq_of_parameters
         ]
 
+        # Convert to the expected type for execute_many
+        converted_seq: t.List[t.List[t.Any]] = []
+        for params in processed_seq:
+            if params is None:
+                converted_seq.append([])
+            elif isinstance(params, dict):
+                converted_seq.append(list(params.values()))
+            elif isinstance(params, (list, tuple)):
+                converted_seq.append(list(params))
+            else:
+                converted_seq.append([params])
+
         return await self._connection.execute_many(
             operation,
-            processed_seq,
+            converted_seq,
             True,
         )
 
@@ -360,7 +374,7 @@ class AsyncAdapt_psqlpy_ss_cursor(
 ):
     """Enhanced server-side cursor with better async iteration support"""
 
-    _cursor: psqlpy.Cursor
+    _cursor: t.Optional[psqlpy.Cursor]
 
     def __init__(
         self, adapt_connection: "AsyncAdapt_psqlpy_connection"
@@ -377,7 +391,7 @@ def _convert_result(
     ) -> Tuple[Tuple[Any, ...], ...]:
         """Enhanced result conversion with better error handling"""
         if result is None:
-            return tuple()
+            return ()
 
         try:
             return tuple(
@@ -386,7 +400,7 @@ def _convert_result(
             )
         except Exception:
             # Return empty tuple on conversion error
-            return tuple()
+            return ()
 
     def close(self) -> None:
         """Enhanced close with proper state management"""
@@ -456,6 +470,7 @@ class AsyncAdapt_psqlpy_connection(AsyncAdapt_dbapi_connection):
     _ss_cursor_cls = AsyncAdapt_psqlpy_ss_cursor
 
     _connection: psqlpy.Connection
+    _transaction: t.Optional[psqlpy.Transaction]
 
     __slots__ = (
         "_invalidate_schema_cache_asof",
@@ -513,7 +528,7 @@ def rollback(self) -> None:
             if self._transaction is not None:
                 await_only(self._transaction.rollback())
             else:
-                await_only(self._connection.rollback())
+                await_only(self._connection.rollback())  # type: ignore[attr-defined]
             self._performance_stats["transactions_rolled_back"] += 1
         except Exception:
             self._performance_stats["connection_errors"] += 1
@@ -530,16 +545,14 @@ def commit(self) -> None:
             if self._transaction is not None:
                 await_only(self._transaction.commit())
             else:
-                await_only(self._connection.commit())
+                await_only(self._connection.commit())  # type: ignore[attr-defined]
             self._performance_stats["transactions_committed"] += 1
         except Exception as e:
             self._performance_stats["connection_errors"] += 1
             self._connection_valid = False
             # On commit failure, try to rollback
-            try:
+            with contextlib.suppress(Exception):
                 self.rollback()
-            except Exception:
-                pass
             raise e
         finally:
             self._transaction = None
@@ -549,7 +562,7 @@ def is_valid(self) -> bool:
         """Check if connection is valid"""
         return self._connection_valid and self._connection is not None
 
-    def ping(self) -> bool:
+    def ping(self, reconnect: t.Any = None) -> t.Any:
         """Ping the connection to check if it's alive"""
         import time
 

psqlpy_sqlalchemy/dbapi.py

@@ -38,10 +38,8 @@ def connect(
 
         # Handle server_settings parameter that SQLAlchemy might pass
         server_settings = kw.pop("server_settings", None)
-        if server_settings:
-            # Map server_settings to individual psqlpy parameters
-            if "application_name" in server_settings:
-                kw["application_name"] = server_settings["application_name"]
+        if server_settings and "application_name" in server_settings:
+            kw["application_name"] = server_settings["application_name"]
             # Add other server_settings mappings as needed
 
         # Filter out any other unsupported parameters that SQLAlchemy might pass

psqlpy_sqlalchemy/dialect.py

@@ -7,6 +7,7 @@
 from sqlalchemy import URL, util
 from sqlalchemy.dialects.postgresql.base import INTERVAL, UUID, PGDialect
 from sqlalchemy.dialects.postgresql.json import JSONPathType
+from sqlalchemy.engine.interfaces import DBAPIConnection
 from sqlalchemy.pool import AsyncAdaptedQueuePool, NullPool
 from sqlalchemy.sql import operators, sqltypes
 from sqlalchemy.sql.functions import GenericFunction
@@ -41,80 +42,80 @@ def __init__(
 
 
 # JSONB aggregation functions
-class jsonb_agg(GenericFunction):
+class jsonb_agg(GenericFunction[t.Any]):
     """JSONB aggregation function"""
 
-    type = sqltypes.JSON
+    type_ = sqltypes.JSON
     name = "jsonb_agg"
 
 
-class jsonb_object_agg(GenericFunction):
+class jsonb_object_agg(GenericFunction[t.Any]):
     """JSONB object aggregation function"""
 
-    type = sqltypes.JSON
+    type_ = sqltypes.JSON
     name = "jsonb_object_agg"
 
 
-class jsonb_build_array(GenericFunction):
+class jsonb_build_array(GenericFunction[t.Any]):
     """JSONB build array function"""
 
-    type = sqltypes.JSON
+    type_ = sqltypes.JSON
     name = "jsonb_build_array"
 
 
-class jsonb_build_object(GenericFunction):
+class jsonb_build_object(GenericFunction[t.Any]):
     """JSONB build object function"""
 
-    type = sqltypes.JSON
+    type_ = sqltypes.JSON
     name = "jsonb_build_object"
 
 
-class jsonb_extract_path(GenericFunction):
+class jsonb_extract_path(GenericFunction[t.Any]):
     """JSONB extract path function"""
 
-    type = sqltypes.JSON
+    type_ = sqltypes.JSON
     name = "jsonb_extract_path"
 
 
-class jsonb_extract_path_text(GenericFunction):
+class jsonb_extract_path_text(GenericFunction[t.Any]):
     """JSONB extract path as text function"""
 
-    type = sqltypes.Text
+    type_ = sqltypes.Text
     name = "jsonb_extract_path_text"
 
 
-class jsonb_path_exists(GenericFunction):
+class jsonb_path_exists(GenericFunction[t.Any]):
     """JSONB path exists function"""
 
-    type = sqltypes.Boolean
+    type_ = sqltypes.Boolean
     name = "jsonb_path_exists"
 
 
-class jsonb_path_match(GenericFunction):
+class jsonb_path_match(GenericFunction[t.Any]):
     """JSONB path match function"""
 
-    type = sqltypes.Boolean
+    type_ = sqltypes.Boolean
     name = "jsonb_path_match"
 
 
-class jsonb_path_query(GenericFunction):
+class jsonb_path_query(GenericFunction[t.Any]):
     """JSONB path query function"""
 
-    type = sqltypes.JSON
+    type_ = sqltypes.JSON
     name = "jsonb_path_query"
 
 
-class jsonb_path_query_array(GenericFunction):
+class jsonb_path_query_array(GenericFunction[t.Any]):
     """JSONB path query array function"""
 
-    type = sqltypes.JSON
+    type_ = sqltypes.JSON
     name = "jsonb_path_query_array"
 
 
-class jsonb_path_query_first(GenericFunction):
+class jsonb_path_query_first(GenericFunction[t.Any]):
     """JSONB path query first function"""
 
-    type = sqltypes.JSON
+    type_ = sqltypes.JSON
     name = "jsonb_path_query_first"
 
 
@@ -143,10 +144,10 @@ class _PGJSONB(sqltypes.JSON):
     __visit_name__ = "JSONB"
     render_bind_cast = True
 
-    class Comparator(sqltypes.JSON.Comparator):
+    class Comparator(sqltypes.JSON.Comparator[t.Any]):
         """Enhanced comparator with JSONB-specific operators"""
 
-        def contains(self, other: t.Any) -> t.Any:
+        def contains(self, other: t.Any, **kw: t.Any) -> t.Any:
             """JSONB containment operator @>"""
             return self.operate(operators.custom_op("@>"), other)
 
@@ -225,7 +226,7 @@ class _PGNullType(sqltypes.NullType):
     render_bind_cast = True
 
 
-class _PGUUID(UUID):
+class _PGUUID(UUID[t.Any]):
     """PostgreSQL UUID type with proper parameter binding for psqlpy."""
 
     def bind_processor(
@@ -339,10 +340,13 @@ def create_connect_args(
 
     def set_isolation_level(
         self,
-        dbapi_connection: AsyncAdapt_psqlpy_connection,
+        dbapi_connection: DBAPIConnection,
         level: t.Any,
     ) -> None:
-        dbapi_connection.set_isolation_level(self._isolation_lookup[level])
+        psqlpy_connection = t.cast(
+            AsyncAdapt_psqlpy_connection, dbapi_connection
+        )
+        psqlpy_connection.set_isolation_level(self._isolation_lookup[level])
 
     def set_readonly(self, connection: t.Any, value: t.Any) -> None:
         if value is True:

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "psqlpy-sqlalchemy"
-version = "0.1.0a9"
+version = "0.1.0a10"
 description = "SQLAlchemy dialect for psqlpy PostgreSQL driver"
 readme = "README.md"
 license = {text = "MIT"}
@@ -68,14 +68,23 @@ target-version = "py38"
 
 [tool.ruff.lint]
 select = [
-    "E",  # pycodestyle errors
-    "W",  # pycodestyle warnings
-    "F",  # pyflakes
-    "I",  # isort
-    "UP", # pyupgrade
+    "UP",    # pyupgrade
+    "E",     # pycodestyle errors
+    "W",     # pycodestyle warnings
+    "F",     # pyflakes
+    "I",     # isort
+    "C",     # flake8-comprehensions
+    "B",     # flake8-bugbear
+    "PTH",   # flake8-use-pathlib
+    "ASYNC", # flake8-async
+    "SIM",   # flake8-simplify
+    "RET",   # flake8-return
 ]
 ignore = [
-    "E501", # line too long (handled by formatter)
+    "E501",   # line too long
+    "C901",   # too complex
+    "B008",   # do not perform function calls in argument defaults
+    "B904",   # Within an `except` clause, raise exceptions with `raise ... from err`
 ]
 
 [tool.ruff.format]
@@ -108,6 +117,7 @@ ignore_missing_imports = true
 
 [[tool.mypy.overrides]]
 module = "psqlpy_sqlalchemy.*"
+warn_unused_ignores = false
 
 [[tool.mypy.overrides]]
 module = "tests.*"