Typing highlevel open tcp stream

pyproject.toml

@@ -52,6 +52,7 @@ module = [
     "trio._core._local",
     "trio._sync",
     "trio._file_io",
+    "trio._highlevel_open_tcp_stream.py",
 ]
 disallow_incomplete_defs = true
 disallow_untyped_defs = true

trio/_highlevel_open_tcp_stream.py

@@ -1,9 +1,13 @@
+from __future__ import annotations
+
 import sys
+from collections.abc import Generator
 from contextlib import contextmanager
+from socket import AddressFamily, SocketKind
 
 import trio
 from trio._core._multierror import MultiError
-from trio.socket import SOCK_STREAM, getaddrinfo, socket
+from trio.socket import SOCK_STREAM, _SocketType, getaddrinfo, socket
 
 if sys.version_info < (3, 11):
     from exceptiongroup import ExceptionGroup
@@ -109,8 +113,8 @@
 
 
 @contextmanager
-def close_all():
-    sockets_to_close = set()
+def close_all() -> Generator[set[_SocketType], None, None]:
+    sockets_to_close: set[_SocketType] = set()
     try:
         yield sockets_to_close
     finally:
@@ -126,7 +130,17 @@ def close_all():
             raise MultiError(errs)
 
 
-def reorder_for_rfc_6555_section_5_4(targets):
+def reorder_for_rfc_6555_section_5_4(
+    targets: list[
+        tuple[
+            AddressFamily,
+            SocketKind,
+            int,
+            str,
+            tuple[str, int] | tuple[str, int, int, int],
+        ]
+    ]
+) -> None:
     # RFC 6555 section 5.4 says that if getaddrinfo returns multiple address
     # families (e.g. IPv4 and IPv6), then you should make sure that your first
     # and second attempts use different families:
@@ -144,7 +158,7 @@ def reorder_for_rfc_6555_section_5_4(targets):
             break
 
 
-def format_host_port(host, port):
+def format_host_port(host: str | bytes, port: int) -> str:
     host = host.decode("ascii") if isinstance(host, bytes) else host
     if ":" in host:
         return f"[{host}]:{port}"
@@ -173,8 +187,12 @@ def format_host_port(host, port):
 #   AF_INET6: "..."}
 # this might be simpler after
 async def open_tcp_stream(
-    host, port, *, happy_eyeballs_delay=DEFAULT_DELAY, local_address=None
-):
+    host: str | bytes,
+    port: int,
+    *,
+    happy_eyeballs_delay: float | None = DEFAULT_DELAY,
+    local_address: str | None = None,
+) -> trio.abc.Stream:
     """Connect to the given host and port over TCP.
 
     If the given ``host`` has multiple IP addresses associated with it, then
@@ -212,9 +230,9 @@ async def open_tcp_stream(
 
       port (int): The port to connect to.
 
-      happy_eyeballs_delay (float): How many seconds to wait for each
+      happy_eyeballs_delay (float or None): How many seconds to wait for each
           connection attempt to succeed or fail before getting impatient and
-          starting another one in parallel. Set to `math.inf` if you want
+          starting another one in parallel. Set to `None` if you want
           to limit to only one connection attempt at a time (like
           :func:`socket.create_connection`). Default: 0.25 (250 ms).
 
@@ -274,7 +292,7 @@ async def open_tcp_stream(
 
     # Keeps track of the socket that we're going to complete with,
     # need to make sure this isn't automatically closed
-    winning_socket = None
+    winning_socket: _SocketType | None = None
 
     # Try connecting to the specified address. Possible outcomes:
     # - success: record connected socket in winning_socket and cancel
@@ -283,7 +301,11 @@ async def open_tcp_stream(
     #   the next connection attempt to start early
     # code needs to ensure sockets can be closed appropriately in the
     # face of crash or cancellation
-    async def attempt_connect(socket_args, sockaddr, attempt_failed):
+    async def attempt_connect(
+        socket_args: tuple[AddressFamily, SocketKind],
+        sockaddr: tuple[str, int] | tuple[str, int, int, int],
+        attempt_failed: trio.Event,
+    ) -> None:
         nonlocal winning_socket
 
         try:

trio/_tests/verify_types.json

@@ -7,11 +7,11 @@
     "warningCount": 0
   },
   "typeCompleteness": {
-    "completenessScore": 0.9072,
+    "completenessScore": 0.9088,
     "exportedSymbolCounts": {
       "withAmbiguousType": 1,
-      "withKnownType": 567,
-      "withUnknownType": 57
+      "withKnownType": 568,
+      "withUnknownType": 56
     },
     "ignoreUnknownTypesFromImports": true,
     "missingClassDocStringCount": 1,
@@ -122,7 +122,6 @@
       "trio.open_ssl_over_tcp_listeners",
       "trio.open_ssl_over_tcp_stream",
       "trio.open_tcp_listeners",
-      "trio.open_tcp_stream",
       "trio.open_unix_socket",
       "trio.run",
       "trio.run_process",