Typing SocketType, test_socket, test_highlevel_[socket, open_tcp_stream, open_tcp_listeners] (#2774)

* Removes _SocketType as a public type, expanding SocketType to have "abstract" methods raising `NotImplementedError`.
* Also replaces several methods taking `socket` as a type with taking `HasFileNo`.
* Removes `Address` type alias due to complexity, using `Any` instead.
* Add types to test_socket, test_highlevel_socket, test_highlevel_open_tcp_stream, test_highlevel_open_tcp_listeners.
* Update `FakeSocket`s and `FakeSocketFactory`s to be compatible with signatures in `SocketType`.
---------

Co-authored-by: CoolCat467 <52022020+CoolCat467@users.noreply.github.com>

Author: jakkdl

.coveragerc

@@ -25,6 +25,7 @@ exclude_lines =
   if t.TYPE_CHECKING:
   @overload
   class .*\bProtocol\b.*\):
+  raise NotImplementedError
 
 partial_branches =
   pragma: no branch

docs/source/conf.py

@@ -54,8 +54,6 @@
     ("py:class", "sync function"),
     # why aren't these found in stdlib?
     ("py:class", "types.FrameType"),
-    # TODO: temporary type
-    ("py:class", "_SocketType"),
     # these are not defined in https://docs.python.org/3/objects.inv
     ("py:class", "socket.AddressFamily"),
     ("py:class", "socket.SocketKind"),

docs/source/reference-io.rst

@@ -504,13 +504,6 @@ Socket objects
    * :meth:`~socket.socket.set_inheritable`
    * :meth:`~socket.socket.get_inheritable`
 
-The internal SocketType
-~~~~~~~~~~~~~~~~~~~~~~~~~~
-.. autoclass:: _SocketType
-..
-    TODO: adding `:members:` here gives error due to overload+_wraps on `sendto`
-    TODO: rewrite ... all of the above when fixing _SocketType vs SocketType
-
 
 .. currentmodule:: trio
 

pyproject.toml

@@ -79,15 +79,11 @@ module = [
 "trio/_tests/test_exports",
 "trio/_tests/test_file_io",
 "trio/_tests/test_highlevel_generic",
-"trio/_tests/test_highlevel_open_tcp_listeners",
-"trio/_tests/test_highlevel_open_tcp_stream",
 "trio/_tests/test_highlevel_open_unix_stream",
 "trio/_tests/test_highlevel_serve_listeners",
-"trio/_tests/test_highlevel_socket",
 "trio/_tests/test_highlevel_ssl_helpers",
 "trio/_tests/test_path",
 "trio/_tests/test_scheduler_determinism",
-"trio/_tests/test_socket",
 "trio/_tests/test_ssl",
 "trio/_tests/test_subprocess",
 "trio/_tests/test_sync",

trio/_abc.py

@@ -12,7 +12,7 @@
     from typing_extensions import Self
 
     # both of these introduce circular imports if outside a TYPE_CHECKING guard
-    from ._socket import _SocketType
+    from ._socket import SocketType
     from .lowlevel import Task
 
 
@@ -211,10 +211,10 @@ class SocketFactory(metaclass=ABCMeta):
     @abstractmethod
     def socket(
         self,
-        family: socket.AddressFamily | int | None = None,
-        type: socket.SocketKind | int | None = None,
-        proto: int | None = None,
-    ) -> _SocketType:
+        family: socket.AddressFamily | int = socket.AF_INET,
+        type: socket.SocketKind | int = socket.SOCK_STREAM,
+        proto: int = 0,
+    ) -> SocketType:
         """Create and return a socket object.
 
         Your socket object must inherit from :class:`trio.socket.SocketType`,

trio/_core/_generated_io_epoll.py

@@ -13,11 +13,10 @@
 from ._wakeup_socketpair import WakeupSocketpair
 
 if TYPE_CHECKING:
-    from socket import socket
-
     from typing_extensions import TypeAlias
 
     from .._core import Abort, RaiseCancelT
+    from .._file_io import _HasFileNo
 
 
 @attr.s(slots=True, eq=False)
@@ -290,7 +289,7 @@ def _update_registrations(self, fd: int) -> None:
         if not wanted_flags:
             del self._registered[fd]
 
-    async def _epoll_wait(self, fd: int | socket, attr_name: str) -> None:
+    async def _epoll_wait(self, fd: int | _HasFileNo, attr_name: str) -> None:
         if not isinstance(fd, int):
             fd = fd.fileno()
         waiters = self._registered[fd]
@@ -309,15 +308,15 @@ def abort(_: RaiseCancelT) -> Abort:
         await _core.wait_task_rescheduled(abort)
 
     @_public
-    async def wait_readable(self, fd: int | socket) -> None:
+    async def wait_readable(self, fd: int | _HasFileNo) -> None:
         await self._epoll_wait(fd, "read_task")
 
     @_public
-    async def wait_writable(self, fd: int | socket) -> None:
+    async def wait_writable(self, fd: int | _HasFileNo) -> None:
         await self._epoll_wait(fd, "write_task")
 
     @_public
-    def notify_closing(self, fd: int | socket) -> None:
+    def notify_closing(self, fd: int | _HasFileNo) -> None:
         if not isinstance(fd, int):
             fd = fd.fileno()
         wake_all(

trio/_core/_generated_io_kqueue.py

@@ -14,11 +14,10 @@
 from ._wakeup_socketpair import WakeupSocketpair
 
 if TYPE_CHECKING:
-    from socket import socket
-
     from typing_extensions import TypeAlias
 
     from .._core import Abort, RaiseCancelT, Task, UnboundedQueue
+    from .._file_io import _HasFileNo
 
 assert not TYPE_CHECKING or (sys.platform != "linux" and sys.platform != "win32")
 
@@ -149,7 +148,7 @@ def abort(raise_cancel: RaiseCancelT) -> Abort:
         # wait_task_rescheduled does not have its return type typed
         return await _core.wait_task_rescheduled(abort)  # type: ignore[no-any-return]
 
-    async def _wait_common(self, fd: int | socket, filter: int) -> None:
+    async def _wait_common(self, fd: int | _HasFileNo, filter: int) -> None:
         if not isinstance(fd, int):
             fd = fd.fileno()
         flags = select.KQ_EV_ADD | select.KQ_EV_ONESHOT
@@ -181,15 +180,15 @@ def abort(_: RaiseCancelT) -> Abort:
         await self.wait_kevent(fd, filter, abort)
 
     @_public
-    async def wait_readable(self, fd: int | socket) -> None:
+    async def wait_readable(self, fd: int | _HasFileNo) -> None:
         await self._wait_common(fd, select.KQ_FILTER_READ)
 
     @_public
-    async def wait_writable(self, fd: int | socket) -> None:
+    async def wait_writable(self, fd: int | _HasFileNo) -> None:
         await self._wait_common(fd, select.KQ_FILTER_WRITE)
 
     @_public
-    def notify_closing(self, fd: int | socket) -> None:
+    def notify_closing(self, fd: int | _HasFileNo) -> None:
         if not isinstance(fd, int):
             fd = fd.fileno()
 

trio/_core/_io_epoll.py

@@ -43,26 +43,26 @@
     from OpenSSL.SSL import Context
     from typing_extensions import Self, TypeAlias
 
-    from trio.socket import Address, _SocketType
+    from trio.socket import SocketType
 
 MAX_UDP_PACKET_SIZE = 65527
 
 
-def packet_header_overhead(sock: _SocketType) -> int:
+def packet_header_overhead(sock: SocketType) -> int:
     if sock.family == trio.socket.AF_INET:
         return 28
     else:
         return 48
 
 
-def worst_case_mtu(sock: _SocketType) -> int:
+def worst_case_mtu(sock: SocketType) -> int:
     if sock.family == trio.socket.AF_INET:
         return 576 - packet_header_overhead(sock)
     else:
         return 1280 - packet_header_overhead(sock)
 
 
-def best_guess_mtu(sock: _SocketType) -> int:
+def best_guess_mtu(sock: SocketType) -> int:
     return 1500 - packet_header_overhead(sock)
 
 
@@ -563,7 +563,7 @@ def _signable(*fields: bytes) -> bytes:
 
 
 def _make_cookie(
-    key: bytes, salt: bytes, tick: int, address: Address, client_hello_bits: bytes
+    key: bytes, salt: bytes, tick: int, address: Any, client_hello_bits: bytes
 ) -> bytes:
     assert len(salt) == SALT_BYTES
     assert len(key) == KEY_BYTES
@@ -581,7 +581,7 @@ def _make_cookie(
 
 
 def valid_cookie(
-    key: bytes, cookie: bytes, address: Address, client_hello_bits: bytes
+    key: bytes, cookie: bytes, address: Any, client_hello_bits: bytes
 ) -> bool:
     if len(cookie) > SALT_BYTES:
         salt = cookie[:SALT_BYTES]
@@ -603,7 +603,7 @@ def valid_cookie(
 
 
 def challenge_for(
-    key: bytes, address: Address, epoch_seqno: int, client_hello_bits: bytes
+    key: bytes, address: Any, epoch_seqno: int, client_hello_bits: bytes
 ) -> bytes:
     salt = os.urandom(SALT_BYTES)
     tick = _current_cookie_tick()
@@ -664,7 +664,7 @@ def _read_loop(read_fn: Callable[[int], bytes]) -> bytes:
 
 
 async def handle_client_hello_untrusted(
-    endpoint: DTLSEndpoint, address: Address, packet: bytes
+    endpoint: DTLSEndpoint, address: Any, packet: bytes
 ) -> None:
     if endpoint._listening_context is None:
         return
@@ -739,7 +739,7 @@ async def handle_client_hello_untrusted(
 
 
 async def dtls_receive_loop(
-    endpoint_ref: ReferenceType[DTLSEndpoint], sock: _SocketType
+    endpoint_ref: ReferenceType[DTLSEndpoint], sock: SocketType
 ) -> None:
     try:
         while True:
@@ -829,7 +829,7 @@ class DTLSChannel(trio.abc.Channel[bytes], metaclass=NoPublicConstructor):
 
     """
 
-    def __init__(self, endpoint: DTLSEndpoint, peer_address: Address, ctx: Context):
+    def __init__(self, endpoint: DTLSEndpoint, peer_address: Any, ctx: Context):
         self.endpoint = endpoint
         self.peer_address = peer_address
         self._packets_dropped_in_trio = 0
@@ -1180,7 +1180,7 @@ class DTLSEndpoint:
 
     """
 
-    def __init__(self, socket: _SocketType, *, incoming_packets_buffer: int = 10):
+    def __init__(self, socket: SocketType, *, incoming_packets_buffer: int = 10):
         # We do this lazily on first construction, so only people who actually use DTLS
         # have to install PyOpenSSL.
         global SSL
@@ -1191,7 +1191,7 @@ def __init__(self, socket: _SocketType, *, incoming_packets_buffer: int = 10):
         if socket.type != trio.socket.SOCK_DGRAM:
             raise ValueError("DTLS requires a SOCK_DGRAM socket")
         self._initialized = True
-        self.socket: _SocketType = socket
+        self.socket: SocketType = socket
 
         self.incoming_packets_buffer = incoming_packets_buffer
         self._token = trio.lowlevel.current_trio_token()
@@ -1200,7 +1200,7 @@ def __init__(self, socket: _SocketType, *, incoming_packets_buffer: int = 10):
         # as a peer provides a valid cookie, we can immediately tear down the
         # old connection.
         # {remote address: DTLSChannel}
-        self._streams: WeakValueDictionary[Address, DTLSChannel] = WeakValueDictionary()
+        self._streams: WeakValueDictionary[Any, DTLSChannel] = WeakValueDictionary()
         self._listening_context: Context | None = None
         self._listening_key: bytes | None = None
         self._incoming_connections_q = _Queue[DTLSChannel](float("inf"))