typecheck _socket and _core._local (#2705)

* add type hints to _socket and _core.local

---------

Co-authored-by: Spencer Brown <[email protected]>
Co-authored-by: EXPLOSION <[email protected]>

docs/source/conf.py

@@ -62,6 +62,8 @@
     ("py:obj", "trio._abc.SendType"),
     ("py:obj", "trio._abc.T"),
     ("py:obj", "trio._abc.T_resource"),
+    ("py:class", "trio._threads.T"),
+    # why aren't these found in stdlib?
     ("py:class", "types.FrameType"),
     ("py:class", "P.args"),
     ("py:class", "P.kwargs"),

docs/source/reference-io.rst

@@ -504,6 +504,14 @@ 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

@@ -48,20 +48,25 @@ disallow_untyped_defs = false
 # downstream and users have to deal with them.
 [[tool.mypy.overrides]]
 module = [
-    "trio._path",
+    "trio._socket",
+    "trio._core._local",
+    "trio._sync",
     "trio._file_io",
 ]
+disallow_incomplete_defs = true
 disallow_untyped_defs = true
+disallow_any_generics = true
+disallow_any_decorated = true
+disallow_subclassing_any = true
 
 [[tool.mypy.overrides]]
 module = [
-    "trio._dtls",
-    "trio._abc"
+    "trio._path",
 ]
 disallow_incomplete_defs = true
 disallow_untyped_defs = true
-disallow_any_generics = true
-disallow_any_decorated = true
+#disallow_any_generics = true
+#disallow_any_decorated = true
 disallow_subclassing_any = true
 
 [tool.pytest.ini_options]

trio/_core/_local.py

@@ -1,25 +1,34 @@
+from __future__ import annotations
+
+from typing import Generic, TypeVar, final
+
 # Runvar implementations
 import attr
 
-from .._util import Final
+from .._util import Final, NoPublicConstructor
 from . import _run
 
+T = TypeVar("T")
+
+
+@final
+class _NoValue(metaclass=Final):
+    ...
 
-@attr.s(eq=False, hash=False, slots=True)
-class _RunVarToken:
-    _no_value = object()
 
-    _var = attr.ib()
-    previous_value = attr.ib(default=_no_value)
-    redeemed = attr.ib(default=False, init=False)
+@attr.s(eq=False, hash=False, slots=False)
+class RunVarToken(Generic[T], metaclass=NoPublicConstructor):
+    _var: RunVar[T] = attr.ib()
+    previous_value: T | type[_NoValue] = attr.ib(default=_NoValue)
+    redeemed: bool = attr.ib(default=False, init=False)
 
     @classmethod
-    def empty(cls, var):
-        return cls(var)
+    def _empty(cls, var: RunVar[T]) -> RunVarToken[T]:
+        return cls._create(var)
 
 
 @attr.s(eq=False, hash=False, slots=True)
-class RunVar(metaclass=Final):
+class RunVar(Generic[T], metaclass=Final):
     """The run-local variant of a context variable.
 
     :class:`RunVar` objects are similar to context variable objects,
@@ -28,44 +37,45 @@ class RunVar(metaclass=Final):
 
     """
 
-    _NO_DEFAULT = object()
-    _name = attr.ib()
-    _default = attr.ib(default=_NO_DEFAULT)
+    _name: str = attr.ib()
+    _default: T | type[_NoValue] = attr.ib(default=_NoValue)
 
-    def get(self, default=_NO_DEFAULT):
+    def get(self, default: T | type[_NoValue] = _NoValue) -> T:
         """Gets the value of this :class:`RunVar` for the current run call."""
         try:
-            return _run.GLOBAL_RUN_CONTEXT.runner._locals[self]
+            # not typed yet
+            return _run.GLOBAL_RUN_CONTEXT.runner._locals[self]  # type: ignore[return-value, index]
         except AttributeError:
             raise RuntimeError("Cannot be used outside of a run context") from None
         except KeyError:
             # contextvars consistency
-            if default is not self._NO_DEFAULT:
-                return default
+            # `type: ignore` awaiting https://github.com/python/mypy/issues/15553 to be fixed & released
+            if default is not _NoValue:
+                return default  # type: ignore[return-value]
 
-            if self._default is not self._NO_DEFAULT:
-                return self._default
+            if self._default is not _NoValue:
+                return self._default  # type: ignore[return-value]
 
             raise LookupError(self) from None
 
-    def set(self, value):
+    def set(self, value: T) -> RunVarToken[T]:
         """Sets the value of this :class:`RunVar` for this current run
         call.
 
         """
         try:
             old_value = self.get()
         except LookupError:
-            token = _RunVarToken.empty(self)
+            token = RunVarToken._empty(self)
         else:
-            token = _RunVarToken(self, old_value)
+            token = RunVarToken[T]._create(self, old_value)
 
         # This can't fail, because if we weren't in Trio context then the
         # get() above would have failed.
-        _run.GLOBAL_RUN_CONTEXT.runner._locals[self] = value
+        _run.GLOBAL_RUN_CONTEXT.runner._locals[self] = value  # type: ignore[assignment, index]
         return token
 
-    def reset(self, token):
+    def reset(self, token: RunVarToken[T]) -> None:
         """Resets the value of this :class:`RunVar` to what it was
         previously specified by the token.
 
@@ -81,14 +91,14 @@ def reset(self, token):
 
         previous = token.previous_value
         try:
-            if previous is _RunVarToken._no_value:
-                _run.GLOBAL_RUN_CONTEXT.runner._locals.pop(self)
+            if previous is _NoValue:
+                _run.GLOBAL_RUN_CONTEXT.runner._locals.pop(self)  # type: ignore[arg-type]
             else:
-                _run.GLOBAL_RUN_CONTEXT.runner._locals[self] = previous
+                _run.GLOBAL_RUN_CONTEXT.runner._locals[self] = previous  # type: ignore[index, assignment]
         except AttributeError:
             raise RuntimeError("Cannot be used outside of a run context")
 
         token.redeemed = True
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<RunVar name={self._name!r}>"