[asyncio] Use tagged union for getaddrinfo() (#16010)

Author: Dreamsorcerer

stdlib/@tests/test_cases/asyncio/check_getaddrinfo.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import asyncio
+from socket import AddressFamily
+from typing_extensions import assert_type
+
+
+async def check_getaddrinfo(loop: asyncio.AbstractEventLoop, base_loop: asyncio.BaseEventLoop) -> None:
+    # The address family (item 0) is a tag that discriminates the sockaddr (item 4).
+    for info in await loop.getaddrinfo("localhost", 80):
+        if info[0] == AddressFamily.AF_INET:
+            assert_type(info[4], "tuple[str, int]")
+        elif info[0] == AddressFamily.AF_INET6:
+            assert_type(info[4], "tuple[str, int, int, int] | tuple[int, bytes]")
+
+    for info in await base_loop.getaddrinfo("localhost", 80):
+        if info[0] == AddressFamily.AF_INET:
+            assert_type(info[4], "tuple[str, int]")
+        elif info[0] == AddressFamily.AF_INET6:
+            assert_type(info[4], "tuple[str, int, int, int] | tuple[int, bytes]")

stdlib/@tests/test_cases/check_socket.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import socket
+from typing_extensions import assert_type
+
+
+def check_getaddrinfo() -> None:
+    # The address family (item 0) is a tag that discriminates the sockaddr (item 4).
+    for info in socket.getaddrinfo("localhost", 80):
+        if info[0] == socket.AddressFamily.AF_INET:
+            assert_type(info[4], "tuple[str, int]")
+        elif info[0] == socket.AddressFamily.AF_INET6:
+            assert_type(info[4], "tuple[str, int, int, int] | tuple[int, bytes]")

stdlib/asyncio/base_events.pyi

@@ -10,7 +10,7 @@ from asyncio.transports import BaseTransport, DatagramTransport, ReadTransport,
 from collections.abc import Callable, Iterable, Sequence
 from concurrent.futures import Executor, ThreadPoolExecutor
 from contextvars import Context
-from socket import AddressFamily, AddressInfo, SocketKind, _Address, _RetAddress, socket
+from socket import AddressFamily, AddressInfo, _Address, _GetAddrInfoResult, _RetAddress, socket
 from typing import IO, Any, Literal, TypeAlias, TypeVar, overload
 from typing_extensions import TypeVarTuple, Unpack
 
@@ -115,7 +115,7 @@ class BaseEventLoop(AbstractEventLoop):
         type: int = 0,
         proto: int = 0,
         flags: int = 0,
-    ) -> list[tuple[AddressFamily, SocketKind, int, str, tuple[str, int] | tuple[str, int, int, int] | tuple[int, bytes]]]: ...
+    ) -> _GetAddrInfoResult: ...
     async def getnameinfo(self, sockaddr: tuple[str, int] | tuple[str, int, int, int], flags: int = 0) -> tuple[str, str]: ...
 
     if sys.version_info >= (3, 12):

stdlib/asyncio/events.pyi

@@ -11,7 +11,7 @@ from abc import ABCMeta, abstractmethod
 from collections.abc import Callable, Sequence
 from concurrent.futures import Executor
 from contextvars import Context
-from socket import AddressFamily, AddressInfo, SocketKind, _Address, _RetAddress, socket
+from socket import AddressFamily, AddressInfo, _Address, _GetAddrInfoResult, _RetAddress, socket
 from typing import IO, Any, Literal, Protocol, TypeAlias, TypeVar, overload, type_check_only
 from typing_extensions import Self, TypeVarTuple, Unpack, deprecated
 
@@ -205,7 +205,7 @@ class AbstractEventLoop:
         type: int = 0,
         proto: int = 0,
         flags: int = 0,
-    ) -> list[tuple[AddressFamily, SocketKind, int, str, tuple[str, int] | tuple[str, int, int, int] | tuple[int, bytes]]]: ...
+    ) -> _GetAddrInfoResult: ...
     @abstractmethod
     async def getnameinfo(self, sockaddr: tuple[str, int] | tuple[str, int, int, int], flags: int = 0) -> tuple[str, str]: ...
 

stdlib/socket.pyi

@@ -137,7 +137,7 @@ from _typeshed import ReadableBuffer, Unused, WriteableBuffer
 from collections.abc import Iterable
 from enum import IntEnum, IntFlag
 from io import BufferedReader, BufferedRWPair, BufferedWriter, IOBase, RawIOBase, TextIOWrapper
-from typing import Any, Final, Literal, Protocol, SupportsIndex, overload, type_check_only
+from typing import Any, Final, Literal, Protocol, SupportsIndex, TypeAlias, overload, type_check_only
 from typing_extensions import Self
 
 __all__ = [
@@ -1576,6 +1576,11 @@ def create_server(
 ) -> socket: ...
 
 # The 5th tuple item is the socket address, for IP4, IP6, or IP6 if Python is compiled with --disable-ipv6, respectively.
+_GetAddrInfoResult: TypeAlias = list[
+    tuple[Literal[AddressFamily.AF_INET], SocketKind, int, str, tuple[str, int]]
+    | tuple[Literal[AddressFamily.AF_INET6], SocketKind, int, str, tuple[str, int, int, int] | tuple[int, bytes]]
+]
+
 def getaddrinfo(
     host: bytes | str | None, port: bytes | str | int | None, family: int = 0, type: int = 0, proto: int = 0, flags: int = 0
-) -> list[tuple[AddressFamily, SocketKind, int, str, tuple[str, int] | tuple[str, int, int, int] | tuple[int, bytes]]]: ...
+) -> _GetAddrInfoResult: ...