Fixes for pyright

Author: bryanforbes

asyncpg/__init__.py

@@ -17,10 +17,5 @@
 from ._version import __version__  # NOQA
 
 
-__all__ = (
-    'connect',
-    'create_pool',
-    'Pool',
-    'Record',
-    'Connection'
-) + exceptions.__all__
+__all__ = ['connect', 'create_pool', 'Pool', 'Record', 'Connection']
+__all__ += exceptions.__all__

asyncpg/cluster.py

@@ -18,7 +18,12 @@
 import textwrap
 import time
 import typing
-import typing_extensions
+
+# Work around https://github.com/microsoft/pyright/issues/3012
+if sys.version_info >= (3, 8):
+    from typing import Final
+else:
+    from typing_extensions import Final
 
 import asyncpg
 from asyncpg import compat
@@ -35,7 +40,7 @@ class _ConnectionSpec(compat.TypedDict):
     port: str
 
 
-_system: typing_extensions.Final = platform.uname().system
+_system: Final = platform.uname().system
 
 if _system == 'Windows':
     def platform_exe(name: str) -> str:
@@ -338,7 +343,7 @@ def reset_wal(self, *, oid: typing.Optional[int] = None,
             raise ClusterError(
                 'cannot modify WAL status: cluster is running')
 
-        opts = []
+        opts: typing.List[str] = []
         if oid is not None:
             opts.extend(['-o', str(oid)])
         if xid is not None:
@@ -526,14 +531,15 @@ def _test_connection(self, timeout: int = 60) -> str:
                         continue
 
                 try:
-                    con = loop.run_until_complete(
-                        asyncpg.connect(  # type: ignore[arg-type]  # noqa: E501
-                            database='postgres',
-                            user='postgres',
-                            timeout=5, loop=loop,
-                            **self._connection_addr
+                    con: 'connection.Connection[typing.Any]' = \
+                        loop.run_until_complete(
+                            asyncpg.connect(  # type: ignore[arg-type]  # noqa: E501
+                                database='postgres',
+                                user='postgres',
+                                timeout=5, loop=loop,
+                                **self._connection_addr
+                            )
                         )
-                    )
                 except (OSError, asyncio.TimeoutError,
                         exceptions.CannotConnectNowError,
                         exceptions.PostgresConnectionError):
@@ -562,7 +568,7 @@ def _run_pg_config(self, pg_config_path: str) -> typing.Dict[str, str]:
                 'pg_config exited with status {:d}: {!r}'.format(
                     process.returncode, stderr))
         else:
-            config = {}
+            config: typing.Dict[str, str] = {}
 
             for line in stdout.splitlines():
                 k, eq, v = line.decode('utf-8').partition('=')

asyncpg/compat.py

@@ -9,7 +9,6 @@
 import pathlib
 import sys
 import typing
-import typing_extensions
 
 
 if sys.version_info >= (3, 8):
@@ -27,7 +26,6 @@
 
 
 _T = typing.TypeVar('_T')
-PY_37: typing_extensions.Final = sys.version_info >= (3, 7)
 
 
 if sys.platform == 'win32':

asyncpg/connect_utils.py

@@ -20,11 +20,16 @@
 import sys
 import time
 import typing
-import typing_extensions
 import urllib.parse
 import warnings
 import inspect
 
+# Work around https://github.com/microsoft/pyright/issues/3012
+if sys.version_info >= (3, 8):
+    from typing import Final
+else:
+    from typing_extensions import Final
+
 from . import compat
 from . import exceptions
 from . import protocol
@@ -37,10 +42,13 @@
     bound='connection.Connection[typing.Any]'
 )
 _Protocol = typing.TypeVar('_Protocol', bound='protocol.Protocol[typing.Any]')
+_AsyncProtocol = typing.TypeVar(
+    '_AsyncProtocol', bound='asyncio.protocols.Protocol'
+)
 _Record = typing.TypeVar('_Record', bound=protocol.Record)
 _SSLMode = typing.TypeVar('_SSLMode', bound='SSLMode')
 
-_TPTupleType = typing.Tuple[asyncio.WriteTransport, _Protocol]
+_TPTupleType = typing.Tuple[asyncio.WriteTransport, _AsyncProtocol]
 AddrType = typing.Union[typing.Tuple[str, int], str]
 SSLStringValues = compat.Literal[
     'disable', 'prefer', 'allow', 'require', 'verify-ca', 'verify-full'
@@ -76,7 +84,12 @@ def parse(
 
 class _ConnectionParameters(typing.NamedTuple):
     user: str
-    password: typing.Optional[str]
+    password: typing.Union[
+        str,
+        typing.Callable[[], str],
+        typing.Callable[[], typing.Awaitable[str]],
+        None
+    ]
     database: str
     ssl: typing.Optional[_ParsedSSLType]
     sslmode: typing.Optional[SSLMode]
@@ -92,8 +105,8 @@ class _ClientConfiguration(typing.NamedTuple):
     max_cacheable_statement_size: int
 
 
-_system: typing_extensions.Final = platform.uname().system
-PGPASSFILE: typing_extensions.Final = (
+_system: Final = platform.uname().system
+PGPASSFILE: Final = (
     'pgpass.conf' if _system == 'Windows' else '.pgpass'
 )
 
@@ -102,7 +115,7 @@ def _read_password_file(
     passfile: pathlib.Path
 ) -> typing.List[typing.Tuple[str, ...]]:
 
-    passtab = []
+    passtab: typing.List[typing.Tuple[str, ...]] = []
 
     try:
         if not passfile.exists():
@@ -295,7 +308,8 @@ def _parse_connect_dsn_and_args(*, dsn: typing.Optional[str],
                                 ssl: typing.Optional[SSLType],
                                 direct_tls: bool,
                                 connect_timeout: float,
-                                server_settings: typing.Dict[str, str]) \
+                                server_settings: typing.Optional[
+                                    typing.Dict[str, str]]) \
         -> typing.Tuple[typing.List[typing.Union[typing.Tuple[str, int], str]],
                         _ConnectionParameters]:
     # `auth_hosts` is the version of host information for the purposes
@@ -672,7 +686,8 @@ def _parse_connect_arguments(*, dsn: typing.Optional[str],
                              max_cacheable_statement_size: int,
                              ssl: typing.Optional[SSLType],
                              direct_tls: bool,
-                             server_settings: typing.Dict[str, str]) \
+                             server_settings: typing.Optional[
+                                 typing.Dict[str, str]]) \
         -> typing.Tuple[typing.List[AddrType], _ConnectionParameters,
                         _ClientConfiguration]:
 
@@ -786,7 +801,7 @@ async def _create_ssl_connection(
     loop: asyncio.AbstractEventLoop,
     ssl_context: ssl_module.SSLContext,
     ssl_is_advisory: typing.Optional[bool] = False
-) -> _TPTupleType[typing.Union[_Protocol, '_CancelProto']]:
+) -> _TPTupleType[typing.Any]:
 
     tr, pr = typing.cast(
         typing.Tuple[asyncio.WriteTransport, TLSUpgradeProto],
@@ -915,14 +930,17 @@ async def __connect_addr(
 ) -> _Connection:
     connected = _create_future(loop)
 
-    proto_factory = lambda: protocol.Protocol(
+    proto_factory: typing.Callable[
+        [], 'protocol.Protocol[_Record]'
+    ] = lambda: protocol.Protocol(
         addr, connected, params, record_class, loop)
 
     if isinstance(addr, str):
         # UNIX socket
         connector = typing.cast(
-            typing.Coroutine[typing.Any, None,
-                             _TPTupleType['protocol.Protocol[_Record]']],
+            typing.Coroutine[
+                typing.Any, None, _TPTupleType['protocol.Protocol[_Record]']
+            ],
             loop.create_unix_connection(proto_factory, addr)
         )
 
@@ -939,9 +957,11 @@ async def __connect_addr(
             ssl_is_advisory=params.sslmode == SSLMode.prefer)
     else:
         connector = typing.cast(
-            typing.Coroutine[typing.Any, None,
-                             _TPTupleType['protocol.Protocol[_Record]']],
-            loop.create_connection(proto_factory, *addr))
+            typing.Coroutine[
+                typing.Any, None, _TPTupleType['protocol.Protocol[_Record]']
+            ],
+            loop.create_connection(proto_factory, *addr)
+        )
 
     connector_future = asyncio.ensure_future(connector)
     before = time.monotonic()
@@ -1043,13 +1063,12 @@ async def _cancel(*, loop: asyncio.AbstractEventLoop,
                   params: _ConnectionParameters,
                   backend_pid: int, backend_secret: str) -> None:
 
-    proto_factory = lambda: _CancelProto(loop)
+    proto_factory: typing.Callable[
+        [], _CancelProto
+    ] = lambda: _CancelProto(loop)
 
     if isinstance(addr, str):
-        tr, pr = typing.cast(
-            typing.Tuple[asyncio.WriteTransport, _CancelProto],
-            await loop.create_unix_connection(proto_factory, addr)
-        )
+        tr, pr = await loop.create_unix_connection(proto_factory, addr)
     else:
         if params.ssl and params.sslmode != SSLMode.allow:
             tr, pr = await _create_ssl_connection(
@@ -1059,17 +1078,15 @@ async def _cancel(*, loop: asyncio.AbstractEventLoop,
                 ssl_context=params.ssl,
                 ssl_is_advisory=params.sslmode == SSLMode.prefer)
         else:
-            tr, pr = typing.cast(
-                typing.Tuple[asyncio.WriteTransport, _CancelProto],
-                await loop.create_connection(proto_factory, *addr))
+            tr, pr = await loop.create_connection(proto_factory, *addr)
             _set_nodelay(_get_socket(tr))
 
     # Pack a CancelRequest message
     msg = struct.pack('!llll', 16, 80877102, backend_pid, backend_secret)
 
     try:
-        tr.write(msg)
-        await pr.on_disconnect
+        typing.cast(typing.Any, tr).write(msg)
+        await typing.cast(typing.Any, pr).on_disconnect
     finally:
         tr.close()
 

asyncpg/connection.py

@@ -44,7 +44,6 @@
 _OtherRecord = typing.TypeVar('_OtherRecord', bound=protocol.Record)
 _RecordsType = typing.List[_Record]
 _RecordsExtraType = typing.Tuple[_RecordsType[_Record], bytes, bool]
-_AnyCallable = typing.Callable[..., typing.Any]
 
 OutputType = typing.Union['os.PathLike[typing.Any]',
                           typing.BinaryIO,
@@ -1060,7 +1059,8 @@ async def copy_from_table(self, table_name: str, *,
                               header: typing.Optional[bool] = None,
                               quote: typing.Optional[str] = None,
                               escape: typing.Optional[str] = None,
-                              force_quote: typing.Optional[bool] = None,
+                              force_quote: typing.Union[
+                                  bool, typing.Iterable[str], None] = None,
                               encoding: typing.Optional[str] = None) -> str:
         """Copy table contents to a file or file-like object.
 
@@ -1139,7 +1139,8 @@ async def copy_from_query(self, query: str, *args: typing.Any,
                               header: typing.Optional[bool] = None,
                               quote: typing.Optional[str] = None,
                               escape: typing.Optional[str] = None,
-                              force_quote: typing.Optional[bool] = None,
+                              force_quote: typing.Union[
+                                  bool, typing.Iterable[str], None] = None,
                               encoding: typing.Optional[str] = None) -> str:
         """Copy the results of a query to a file or file-like object.
 
@@ -1212,9 +1213,12 @@ async def copy_to_table(self, table_name: str, *,
                             header: typing.Optional[bool] = None,
                             quote: typing.Optional[str] = None,
                             escape: typing.Optional[str] = None,
-                            force_quote: typing.Optional[bool] = None,
-                            force_not_null: typing.Optional[bool] = None,
-                            force_null: typing.Optional[bool] = None,
+                            force_quote: typing.Union[
+                                bool, typing.Iterable[str], None] = None,
+                            force_not_null: typing.Union[
+                                bool, typing.Iterable[str], None] = None,
+                            force_null: typing.Union[
+                                bool, typing.Iterable[str], None] = None,
                             encoding: typing.Optional[str] = None) -> str:
         """Copy data to the specified table.
 
@@ -1390,13 +1394,16 @@ def _format_copy_opts(self, *,
                           header: typing.Optional[bool] = None,
                           quote: typing.Optional[str] = None,
                           escape: typing.Optional[str] = None,
-                          force_quote: typing.Optional[bool] = None,
-                          force_not_null: typing.Optional[bool] = None,
-                          force_null: typing.Optional[bool] = None,
+                          force_quote: typing.Union[
+                              bool, typing.Iterable[str], None] = None,
+                          force_not_null: typing.Union[
+                              bool, typing.Iterable[str], None] = None,
+                          force_null: typing.Union[
+                              bool, typing.Iterable[str], None] = None,
                           encoding: typing.Optional[str] = None) -> str:
         kwargs = dict(locals())
         kwargs.pop('self')
-        opts = []
+        opts: typing.List[str] = []
 
         if force_quote is not None and isinstance(force_quote, bool):
             kwargs.pop('force_quote')
@@ -1975,7 +1982,7 @@ def _get_reset_query(self) -> str:
 
         caps = self._server_caps
 
-        _reset_query = []
+        _reset_query: typing.List[str] = []
         if caps.advisory_locks:
             _reset_query.append('SELECT pg_advisory_unlock_all();')
         if caps.sql_close_all:
@@ -2337,7 +2344,8 @@ async def __execute(
     ) -> typing.Tuple[
         typing.Any, '_cprotocol.PreparedStatementState[typing.Any]'
     ]:
-        executor = lambda stmt, timeout: self._protocol.bind_execute(
+        executor: Executor[_OtherRecord] = \
+            lambda stmt, timeout: self._protocol.bind_execute(
             stmt, args, '', limit, return_status, timeout)
         timeout = self._protocol._get_timeout(timeout)
         return await self._do_execute(
@@ -2351,7 +2359,8 @@ async def __execute(
     async def _executemany(self, query: str,
                            args: typing.Iterable[typing.Sequence[typing.Any]],
                            timeout: typing.Optional[float]) -> None:
-        executor = lambda stmt, timeout: self._protocol.bind_execute_many(
+        executor: Executor[_Record] = \
+            lambda stmt, timeout: self._protocol.bind_execute_many(
             stmt, args, '', timeout)
         timeout = self._protocol._get_timeout(timeout)
         with self._stmt_exclusive_section:
@@ -3222,7 +3231,7 @@ def _extract_stack(limit: int = 10) -> str:
     finally:
         del frame
 
-    apg_path = asyncpg.__path__[0]
+    apg_path = list(asyncpg.__path__)[0]
     i = 0
     while i < len(stack) and stack[i][0].startswith(apg_path):
         i += 1

asyncpg/connresource.py

@@ -16,16 +16,14 @@
     from . import connection as _connection
 
 
-_ConnectionResource = typing.TypeVar('_ConnectionResource',
-                                     bound='ConnectionResource')
 _Callable = typing.TypeVar('_Callable', bound=typing.Callable[..., typing.Any])
 
 
 def guarded(meth: _Callable) -> _Callable:
     """A decorator to add a sanity check to ConnectionResource methods."""
 
     @functools.wraps(meth)
-    def _check(self: _ConnectionResource,
+    def _check(self: 'ConnectionResource',
                *args: typing.Any,
                **kwargs: typing.Any) -> typing.Any:
         self._check_conn_validity(meth.__name__)

asyncpg/cursor.py

@@ -235,7 +235,7 @@ async def _exec(self, n: int,
         return buffer
 
     def __repr__(self) -> str:
-        attrs = []
+        attrs: typing.List[str] = []
         if self._exhausted:
             attrs.append('exhausted')
         attrs.append('')  # to separate from id