Add possibility to add a cleanup hook to modules

- may be needed for modules that cannot cleanly reload
- used for pandas.core.arrays.arrow.extension_types, see #947

Author: mrbean-bremen

.github/workflows/testsuite.yml

@@ -14,7 +14,7 @@ jobs:
       with:
         python-version: "3.10"
     - name: install pytype
-      run: pip install setuptools pytype pytest scandir pathlib2 pandas xlrd django
+      run: pip install setuptools pytype pytest scandir pathlib2 pandas xlrd django pyarrow
     - name: Run pytype
       run: |
         pytype pyfakefs --keep-going --exclude pyfakefs/tests/* --exclude pyfakefs/pytest_tests/*
@@ -114,7 +114,7 @@ jobs:
         run: |
           pip install -r requirements.txt
           pip install -U pytest==${{ matrix.pytest-version }}
-          pip install opentimelineio undefined
+          pip install opentimelineio undefined pandas parquet pyarrow
           pip install -e .
           if [[ '${{ matrix.pytest-version }}' == '4.0.2' ]]; then
              pip install -U attrs==19.1.0

CHANGES.md

@@ -1,6 +1,12 @@
 # pyfakefs Release Notes
 The released versions correspond to PyPI releases.
 
+## Unreleased
+
+### Fixes
+* Fixed a specific problem on reloading a pandas-related module (see [#947](../../issues/947)),
+  added possibility for unload hooks for specific modules
+
 ## [Version 5.3.5](https://pypi.python.org/pypi/pyfakefs/5.3.5) (2024-01-30)
 Fixes a regression.
 

docs/modules.rst

@@ -34,7 +34,7 @@ Unittest module classes
 .. autoclass:: pyfakefs.fake_filesystem_unittest.TestCase
 
 .. autoclass:: pyfakefs.fake_filesystem_unittest.Patcher
-    :members: setUp, tearDown, pause, resume
+    :members: setUp, tearDown, pause, resume, register_cleanup_handler
 
 .. automodule:: pyfakefs.fake_filesystem_unittest
     :members: patchfs

pyfakefs/fake_filesystem_unittest.py

@@ -619,18 +619,21 @@ def __init__(
         self.use_cache = use_cache
         self.use_dynamic_patch = use_dynamic_patch
         self.module_cleanup_mode = module_cleanup_mode
+        self.cleanup_handlers: Dict[str, Callable[[], bool]] = {}
 
         if use_known_patches:
             from pyfakefs.patched_packages import (
                 get_modules_to_patch,
                 get_classes_to_patch,
                 get_fake_module_classes,
+                get_cleanup_handlers,
             )
 
             modules_to_patch = modules_to_patch or {}
             modules_to_patch.update(get_modules_to_patch())
             self._class_modules.update(get_classes_to_patch())
             self._fake_module_classes.update(get_fake_module_classes())
+            self.cleanup_handlers.update(get_cleanup_handlers())
 
         if modules_to_patch is not None:
             for name, fake_module in modules_to_patch.items():
@@ -680,6 +683,22 @@ def clear_cache(self) -> None:
         """Clear the module cache (convenience instance method)."""
         self.__class__.clear_fs_cache()
 
+    def register_cleanup_handler(self, name: str, handler: Callable[[], bool]):
+        """Register a handler for cleaning up a module after it had been loaded by
+        the dynamic patcher. This allows to handle modules that cannot be reloaded
+        without unwanted side effects.
+
+        Args:
+            name: The fully qualified module name.
+            handler: A callable that may do any module cleanup, or do nothing
+                and return `True` in case reloading shall be prevented.
+
+        Returns:
+            `True` if no further cleanup/reload shall occur after the handler is
+                executed, `False` if the cleanup/reload shall still happen.
+        """
+        self.cleanup_handlers[name] = handler
+
     def _init_fake_module_classes(self) -> None:
         # IMPORTANT TESTING NOTE: Whenever you add a new module below, test
         # it by adding an attribute in fixtures/module_with_attributes.py
@@ -946,7 +965,9 @@ def start_patching(self) -> None:
             self.patch_functions()
             self.patch_defaults()
 
-            self._dyn_patcher = DynamicPatcher(self)
+            self._dyn_patcher = DynamicPatcher(
+                self, cleanup_handlers=self.cleanup_handlers
+            )
             sys.meta_path.insert(0, self._dyn_patcher)
             for module in self.modules_to_reload:
                 if sys.modules.get(module.__name__) is module:
@@ -1106,11 +1127,16 @@ class DynamicPatcher(MetaPathFinder, Loader):
     Implements the protocol needed for import hooks.
     """
 
-    def __init__(self, patcher: Patcher) -> None:
+    def __init__(
+        self,
+        patcher: Patcher,
+        cleanup_handlers: Optional[Dict[str, Callable[[], bool]]] = None,
+    ) -> None:
         self._patcher = patcher
         self.sysmodules = {}
         self.modules = self._patcher.fake_modules
         self._loaded_module_names: Set[str] = set()
+        self.cleanup_handlers = cleanup_handlers or {}
 
         # remove all modules that have to be patched from `sys.modules`,
         # otherwise the find_... methods will not be called
@@ -1143,6 +1169,8 @@ def cleanup(self, cleanup_mode: ModuleCleanupMode) -> None:
                 cleanup_mode = ModuleCleanupMode.DELETE
         for name in self._loaded_module_names:
             if name in sys.modules and name not in reloaded_module_names:
+                if name in self.cleanup_handlers and self.cleanup_handlers[name]():
+                    continue
                 if cleanup_mode == ModuleCleanupMode.RELOAD:
                     try:
                         reload(sys.modules[name])

pyfakefs/patched_packages.py

@@ -18,10 +18,16 @@
 
 try:
     import pandas as pd
-    import pandas.io.parsers as parsers
+
+    try:
+        import pandas.io.parsers as parsers
+    except ImportError:
+        parsers = None
 except ImportError:
+    pd = None
     parsers = None
 
+
 try:
     import xlrd
 except ImportError:
@@ -57,6 +63,15 @@ def get_classes_to_patch():
     return classes_to_patch
 
 
+def get_cleanup_handlers():
+    handlers = {}
+    if pd is not None:
+        handlers["pandas.core.arrays.arrow.extension_types"] = (
+            handle_extension_type_cleanup
+        )
+    return handlers
+
+
 def get_fake_module_classes():
     fake_module_classes = {}
     if patch_pandas:
@@ -134,6 +149,23 @@ def __getattr__(self, name):
             return getattr(self._parsers_module, name)
 
 
+if pd is not None:
+
+    def handle_extension_type_cleanup():
+        # the module registers two extension types on load
+        # on reload it raises if the extensions have not been unregistered before
+        try:
+            import pyarrow
+
+            # the code to register these types has been in the module
+            # since it was created (in pandas 1.5)
+            pyarrow.unregister_extension_type("pandas.interval")
+            pyarrow.unregister_extension_type("pandas.period")
+        except ImportError:
+            pass
+        return False
+
+
 if locks is not None:
 
     class FakeLocks:

pyfakefs/pytest_tests/data/test.parquet

@@ -0,0 +1,18 @@
+"""Regression test for #947.
+Ensures that reloading the `pandas.core.arrays.arrow.extension_types` module succeeds.
+"""
+
+from pathlib import Path
+
+import pandas as pd
+
+
+def test_1(fs):
+    dir_ = Path(__file__).parent / "data"
+    fs.add_real_directory(dir_)
+    pd.read_parquet(dir_ / "test.parquet")
+
+
+def test_2():
+    dir_ = Path(__file__).parent / "data"
+    pd.read_parquet(dir_ / "test.parquet")