Another try; now when pickling we specify a directory wrt we store all absolute filepaths

Author: elinscott

src/koopmans/calculators/_calculator.py

@@ -164,6 +164,14 @@ def absolute_directory(self) -> Path:
             raise AttributeError(f'Expected `{obj.__class__.__name__}` instance to have a `base_directory` attribute')
         return obj.base_directory / path
 
+    @property
+    def base_directory(self) -> Path:
+        if self.parent is not None:
+            return self.parent.base_directory
+        else:
+            raise ValueError(
+                f'{self.__class__.__name__} does not have a parent, so `base_directory` cannot be accessed')
+
     def calculate(self):
         """Generic function for running a calculator"""
 

src/koopmans/files.py

@@ -72,43 +72,87 @@ def __eq__(self, other):
 
     def __reduce__(self):
         # We don't want to store the entire parent object in the database; we only need the directory information
-        abs_dir = self.parent.absolute_directory
-        dummy_parent = ParentPlaceholder(directory=self.parent.directory,
-                                         absolute_directory=self.parent.absolute_directory)
+        dummy_parent = ParentPlaceholder.fromobj(self.parent)
         return (FilePointer, (dummy_parent, self.name))
 
 
 class ParentPlaceholder:
-    # Move into the test suite?
-    def __init__(self, directory, absolute_directory):
+    # Placeholder parent for FilePointers that don't have a Workflow/Process/Calculator as a parent
+    def __init__(self, parent, directory, base_directory=None):
+        self.parent = parent
         self.directory = directory
-        self._absolute_directory = absolute_directory
+        if self.parent is None:
+            self.base_directory = base_directory
+        else:
+            self._base_directory = None
 
-    @property
-    def absolute_directory(self):
-        # This is a property in order to follow the HasDirectoryInfo protocol
-        return Path(__file__).parents[2] / self._absolute_directory
+    def __repr__(self):
+        return f'ParentPlaceholder(directory={self.absolute_directory})'
 
+    @classmethod
+    def fromobj(cls, obj, replace_parents_with_placeholders=False):
+        if replace_parents_with_placeholders:
+            if obj.parent is None:
+                parent = None
+            else:
+                parent = cls.fromobj(obj, replace_parents_with_placeholders=True)
+        else:
+            parent = obj.parent
 
-class AbsolutePath:
-    """ A class that can stand in as a parent in a FilePointer when the file is unattached to a Calculator or Process """
+        if parent is None:
+            base_directory = obj.base_directory
+        else:
+            base_directory = None
+        directory = obj.directory
+        new_obj = cls(parent, directory, base_directory)
 
-    def __init__(self, directory: Path | str | None):
-        directory = Path(directory) if isinstance(directory, str) else directory
-        self.directory: Path | None = directory
+        # Sanity checks
+        assert new_obj.directory == obj.directory
+        assert new_obj.absolute_directory == obj.absolute_directory
+        assert new_obj.base_directory == obj.base_directory
 
-    def __eq__(self, other):
-        if not isinstance(other, AbsolutePath):
-            return False
-        return self.directory == other.directory
+        return new_obj
+
+    @classmethod
+    def frompath(cls, path: Path):
+        return cls(None, Path(), path)
 
     @property
-    def absolute_directory(self) -> Path | None:
-        return self.directory
+    def base_directory(self) -> Path:
+        if self.parent is not None:
+            return self.parent.base_directory
+        else:
+            if self._base_directory is None:
+                raise ValueError(f'{self.__class__.__name__}.base_directory has not been set')
+            return self._base_directory
+
+    @base_directory.setter
+    def base_directory(self, value: Path):
+        if self.parent is not None:
+            raise ValueError(f'{self.__class__.__name__}.base_directory should not be set for processes with parents')
+        self._base_directory = value.resolve()
+
+    @property
+    def absolute_directory(self) -> Path:
+        assert self.directory is not None
+        if self.parent is None:
+            abs_dir = Path(self.base_directory / self.directory)
+            assert abs_dir.is_absolute()
+            return abs_dir
+        path = self.parent.directory / self.directory
+
+        # Recursive through the parents, adding their directories to path (these are all relative paths)...
+        obj = self.parent
+        while getattr(obj, 'parent', None):
+            assert obj.parent is not None
+            path = obj.parent.directory / path
+            obj = obj.parent
+
+        # Finally, 'path' is relative to self.base_directory
+        return self.base_directory / path
 
 
 def AbsoluteFilePointer(path: Path | str) -> FilePointer:
     path = path if isinstance(path, Path) else Path(path)
-    path = path.resolve()
-    parent = AbsolutePath(directory=path.parent)
+    parent = ParentPlaceholder.frompath(path.parent)
     return FilePointer(parent=parent, name=Path(path.name))

src/koopmans/io/_dill.py

@@ -1,45 +1,68 @@
 import os
 from pathlib import Path
-from typing import Any
+from typing import Any, Optional
 
 import dill
 
+BASE_PLACEHOLDER = '/BASE_DIRECTORY_PLACEHOLDER/'
 
-class PicklerWithRelativePaths(dill.Pickler):
-    # A pickler that converts absolute paths to relative paths, used during the test suite in combination
-    # with a monkeypatching of the Path.__eq__ method
+
+class CustomPicklerUsePlaceholder(dill.Pickler):
     def __init__(self, base_directory, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.base_directory = base_directory
 
-    def reducer_override(self, obj):
-        if isinstance(obj, Path):
-            return Path, pickle_path(obj, self.base_directory)
-        else:
-            return NotImplemented
+    def persistent_id(self, obj):
+        """Called when pickling objects. If obj is a Path, return a unique identifier."""
+        if isinstance(obj, Path) and obj.is_absolute():
+            if self.base_directory == obj:
+                return ('Path', BASE_PLACEHOLDER)
+            # Return a tuple to identify this as a Path that needs special handling
+            tup = ('Path', str((Path(BASE_PLACEHOLDER) / os.path.relpath(obj, self.base_directory))))
+            if 'data' in tup[1]:
+                assert '..' in tup[1]
+            return tup
+        return None  # No special handling required for other objects
+
+
+class CustomUnpicklerReplacePlaceholder(dill.Unpickler):
+    def __init__(self, base_directory, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.base_directory = base_directory.resolve()
+
+    def persistent_load(self, pid):
+        """Called when unpickling to resolve persistent ids."""
+        if isinstance(pid, tuple) and pid[0] == 'Path' and pid[1].startswith(BASE_PLACEHOLDER):
+            # Replace the placeholder with the actual current working directory
+            path = (self.base_directory / pid[1][len(BASE_PLACEHOLDER):]).resolve()
+            if 'data' in pid[1]:
+                assert 'tmp' not in str(path)
+            return path
+        return pid  # If it's not a special persistent id, return as is
 
 
-def pickle_path(path: Path, base_directory: Path):
-    # Convert the path to relative if it's absolute
-    base_directory = base_directory.resolve()
-    if path.is_absolute():
-        rel_path = os.path.relpath(path, base_directory)
-        return (str(rel_path),)
-    else:
-        # Return relative paths as-is
-        return (str(path),)
+class CustomUnpicklerKeepPlaceholder(dill.Unpickler):
+    def persistent_load(self, pid):
+        if isinstance(pid, tuple) and pid[0] == 'Path':
+            return Path(pid[1])
 
 
-def read_pkl(filename: Path | str) -> Any:
+def read_pkl(filename: Path | str, base_directory: Optional[Path] = None) -> Any:
     with open(filename, 'rb') as fd:
-        out = dill.load(fd)
+        if base_directory is None:
+            unpickler = CustomUnpicklerKeepPlaceholder(fd)
+        else:
+            unpickler = CustomUnpicklerReplacePlaceholder(base_directory, fd)
+        out = unpickler.load()
+
     return out
 
 
 def write_pkl(obj: Any, filename: Path | str, base_directory: Path | None = None):
     filename = Path(filename) if not isinstance(filename, Path) else filename
-    if base_directory is None:
-        base_directory = filename.parent
     with open(filename, 'wb') as fd:
-        pickler = PicklerWithRelativePaths(base_directory, fd)
+        if base_directory is None:
+            pickler = dill.Pickler(fd)
+        else:
+            pickler = CustomPicklerUsePlaceholder(base_directory, fd)
         pickler.dump(obj)

src/koopmans/processes/_process.py

@@ -140,7 +140,9 @@ def base_directory(self, value: Path):
     def absolute_directory(self) -> Path:
         assert self.directory is not None
         if self.parent is None:
-            return self.base_directory / self.directory
+            abs_dir = Path(self.base_directory / self.directory)
+            assert abs_dir.is_absolute()
+            return abs_dir
         path = self.parent.directory / self.directory
 
         # Recursive through the parents, adding their directories to path (these are all relative paths)...
@@ -150,7 +152,5 @@ def absolute_directory(self) -> Path:
             path = obj.parent.directory / path
             obj = obj.parent
 
-        # ... until we reach the top-level parent, which should have a base_directory attribute (an absolute path)
-        if not hasattr(obj, 'base_directory'):
-            raise AttributeError(f'Expected `{obj.__class__.__name__}` instance to have a `base_directory` attribute')
-        return obj.base_directory / path
+        # Finally, 'path' is relative to self.base_directory
+        return self.base_directory / path

tests/benchmarks/test_compute_power_spectrum_pr0/compute_power_spectrum.pkl

@@ -29,7 +29,7 @@ def test_read_dynG(tio2, tmp_path, datadir, pytestconfig):
 
         if pytestconfig.getoption('generate_benchmark'):
             # Write the dielectric tensor to file
-            write_pkl(eps, benchmark_filename(calc), base_directory=Path(__file__).parents[2])
+            write_pkl(eps, benchmark_filename(calc))
         else:
             # Compare with the dielectric tensor on file
             eps_ref = read_pkl(benchmark_filename(calc))

tests/benchmarks/test_convert_files_from_spin_10/convert_files_from_spin1to2.pkl

@@ -38,5 +38,5 @@ def test_generate_dos(silicon, tmp_path, datadir, pytestconfig):
             write_pkl(dos, benchmark_filename(calc))
         else:
             # Compare with the DOS on file
-            dos_ref = read_pkl(benchmark_filename(calc), base_directory=Path(__file__).parents[2])
+            dos_ref = read_pkl(benchmark_filename(calc))
             assert dos == dos_ref

tests/benchmarks/test_convert_files_from_spin_20/convert_files_from_spin2to1.pkl

@@ -5,8 +5,7 @@
 import koopmans
 
 from .helpers.os import datadir, sys2file
-from .helpers.patches import (check_patch, espresso_patch,
-                              patch_path_comparison, tutorial_patch,
+from .helpers.patches import (check_patch, espresso_patch, tutorial_patch,
                               workflow_patch)
 from .helpers.strategies import ase_cells, bandpaths, kpoints
 from .helpers.systems import gaas, ozone, silicon, tio2, water, water_snapshots

tests/benchmarks/test_extract_coefficients_from0/extract_coefficients_from_xml.pkl

@@ -6,21 +6,6 @@
 from ._utils import benchmark_filename
 
 
-@pytest.fixture(autouse=True)
-def patch_path_comparison(monkeypatch):
-    def path_eq(self, other):
-        if not isinstance(other, type(self)):
-            return False
-        if self.is_absolute() == other.is_absolute():
-            return str(self) == str(other)
-        else:
-            # In some cases we store paths relative to a base directory. Without knowing
-            # the base directory, the best we can do is check that the relative path
-            # ends the same as the absolute path
-            return str(self.resolve()).endswith(str(other)) or str(other).endswith(str(self))
-    monkeypatch.setattr('pathlib.Path.__eq__', path_eq)
-
-
 def monkeypatch_stumble(monkeypatch):
     from ._stumble import (StumblingConvergenceWorkflow,
                            StumblingDeltaSCFWorkflow, StumblingDFTCPWorkflow,

tests/benchmarks/test_ui_si0/unfold_and_interpolate.pkl

@@ -45,7 +45,7 @@ def _calculate(self):
 
         self.parent, parent = None, self.parent
         self.linked_files, linked_files = [], self.linked_files
-        write_pkl(self, filename, base_directory=base_directory)
+        write_pkl(self, filename, self.base_directory)
         self.parent = parent
         self.linked_files = linked_files
 
@@ -91,29 +91,15 @@ def _run(self):
             filename.parent.mkdir(parents=True)
         # Temporarily wipe the parent attribute so that the entire workflow doesn't get pickled
         self.parent, parent = None, self.parent
-        write_pkl(self, filename, base_directory=base_directory)
+        write_pkl(self, filename, self.base_directory)
         self.parent = parent
 
         # Copy over all files that are outputs of the process that need to be read
         for filepath in recursively_find_files([o for _, o in self.outputs]):
             if filepath.name in ['power_spectrum.npy']:
                 shutil.copy(filepath, benchmark_filename(self).parent / filepath.name)
 
-    # Patching the absolute_directory property
-    unpatched_absolute_directory = p.absolute_directory
-
-    def absolute_directory(self) -> Path:
-        if self.parent is None:
-            # Because we wipe parents when storing benchmarks (see above), this prevents us from being able to construct
-            # an absolute directory to locate files. Usually, this would raise an error. For the purposes of the test suite,
-            # instead simply use the base directory of the repo
-            return Path().resolve().relative_to(base_directory)
-        else:
-            # Default behavior
-            return unpatched_absolute_directory.__get__(self)
-
     monkeypatch.setattr(p, '_run', _run)
-    monkeypatch.setattr(p, 'absolute_directory', property(absolute_directory))
 
 
 def monkeypatch_bench(monkeypatch):

tests/calculators/test_ph.py

@@ -341,7 +341,7 @@ def patch_process(p, monkeypatch):
 
     def _run(self):
         # Load the benchmark
-        bench_process = read_pkl(benchmark_filename(self))
+        bench_process = read_pkl(benchmark_filename(self), base_directory=self.base_directory)
 
         # Compare the inputs
         assert bench_process.inputs == self.inputs

tests/calculators/test_projwfc.py

@@ -31,7 +31,7 @@ def mock_calculator__calculate(self):
     with chdir(self.directory):
         # By moving into the directory where the calculation was run, we ensure when we read in the settings that
         # paths are interpreted relative to this particular working directory
-        calc = read_pkl(benchmark_filename(self))
+        calc = read_pkl(benchmark_filename(self), base_directory=self.directory)
 
     # Compare the settings
     for key in set(list(self.parameters.keys()) + list(calc.parameters.keys())):
@@ -129,7 +129,7 @@ def generate_dos(self):
 
 def mock_process_run(self):
     # Load the inputs from file
-    bench_process = read_pkl(benchmark_filename(self))
+    bench_process = read_pkl(benchmark_filename(self), base_directory=self.base_directory)
     assert self.inputs == bench_process.inputs
     self.outputs = bench_process.outputs
 

tests/conftest.py

@@ -23,8 +23,12 @@ def benchmark_filename(obj) -> Path:
     tests_dir = base_directory / 'tests'
     name = getattr(obj, 'prefix', getattr(obj, 'name', None))
     assert isinstance(name, str)
-    if tests_dir / 'tmp' in obj.absolute_directory.parents:
-        benchmark_name = obj.absolute_directory.relative_to(tests_dir / 'tmp') / name
+    abs_dir = obj.absolute_directory
+    if tests_dir / 'tmp' in abs_dir.parents:
+        try:
+            benchmark_name = abs_dir.relative_to(tests_dir / 'tmp') / name
+        except:
+            raise ValueError()
     else:
         # This is a tutorial test, which we run in the tutorials directory. Because of this, we don't have access
         # to the tmp_path automatically generated by pytest, so we need to construct something similar

tests/helpers/patches/__init__.py

@@ -0,0 +1,30 @@
+from pathlib import Path
+
+import pytest
+
+from koopmans.io import read_pkl, write_pkl
+from koopmans.utils import chdir
+
+base_directory = Path('/home/user/base_directory')
+paths = [base_directory / 'subpath/to/directory',
+         Path('/home/user/path/that/is/not/a/subpath/of/base_directory/'), Path('relative/path')]
+
+
+@pytest.mark.parametrize('obj', paths)
+def test_roundtrip_pkl_with_placeholder(obj, tmp_path):
+    with chdir(tmp_path):
+        write_pkl(obj, 'test.pkl', base_directory=base_directory)
+
+        new_obj = read_pkl('test.pkl', base_directory=base_directory)
+
+        assert obj == new_obj
+
+
+@pytest.mark.parametrize('obj', paths)
+def test_roundtrip_pkl_without_placeholder(obj, tmp_path):
+    with chdir(tmp_path):
+        write_pkl(obj, 'test.pkl')
+
+        new_obj = read_pkl('test.pkl')
+
+        assert obj == new_obj