gh-3108: avoid materializing f_locals by using weakrefs to code objects instead

Author: graingert

src/trio/_core/_generated_instrumentation.py

@@ -3,10 +3,9 @@
 # *************************************************************
 from __future__ import annotations
 
-import sys
 from typing import TYPE_CHECKING
 
-from ._ki import LOCALS_KEY_KI_PROTECTION_ENABLED
+from ._ki import enable_ki_protection
 from ._run import GLOBAL_RUN_CONTEXT
 
 if TYPE_CHECKING:
@@ -15,6 +14,7 @@
 __all__ = ["add_instrument", "remove_instrument"]
 
 
+@enable_ki_protection
 def add_instrument(instrument: Instrument) -> None:
     """Start instrumenting the current run loop with the given instrument.
 
@@ -24,13 +24,13 @@ def add_instrument(instrument: Instrument) -> None:
     If ``instrument`` is already active, does nothing.
 
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.instruments.add_instrument(instrument)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def remove_instrument(instrument: Instrument) -> None:
     """Stop instrumenting the current run loop with the given instrument.
 
@@ -44,7 +44,6 @@ def remove_instrument(instrument: Instrument) -> None:
           deactivated.
 
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.instruments.remove_instrument(instrument)
     except AttributeError:

src/trio/_core/_generated_io_epoll.py

@@ -6,7 +6,7 @@
 import sys
 from typing import TYPE_CHECKING
 
-from ._ki import LOCALS_KEY_KI_PROTECTION_ENABLED
+from ._ki import enable_ki_protection
 from ._run import GLOBAL_RUN_CONTEXT
 
 if TYPE_CHECKING:
@@ -18,6 +18,7 @@
 __all__ = ["notify_closing", "wait_readable", "wait_writable"]
 
 
+@enable_ki_protection
 async def wait_readable(fd: int | _HasFileNo) -> None:
     """Block until the kernel reports that the given object is readable.
 
@@ -40,13 +41,13 @@ async def wait_readable(fd: int | _HasFileNo) -> None:
         if another task calls :func:`notify_closing` while this
         function is still working.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_readable(fd)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def wait_writable(fd: int | _HasFileNo) -> None:
     """Block until the kernel reports that the given object is writable.
 
@@ -59,13 +60,13 @@ async def wait_writable(fd: int | _HasFileNo) -> None:
         if another task calls :func:`notify_closing` while this
         function is still working.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_writable(fd)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def notify_closing(fd: int | _HasFileNo) -> None:
     """Notify waiters of the given object that it will be closed.
 
@@ -91,7 +92,6 @@ def notify_closing(fd: int | _HasFileNo) -> None:
     step, so other tasks won't be able to tell what order they happened
     in anyway.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.notify_closing(fd)
     except AttributeError:

src/trio/_core/_generated_io_kqueue.py

@@ -6,7 +6,7 @@
 import sys
 from typing import TYPE_CHECKING, Callable, ContextManager
 
-from ._ki import LOCALS_KEY_KI_PROTECTION_ENABLED
+from ._ki import enable_ki_protection
 from ._run import GLOBAL_RUN_CONTEXT
 
 if TYPE_CHECKING:
@@ -29,18 +29,19 @@
 ]
 
 
+@enable_ki_protection
 def current_kqueue() -> select.kqueue:
     """TODO: these are implemented, but are currently more of a sketch than
     anything real. See `#26
     <https://github.com/python-trio/trio/issues/26>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.current_kqueue()
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def monitor_kevent(
     ident: int,
     filter: int,
@@ -49,13 +50,13 @@ def monitor_kevent(
     anything real. See `#26
     <https://github.com/python-trio/trio/issues/26>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.monitor_kevent(ident, filter)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def wait_kevent(
     ident: int,
     filter: int,
@@ -65,7 +66,6 @@ async def wait_kevent(
     anything real. See `#26
     <https://github.com/python-trio/trio/issues/26>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_kevent(
             ident,
@@ -76,6 +76,7 @@ async def wait_kevent(
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def wait_readable(fd: int | _HasFileNo) -> None:
     """Block until the kernel reports that the given object is readable.
 
@@ -98,13 +99,13 @@ async def wait_readable(fd: int | _HasFileNo) -> None:
         if another task calls :func:`notify_closing` while this
         function is still working.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_readable(fd)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def wait_writable(fd: int | _HasFileNo) -> None:
     """Block until the kernel reports that the given object is writable.
 
@@ -117,13 +118,13 @@ async def wait_writable(fd: int | _HasFileNo) -> None:
         if another task calls :func:`notify_closing` while this
         function is still working.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_writable(fd)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def notify_closing(fd: int | _HasFileNo) -> None:
     """Notify waiters of the given object that it will be closed.
 
@@ -149,7 +150,6 @@ def notify_closing(fd: int | _HasFileNo) -> None:
     step, so other tasks won't be able to tell what order they happened
     in anyway.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.notify_closing(fd)
     except AttributeError:

src/trio/_core/_generated_io_windows.py

@@ -6,7 +6,7 @@
 import sys
 from typing import TYPE_CHECKING, ContextManager
 
-from ._ki import LOCALS_KEY_KI_PROTECTION_ENABLED
+from ._ki import enable_ki_protection
 from ._run import GLOBAL_RUN_CONTEXT
 
 if TYPE_CHECKING:
@@ -32,6 +32,7 @@
 ]
 
 
+@enable_ki_protection
 async def wait_readable(sock: _HasFileNo | int) -> None:
     """Block until the kernel reports that the given object is readable.
 
@@ -54,13 +55,13 @@ async def wait_readable(sock: _HasFileNo | int) -> None:
         if another task calls :func:`notify_closing` while this
         function is still working.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_readable(sock)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def wait_writable(sock: _HasFileNo | int) -> None:
     """Block until the kernel reports that the given object is writable.
 
@@ -73,13 +74,13 @@ async def wait_writable(sock: _HasFileNo | int) -> None:
         if another task calls :func:`notify_closing` while this
         function is still working.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_writable(sock)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def notify_closing(handle: Handle | int | _HasFileNo) -> None:
     """Notify waiters of the given object that it will be closed.
 
@@ -105,33 +106,32 @@ def notify_closing(handle: Handle | int | _HasFileNo) -> None:
     step, so other tasks won't be able to tell what order they happened
     in anyway.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.notify_closing(handle)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def register_with_iocp(handle: int | CData) -> None:
     """TODO: these are implemented, but are currently more of a sketch than
     anything real. See `#26
     <https://github.com/python-trio/trio/issues/26>`__ and `#52
     <https://github.com/python-trio/trio/issues/52>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.register_with_iocp(handle)
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def wait_overlapped(handle_: int | CData, lpOverlapped: CData | int) -> object:
     """TODO: these are implemented, but are currently more of a sketch than
     anything real. See `#26
     <https://github.com/python-trio/trio/issues/26>`__ and `#52
     <https://github.com/python-trio/trio/issues/52>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.wait_overlapped(
             handle_,
@@ -141,6 +141,7 @@ async def wait_overlapped(handle_: int | CData, lpOverlapped: CData | int) -> ob
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def write_overlapped(
     handle: int | CData,
     data: Buffer,
@@ -151,7 +152,6 @@ async def write_overlapped(
     <https://github.com/python-trio/trio/issues/26>`__ and `#52
     <https://github.com/python-trio/trio/issues/52>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.write_overlapped(
             handle,
@@ -162,6 +162,7 @@ async def write_overlapped(
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 async def readinto_overlapped(
     handle: int | CData,
     buffer: Buffer,
@@ -172,7 +173,6 @@ async def readinto_overlapped(
     <https://github.com/python-trio/trio/issues/26>`__ and `#52
     <https://github.com/python-trio/trio/issues/52>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return await GLOBAL_RUN_CONTEXT.runner.io_manager.readinto_overlapped(
             handle,
@@ -183,26 +183,26 @@ async def readinto_overlapped(
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def current_iocp() -> int:
     """TODO: these are implemented, but are currently more of a sketch than
     anything real. See `#26
     <https://github.com/python-trio/trio/issues/26>`__ and `#52
     <https://github.com/python-trio/trio/issues/52>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.current_iocp()
     except AttributeError:
         raise RuntimeError("must be called from async context") from None
 
 
+@enable_ki_protection
 def monitor_completion_key() -> ContextManager[tuple[int, UnboundedQueue[object]]]:
     """TODO: these are implemented, but are currently more of a sketch than
     anything real. See `#26
     <https://github.com/python-trio/trio/issues/26>`__ and `#52
     <https://github.com/python-trio/trio/issues/52>`__.
     """
-    sys._getframe().f_locals[LOCALS_KEY_KI_PROTECTION_ENABLED] = True
     try:
         return GLOBAL_RUN_CONTEXT.runner.io_manager.monitor_completion_key()
     except AttributeError: