Remove some ignore comments and clean up typings

Author: bryanforbes

asyncpg/connect_utils.py

@@ -29,7 +29,13 @@
 from . import exceptions
 from . import protocol
 
-_Connection = typing.TypeVar('_Connection')
+if typing.TYPE_CHECKING:
+    from . import connection
+
+_Connection = typing.TypeVar(
+    '_Connection',
+    bound='connection.Connection[typing.Any]'
+)
 _Protocol = typing.TypeVar('_Protocol', bound='protocol.Protocol[typing.Any]')
 _Record = typing.TypeVar('_Record', bound=protocol.Record)
 _SSLMode = typing.TypeVar('_SSLMode', bound='SSLMode')
@@ -979,7 +985,7 @@ async def __connect_addr(
         tr.close()
         raise
 
-    con = connection_class(  # type: ignore[call-arg]
+    con = connection_class(
         pr, tr, loop, addr, config, params_input
     )
     pr.set_connection(con)

asyncpg/connection.py

@@ -525,27 +525,20 @@ async def _get_statement(
         named: typing.Union[bool, str] = False,
         use_cache: bool = True,
         ignore_custom_codec: bool = False,
-        record_class: typing.Optional[typing.Type[_OtherRecord]] = None
+        record_class: typing.Optional[typing.Type[typing.Any]] = None
     ) -> '_cprotocol.PreparedStatementState[typing.Any]':
-        record_cls: typing.Optional[
-            typing.Union[
-                typing.Type[_Record],
-                typing.Type[_OtherRecord]
-            ]
-        ] = record_class
-
-        if record_cls is None:
-            record_cls = self._protocol.get_record_class()
+        if record_class is None:
+            record_class = self._protocol.get_record_class()
         else:
-            _check_record_class(record_cls)
+            _check_record_class(record_class)
 
         if use_cache:
-            statement: typing.Optional[
-                '_cprotocol.PreparedStatementState[typing.Any]'
-            ] = self._stmt_cache.get((query, record_cls, ignore_custom_codec))
+            cache_statement = self._stmt_cache.get(
+                (query, record_class, ignore_custom_codec)
+            )
 
-            if statement is not None:
-                return statement
+            if cache_statement is not None:
+                return cache_statement
 
             # Only use the cache when:
             #  * `statement_cache_size` is greater than 0;
@@ -567,7 +560,7 @@ async def _get_statement(
             stmt_name,
             query,
             timeout,
-            record_class=record_cls,
+            record_class=record_class,
             ignore_custom_codec=ignore_custom_codec,
         )
         need_reprepare = False
@@ -609,12 +602,12 @@ async def _get_statement(
                 query,
                 timeout,
                 state=statement,
-                record_class=record_cls,
+                record_class=record_class,
             )
 
         if use_cache:
             self._stmt_cache.put(
-                (query, record_cls, ignore_custom_codec), statement)
+                (query, record_class, ignore_custom_codec), statement)
 
         # If we've just created a new statement object, check if there
         # are any statements for GC.
@@ -2595,10 +2588,10 @@ async def connect(dsn: typing.Optional[str] = None, *,
                   command_timeout: typing.Optional[float] = None,
                   ssl: typing.Optional[connect_utils.SSLType] = None,
                   direct_tls: bool = False,
-                  connection_class: typing.Type[_Connection] = Connection,  # type: ignore[assignment]  # noqa: E501
-                  record_class: typing.Type[_Record] = protocol.Record,  # type: ignore[assignment]  # noqa: E501
+                  connection_class: typing.Type[typing.Any] = Connection,
+                  record_class: typing.Type[typing.Any] = protocol.Record,
                   server_settings: typing.Optional[
-                      typing.Dict[str, str]] = None) -> _Connection:
+                      typing.Dict[str, str]] = None) -> Connection[typing.Any]:
     r"""A coroutine to establish a connection to a PostgreSQL server.
 
     The connection parameters may be specified either as a connection
@@ -2921,7 +2914,7 @@ async def connect(dsn: typing.Optional[str] = None, *,
 _StatementCacheKey = typing.Tuple[str, typing.Type[_Record], bool]
 
 
-class _StatementCacheEntry:
+class _StatementCacheEntry(typing.Generic[_Record]):
 
     __slots__ = ('_query', '_statement', '_cache', '_cleanup_cb')
 
@@ -2968,8 +2961,8 @@ def __init__(self, *, loop: asyncio.AbstractEventLoop,
         # entries dict, whereas the unused one will group in the
         # beginning of it.
         self._entries: collections.OrderedDict[
-            _StatementCacheKey['_cprotocol.Record'],
-            _StatementCacheEntry
+            _StatementCacheKey[typing.Any],
+            _StatementCacheEntry[typing.Any]
         ] = collections.OrderedDict()
 
     def __len__(self) -> int:
@@ -3004,7 +2997,9 @@ def get(
             # The cache is disabled.
             return None
 
-        entry: typing.Optional[_StatementCacheEntry] = self._entries.get(query)
+        entry: typing.Optional[
+            _StatementCacheEntry[_Record]
+        ] = self._entries.get(query)
         if entry is None:
             return None
 
@@ -3058,7 +3053,9 @@ def clear(self) -> None:
             self._clear_entry_callback(entry)
             self._on_remove(entry._statement)
 
-    def _set_entry_timeout(self, entry: _StatementCacheEntry) -> None:
+    def _set_entry_timeout(
+        self, entry: _StatementCacheEntry[typing.Any]
+    ) -> None:
         # Clear the existing timeout.
         self._clear_entry_callback(entry)
 
@@ -3071,19 +3068,23 @@ def _new_entry(
         self,
         query: _StatementCacheKey[_Record],
         statement: '_cprotocol.PreparedStatementState[_Record]'
-    ) -> _StatementCacheEntry:
+    ) -> _StatementCacheEntry[_Record]:
         entry = _StatementCacheEntry(self, query, statement)
         self._set_entry_timeout(entry)
         return entry
 
-    def _on_entry_expired(self, entry: _StatementCacheEntry) -> None:
+    def _on_entry_expired(
+        self, entry: _StatementCacheEntry[typing.Any]
+    ) -> None:
         # `call_later` callback, called when an entry stayed longer
         # than `self._max_lifetime`.
         if self._entries.get(entry._query) is entry:
             self._entries.pop(entry._query)
             self._on_remove(entry._statement)
 
-    def _clear_entry_callback(self, entry: _StatementCacheEntry) -> None:
+    def _clear_entry_callback(
+        self, entry: _StatementCacheEntry[typing.Any]
+    ) -> None:
         if entry._cleanup_cb is not None:
             entry._cleanup_cb.cancel()
 
@@ -3213,22 +3214,26 @@ def _extract_stack(limit: int = 10) -> str:
     """
     frame = sys._getframe().f_back
     try:
-        stack = traceback.StackSummary.extract(
-            traceback.walk_stack(frame), lookup_lines=False)
+        stack: typing.List[
+            traceback.FrameSummary
+        ] = traceback.StackSummary.extract(
+            traceback.walk_stack(frame), lookup_lines=False
+        )
     finally:
         del frame
 
     apg_path = asyncpg.__path__[0]
     i = 0
     while i < len(stack) and stack[i][0].startswith(apg_path):
         i += 1
-    stack = stack[i:i + limit]  # type: ignore[assignment]
+
+    stack = stack[i:i + limit]
 
     stack.reverse()
     return ''.join(traceback.format_list(stack))
 
 
-def _check_record_class(record_class: typing.Type[protocol.Record]) -> None:
+def _check_record_class(record_class: typing.Type[typing.Any]) -> None:
     if record_class is protocol.Record:
         pass
     elif (

asyncpg/protocol/protocol.pyi

@@ -220,7 +220,7 @@ class BaseProtocol(CoreProtocol, Generic[_Record]):
         state: Optional[_PreparedStatementState] = ...,
         ignore_custom_codec: bool = ...,
         record_class: Optional[type[_OtherRecord]]
-    ) -> Union[_PreparedStatementState, type[_OtherRecord]]: ...
+    ) -> Union[_PreparedStatementState, PreparedStatementState[_OtherRecord]]: ...
     async def query(self, *args: Any, **kwargs: Any) -> str: ...
     def resume_writing(self, *args: Any, **kwargs: Any) -> Any: ...
     def __reduce__(self) -> Any: ...

asyncpg/types.py

@@ -12,6 +12,8 @@
     Box, Line, LineSegment, Circle,
 )
 
+from . import compat
+
 
 __all__ = (
     'Type', 'Attribute', 'Range', 'BitString', 'Point', 'Path', 'Polygon',
@@ -54,16 +56,28 @@ class ServerVersion(typing.NamedTuple):
 
 ServerVersion.__doc__ = 'PostgreSQL server version tuple.'
 
-T = typing.TypeVar('T')
+
+class _RangeValue(compat.Protocol):
+    def __eq__(self, other: object) -> bool:
+        ...
+
+    def __lt__(self, other: '_RangeValue') -> bool:
+        ...
+
+    def __gt__(self, other: '_RangeValue') -> bool:
+        ...
+
+
+_V = typing.TypeVar('_V', bound=_RangeValue)
 
 
-class Range(typing.Generic[T]):
+class Range(typing.Generic[_V]):
     """Immutable representation of PostgreSQL `range` type."""
 
     __slots__ = '_lower', '_upper', '_lower_inc', '_upper_inc', '_empty'
 
-    def __init__(self, lower: typing.Optional[T] = None,
-                 upper: typing.Optional[T] = None, *,
+    def __init__(self, lower: typing.Optional[_V] = None,
+                 upper: typing.Optional[_V] = None, *,
                  lower_inc: bool = True,
                  upper_inc: bool = False,
                  empty: bool = False) -> None:
@@ -78,7 +92,7 @@ def __init__(self, lower: typing.Optional[T] = None,
             self._upper_inc = upper is not None and upper_inc
 
     @property
-    def lower(self) -> typing.Optional[T]:
+    def lower(self) -> typing.Optional[_V]:
         return self._lower
 
     @property
@@ -90,7 +104,7 @@ def lower_inf(self) -> bool:
         return self._lower is None and not self._empty
 
     @property
-    def upper(self) -> typing.Optional[T]:
+    def upper(self) -> typing.Optional[_V]:
         return self._upper
 
     @property
@@ -105,37 +119,37 @@ def upper_inf(self) -> bool:
     def isempty(self) -> bool:
         return self._empty
 
-    def _issubset_lower(self, other: 'Range[T]') -> bool:
+    def _issubset_lower(self, other: 'Range[_V]') -> bool:
         if other._lower is None:
             return True
         if self._lower is None:
             return False
 
-        return self._lower > other._lower or (  # type: ignore[operator]
+        return self._lower > other._lower or (
             self._lower == other._lower
             and (other._lower_inc or not self._lower_inc)
         )
 
-    def _issubset_upper(self, other: 'Range[T]') -> bool:
+    def _issubset_upper(self, other: 'Range[_V]') -> bool:
         if other._upper is None:
             return True
         if self._upper is None:
             return False
 
-        return self._upper < other._upper or (  # type: ignore[operator]
+        return self._upper < other._upper or (
             self._upper == other._upper
             and (other._upper_inc or not self._upper_inc)
         )
 
-    def issubset(self, other: 'Range[T]') -> bool:
+    def issubset(self, other: 'Range[_V]') -> bool:
         if self._empty:
             return True
         if other._empty:
             return False
 
         return self._issubset_lower(other) and self._issubset_upper(other)
 
-    def issuperset(self, other: 'Range[T]') -> bool:
+    def issuperset(self, other: 'Range[_V]') -> bool:
         return other.issubset(self)
 
     def __bool__(self) -> bool: