diff --git a/src/datasets/utils/patching.py b/src/datasets/utils/patching.py index 69563f562e4..37d0a54d512 100644 --- a/src/datasets/utils/patching.py +++ b/src/datasets/utils/patching.py @@ -1,3 +1,4 @@ +import builtins from importlib import import_module from .logging import get_logger @@ -93,8 +94,8 @@ def __enter__(self): if getattr(self.obj, attr) is attr_value: self.original[attr] = getattr(self.obj, attr) setattr(self.obj, attr, self.new) - elif target_attr in globals()["__builtins__"]: # if it'a s builtin like "open" - self.original[target_attr] = globals()["__builtins__"][target_attr] + elif target_attr in builtins.__dict__: # if it's a builtin like "open" + self.original[target_attr] = builtins.__dict__[target_attr] setattr(self.obj, target_attr, self.new) else: raise RuntimeError(f"Tried to patch attribute {target_attr} instead of a submodule.") diff --git a/tests/test_patching.py b/tests/test_patching.py index 42c592648f8..89753749fe2 100644 --- a/tests/test_patching.py +++ b/tests/test_patching.py @@ -101,6 +101,23 @@ def test_patch_submodule_missing_builtin(): assert _test_patching.len is len +def test_patch_submodule_builtin_when_builtins_is_module(monkeypatch): + # In the __main__ module and under PyPy, a module's ``__builtins__`` global + # is the ``builtins`` module rather than its dict. patch_submodule must + # handle both. Simulate the module case for datasets.utils.patching itself. + import builtins + + from datasets.utils import patching + + monkeypatch.setattr(patching, "__builtins__", builtins) + + mock = "__test_patch_submodule_builtins_is_module_mock__" + assert _test_patching.open is open + with patch_submodule(_test_patching, "open", mock): + assert _test_patching.open is mock + assert _test_patching.open is open + + def test_patch_submodule_start_and_stop(): mock = "__test_patch_submodule_start_and_stop_mock__" patch = patch_submodule(_test_patching, "open", mock)