Make patch_submodule robust to __builtins__ being a module (PyPy / __main__)

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.")

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)