Add typing to _ConnectionParameters and related functions

asyncpg/connect_utils.py

@@ -8,7 +8,7 @@
 
 import asyncio
 import collections
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 import enum
 import functools
 import getpass
@@ -41,31 +41,28 @@ class SSLMode(enum.IntEnum):
  verify_full = 5
 
  @classmethod
- def parse(cls, sslmode):
- if isinstance(sslmode, cls):
- return sslmode
- return getattr(cls, sslmode.replace('-', '_'))
+ def parse(cls, sslmode: typing.Union[str, SSLMode]) -> SSLMode:
+ if isinstance(sslmode, str):
+ return getattr(cls, sslmode.replace('-', '_'))
+ return sslmode
 
 
 class SSLNegotiation(compat.StrEnum):
  postgres = "postgres"
  direct = "direct"
 
 
-_ConnectionParameters = collections.namedtuple(
- 'ConnectionParameters',
- [
- 'user',
- 'password',
- 'database',
- 'ssl',
- 'sslmode',
- 'ssl_negotiation',
- 'server_settings',
- 'target_session_attrs',
- 'krbsrvname',
- 'gsslib',
- ])
+class _ConnectionParameters(typing.NamedTuple):
+ user: str
+ password: typing.Optional[str]
+ database: str
+ ssl: typing.Union[ssl_module.SSLContext, bool, str, SSLMode, None]
+ sslmode: SSLMode
+ ssl_negotiation: SSLNegotiation
+ server_settings: typing.Optional[typing.Dict[str, str]]
+ target_session_attrs: SessionAttribute
+ krbsrvname: typing.Optional[str]
+ gsslib: str
 
 
 _ClientConfiguration = collections.namedtuple(
@@ -131,11 +128,13 @@ def _read_password_file(passfile: pathlib.Path) \
 
 
 def _read_password_from_pgpass(
- *, passfile: typing.Optional[pathlib.Path],
- hosts: typing.List[str],
- ports: typing.List[int],
- database: str,
- user: str):
+ *,
+ passfile: pathlib.Path,
+ hosts: Sequence[str],
+ ports: typing.List[int],
+ database: str,
+ user: str
+) -> typing.Optional[str]:
  """Parse the pgpass file and return the matching password.
 
  :return:
@@ -167,7 +166,9 @@ def _read_password_from_pgpass(
  return None
 
 
-def _validate_port_spec(hosts, port):
+def _validate_port_spec(
+ hosts: Sequence[object], port: typing.Union[int, typing.List[int]]
+) -> typing.List[int]:
  if isinstance(port, list):
  # If there is a list of ports, its length must
  # match that of the host list.
@@ -181,15 +182,20 @@ def _validate_port_spec(hosts, port):
  return port
 
 
-def _parse_hostlist(hostlist, port, *, unquote=False):
+def _parse_hostlist(
+ hostlist: str,
+ port: typing.Union[int, typing.List[int]],
+ *,
+ unquote: bool = False,
+) -> typing.Tuple[typing.List[str], typing.List[int]]:
  if ',' in hostlist:
  # A comma-separated list of host addresses.
  hostspecs = hostlist.split(',')
  else:
  hostspecs = [hostlist]
 
- hosts = []
- hostlist_ports = []
+ hosts: typing.List[str] = []
+ hostlist_ports: typing.List[int] = []
 
  if not port:
  portspec = os.environ.get('PGPORT')
@@ -267,10 +273,25 @@ def _dot_postgresql_path(filename) -> typing.Optional[pathlib.Path]:
  return (homedir / '.postgresql' / filename).resolve()
 
 
-def _parse_connect_dsn_and_args(*, dsn, host, port, user,
- password, passfile, database, ssl,
- direct_tls, server_settings,
- target_session_attrs, krbsrvname, gsslib):
+def _parse_connect_dsn_and_args(
+ *,
+ dsn: str,
+ host: typing.Union[str, typing.List[str], typing.Tuple[str]],
+ port: typing.Union[int, typing.List[int]],
+ user: typing.Optional[str],
+ password: typing.Optional[str],
+ passfile: typing.Union[str, pathlib.Path, None],
+ database: typing.Optional[str],
+ ssl: typing.Union[bool, None, str, SSLMode],
+ direct_tls: typing.Optional[bool],
+ server_settings: typing.Optional[typing.Dict[str, str]],
+ target_session_attrs: typing.Optional[str],
+ krbsrvname: typing.Optional[str],
+ gsslib: typing.Optional[str],
+) -> typing.Tuple[
+ typing.List[typing.Union[str, typing.Tuple[str, int]]],
+ _ConnectionParameters,
+]:
  # `auth_hosts` is the version of host information for the purposes
  # of reading the pgpass file.
  auth_hosts = None
@@ -316,10 +337,12 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user,
  password = urllib.parse.unquote(dsn_password)
 
  if parsed.query:
- query = urllib.parse.parse_qs(parsed.query, strict_parsing=True)
- for key, val in query.items():
- if isinstance(val, list):
- query[key] = val[-1]
+ query = {
+ key: val[-1]
+ for key, val in urllib.parse.parse_qs(
+ parsed.query, strict_parsing=True
+ ).items()
+ }
 
  if 'port' in query:
  val = query.pop('port')
@@ -491,7 +514,7 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user,
  database=database, user=user,
  passfile=passfile)
 
- addrs = []
+ addrs: typing.List[typing.Union[str, typing.Tuple[str, int]]] = []
  have_tcp_addrs = False
  for h, p in zip(host, port):
  if h.startswith('/'):