Improve typing for connect(port=)

Author: bryanforbes

asyncpg/cluster.py

@@ -543,11 +543,14 @@ def _test_connection(self, timeout: int = 60) -> str:
                 try:
                     con: 'connection.Connection[typing.Any]' = \
                         loop.run_until_complete(
-                            asyncpg.connect(  # type: ignore[arg-type]  # noqa: E501
+                            asyncpg.connect(
                                 database='postgres',
                                 user='postgres',
                                 timeout=5, loop=loop,
-                                **self._connection_addr
+                                **typing.cast(
+                                    _ConnectionSpec,
+                                    self._connection_addr
+                                )
                             )
                         )
                 except (OSError, asyncio.TimeoutError,

asyncpg/connect_utils.py

@@ -58,7 +58,16 @@
 ]
 SSLType = typing.Union[_ParsedSSLType, SSLStringValues, bool]
 HostType = typing.Union[typing.List[str], str]
-PortType = typing.Union[typing.List[int], int]
+PortListType = typing.Union[
+    typing.List[typing.Union[int, str]],
+    typing.List[int],
+    typing.List[str],
+]
+PortType = typing.Union[
+    PortListType,
+    int,
+    str
+]
 
 
 class SSLMode(enum.IntEnum):
@@ -192,26 +201,42 @@ def _read_password_from_pgpass(
     return None
 
 
-def _validate_port_spec(hosts: typing.List[str],
-                        port: PortType) \
-        -> typing.List[int]:
+@typing.overload
+def _validate_port_spec(
+    hosts: typing.List[str],
+    port: PortListType
+) -> typing.List[int]:
+    ...
+
+
+@typing.overload
+def _validate_port_spec(
+    hosts: typing.List[str],
+    port: typing.Union[int, str]
+) -> typing.List[int]:
+    ...
+
+
+def _validate_port_spec(
+    hosts: typing.List[str],
+    port: PortType
+) -> typing.List[int]:
     if isinstance(port, list):
         # If there is a list of ports, its length must
         # match that of the host list.
         if len(port) != len(hosts):
             raise exceptions.InterfaceError(
                 'could not match {} port numbers to {} hosts'.format(
                     len(port), len(hosts)))
+        return [int(p) for p in port]
     else:
-        port = [port for _ in range(len(hosts))]
-
-    return port
+        return [int(port) for _ in range(len(hosts))]
 
 
 def _parse_hostlist(hostlist: str,
                     port: typing.Optional[PortType],
                     *, unquote: bool = False) \
-        -> typing.Tuple[typing.List[str], typing.List[int]]:
+        -> typing.Tuple[typing.List[str], PortListType]:
     if ',' in hostlist:
         # A comma-separated list of host addresses.
         hostspecs = hostlist.split(',')
@@ -242,7 +267,7 @@ def _parse_hostlist(hostlist: str,
         if hostspec[0] == '/':
             # Unix socket
             addr = hostspec
-            hostspec_port = ''
+            hostspec_port: str = ''
         elif hostspec[0] == '[':
             # IPv6 address
             m = re.match(r'(?:\[([^\]]+)\])(?::([0-9]+))?', hostspec)
@@ -470,13 +495,10 @@ def _parse_connect_dsn_and_args(*, dsn: typing.Optional[str],
         else:
             port = 5432
 
-    elif isinstance(port, (list, tuple)):
-        port = [int(p) for p in port]
-
-    else:
+    elif not isinstance(port, (list, tuple)):
         port = int(port)
 
-    port = _validate_port_spec(host, port)
+    validated_ports = _validate_port_spec(host, port)
 
     if user is None:
         user = os.getenv('PGUSER')
@@ -517,13 +539,13 @@ def _parse_connect_dsn_and_args(*, dsn: typing.Optional[str],
 
         if passfile_path is not None:
             password = _read_password_from_pgpass(
-                hosts=auth_hosts, ports=port,
+                hosts=auth_hosts, ports=validated_ports,
                 database=database, user=user,
                 passfile=passfile_path)
 
     addrs: typing.List[AddrType] = []
     have_tcp_addrs = False
-    for h, p in zip(host, port):
+    for h, p in zip(host, validated_ports):
         if h.startswith('/'):
             # UNIX socket name
             if '.s.PGSQL.' not in h: