typecheck _socket and _core._local

pyproject.toml

@@ -43,6 +43,13 @@ disallow_untyped_defs = false
 
 # DO NOT use `ignore_errors`; it doesn't apply
 # downstream and users have to deal with them.
+[[tool.mypy.overrides]]
+module = [
+    "trio._socket",
+    "trio._core._local",
+    "trio._sync",
+]
+disallow_untyped_defs = true
 
 [tool.pytest.ini_options]
 addopts = ["--strict-markers", "--strict-config"]

trio/_core/_local.py

@@ -1,25 +1,36 @@
+from __future__ import annotations
+
+from typing import Generic, TypeVar, overload
+
 # Runvar implementations
 import attr
 
 from .._util import Final
 from . import _run
 
+T = TypeVar("T")
+C = TypeVar("C", bound="_RunVarToken")
+
+
+class NoValue(object):
+    ...
+
 
 @attr.s(eq=False, hash=False, slots=True)
-class _RunVarToken:
-    _no_value = object()
+class _RunVarToken(Generic[T]):
+    _no_value = NoValue()
 
-    _var = attr.ib()
-    previous_value = attr.ib(default=_no_value)
-    redeemed = attr.ib(default=False, init=False)
+    _var: RunVar[T] = attr.ib()
+    previous_value: T | NoValue = attr.ib(default=_no_value)
+    redeemed: bool = attr.ib(default=False, init=False)
 
     @classmethod
-    def empty(cls, var):
+    def empty(cls: type[C], var: RunVar[T]) -> C:
         return cls(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,14 +39,23 @@ class RunVar(metaclass=Final):
 
     """
 
-    _NO_DEFAULT = object()
-    _name = attr.ib()
-    _default = attr.ib(default=_NO_DEFAULT)
+    _NO_DEFAULT = NoValue()
+    _name: str = attr.ib()
+    _default: T | NoValue = attr.ib(default=_NO_DEFAULT)
+
+    @overload
+    def get(self, default: T) -> T:
+        ...
+
+    @overload
+    def get(self, default: NoValue = _NO_DEFAULT) -> T | NoValue:
+        ...
 
-    def get(self, default=_NO_DEFAULT):
+    def get(self, default: T | NoValue = _NO_DEFAULT) -> T | NoValue:
         """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:
@@ -48,24 +68,24 @@ def get(self, default=_NO_DEFAULT):
 
             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[T] = _RunVarToken.empty(self)
         else:
             token = _RunVarToken(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.
 
@@ -82,13 +102,13 @@ def reset(self, token):
         previous = token.previous_value
         try:
             if previous is _RunVarToken._no_value:
-                _run.GLOBAL_RUN_CONTEXT.runner._locals.pop(self)
+                _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}>"