From 2391d37f4270390fbc98a3878138c832579f0757 Mon Sep 17 00:00:00 2001 From: Rutvik Savaliya <92663812+Rutvik552k@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:15:10 -0400 Subject: [PATCH 1/2] Fix premature hooks for modules still initializing --- src/wrapt/importer.py | 9 +++ tests/core/test_post_import_hooks.py | 89 +++++++++++++++++++++++++++- 2 files changed, 97 insertions(+), 1 deletion(-) diff --git a/src/wrapt/importer.py b/src/wrapt/importer.py index 4c1ea5f6..7c235eef 100644 --- a/src/wrapt/importer.py +++ b/src/wrapt/importer.py @@ -79,6 +79,15 @@ def register_post_import_hook(hook, name): module = sys.modules.get(name, None) + # A watched module is visible in sys.modules before its body finishes. + # Keep additional hooks in its pending queue until our loader notifies + # them. Once notification starts, the queue has been removed, so hooks + # registered by another hook can still run immediately. + if name in _post_import_hooks and getattr( + getattr(module, "__spec__", None), "_initializing", False + ): + module = None + if module is None: _post_import_hooks.setdefault(name, []).append(hook) diff --git a/tests/core/test_post_import_hooks.py b/tests/core/test_post_import_hooks.py index 5dfb9969..80f19db3 100644 --- a/tests/core/test_post_import_hooks.py +++ b/tests/core/test_post_import_hooks.py @@ -1,7 +1,10 @@ -import unittest +import importlib +import pathlib import sys +import tempfile import threading import types +import unittest import wrapt from wrapt.importer import _post_import_hooks @@ -47,6 +50,42 @@ def hook_this(module): self.assertEqual(len(invoked), 1) + def test_hook_registered_during_import_waits_for_module_body(self): + name = "_wrapt_partial_import" + registration = name + "_registration" + invoked = [] + + def first_hook(module): + invoked.append(module.value) + wrapt.register_post_import_hook( + lambda loaded: invoked.append(loaded.value), name + ) + + with tempfile.TemporaryDirectory() as directory: + pathlib.Path(directory, name + ".py").write_text( + "import " + registration + "\nvalue = 42\n", encoding="utf-8" + ) + pathlib.Path(directory, registration + ".py").write_text( + "import wrapt\n" + "@wrapt.when_imported('" + name + "')\n" + "def hook(module):\n" + " module.hook_value = module.value\n", + encoding="utf-8", + ) + sys.path.insert(0, directory) + try: + wrapt.register_post_import_hook(first_hook, name) + module = importlib.import_module(name) + self.assertEqual(module.hook_value, 42) + self.assertEqual(invoked, [42, 42]) + self.assertNotIn(name, _post_import_hooks) + finally: + sys.path.remove(directory) + sys.path_importer_cache.pop(directory, None) + sys.modules.pop(name, None) + sys.modules.pop(registration, None) + _post_import_hooks.pop(name, None) + def test_before_and_after_import(self): invoked_one = [] invoked_two = [] @@ -72,6 +111,54 @@ def hook_this_two(module): self.assertEqual(len(invoked_one), 1) self.assertEqual(len(invoked_two), 1) + def test_hook_registered_from_another_thread_waits_for_module_body(self): + name = "_wrapt_threaded_import" + control_name = name + "_control" + control = types.ModuleType(control_name) + control.entered = threading.Event() + control.release = threading.Event() + invoked = [] + errors = [] + + def import_module(): + try: + importlib.import_module(name) + except BaseException as exc: + errors.append(exc) + + with tempfile.TemporaryDirectory() as directory: + pathlib.Path(directory, name + ".py").write_text( + "import " + control_name + " as control\n" + "control.entered.set()\n" + "assert control.release.wait(5)\n" + "value = 42\n", + encoding="utf-8", + ) + sys.path.insert(0, directory) + sys.modules[control_name] = control + thread = threading.Thread(target=import_module) + try: + wrapt.register_post_import_hook( + lambda module: invoked.append(module.value), name + ) + thread.start() + self.assertTrue(control.entered.wait(5)) + wrapt.register_post_import_hook( + lambda module: invoked.append(module.value), name + ) + self.assertEqual(invoked, []) + finally: + control.release.set() + thread.join(timeout=5) + sys.path.remove(directory) + sys.path_importer_cache.pop(directory, None) + sys.modules.pop(name, None) + sys.modules.pop(control_name, None) + _post_import_hooks.pop(name, None) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + self.assertEqual(invoked, [42, 42]) + def test_remove_from_sys_modules(self): invoked = [] From 39706525024089f36451350ce88ed3954fcc7c4d Mon Sep 17 00:00:00 2001 From: Rutvik Savaliya <92663812+Rutvik552k@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:27:52 -0400 Subject: [PATCH 2/2] Clarify pending hook behavior for loaded modules --- src/wrapt/importer.py | 7 +++---- tests/core/test_post_import_hooks.py | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/src/wrapt/importer.py b/src/wrapt/importer.py index 7c235eef..f16dd437 100644 --- a/src/wrapt/importer.py +++ b/src/wrapt/importer.py @@ -12,10 +12,9 @@ from .__wrapt__ import BaseObjectProxy # The dictionary registering any post import hooks to be triggered once -# the target module has been imported. Once a module has been imported -# and the hooks fired, the list of hooks recorded against the target -# module will be truncated but the list left in the dictionary. This -# acts as a flag to indicate that the module had already been imported. +# the target module has been imported. Notification removes the module's +# entry before invoking its hooks, so hooks registered during notification +# can run immediately. _post_import_hooks: dict[str, list[Callable]] = {} _post_import_hooks_init = False diff --git a/tests/core/test_post_import_hooks.py b/tests/core/test_post_import_hooks.py index 80f19db3..0d2a51eb 100644 --- a/tests/core/test_post_import_hooks.py +++ b/tests/core/test_post_import_hooks.py @@ -86,6 +86,21 @@ def first_hook(module): sys.modules.pop(registration, None) _post_import_hooks.pop(name, None) + def test_loaded_module_with_pending_hook_runs_new_hook_immediately(self): + name = "_wrapt_loaded_module_with_pending_hook" + invoked = [] + module = types.ModuleType(name) + try: + wrapt.register_post_import_hook( + lambda loaded: invoked.append("pending"), name + ) + sys.modules[name] = module + wrapt.register_post_import_hook(lambda loaded: invoked.append(loaded), name) + self.assertEqual(invoked, [module]) + finally: + sys.modules.pop(name, None) + _post_import_hooks.pop(name, None) + def test_before_and_after_import(self): invoked_one = [] invoked_two = []