FileSet.map_nested and the staged_out_dir extras wrapper

fileformats/core/__init__.py

@@ -3,7 +3,13 @@
 from .convertible import ConvertibleTo, ConvertibleToMarker
 from .datatype import DataType, FieldPrimitive
 from .decorators import mtime_cached_property, validated_property
-from .extras import converter, extra, extra_implementation, find_extra_implementation
+from .extras import (
+    converter,
+    extra,
+    extra_implementation,
+    find_extra_implementation,
+    staged_out_dir,
+)
 from .field import Field
 from .fileset import FileSet, FileSetMetadata, FileSetPrimitive
 from .identification import find_matching, from_mime, from_paths, to_mime
@@ -12,9 +18,7 @@
 from .sampling import SampleFileGenerator
 
 __all__ = [
-    "__version__",
     "Classifier",
-    "converter",
     "ConvertibleTo",
     "ConvertibleToMarker",
     "DataType",
@@ -23,18 +27,21 @@
     "FileSet",
     "FileSetMetadata",
     "FileSetPrimitive",
-    "find_extra_implementation",
-    "find_matching",
-    "from_mime",
-    "from_paths",
     "Loaded",
     "LoadedMarker",
     "MockMixin",
-    "mtime_cached_property",
     "SampleFileGenerator",
-    "to_mime",
-    "validated_property",
+    "__version__",
     "check_loaded",
+    "converter",
     "extra",
     "extra_implementation",
+    "find_extra_implementation",
+    "find_matching",
+    "from_mime",
+    "from_paths",
+    "mtime_cached_property",
+    "staged_out_dir",
+    "to_mime",
+    "validated_property",
 ]

fileformats/core/extras.py

@@ -2,10 +2,14 @@
 import importlib
 import inspect
 import logging
+import os
+import shutil
+import tempfile
 import types
 import typing as ty
 import urllib.error
 from itertools import zip_longest
+from pathlib import Path
 from typing import Annotated, Self
 
 import fileformats.core
@@ -116,6 +120,56 @@ def decorated_extra(obj: DataType, *args: ty.Any, **kwargs: ty.Any) -> ty.Any:
     return decorated_extra  # type: ignore[return-value]
 
 
+def staged_out_dir(
+    implementation: ty.Callable[..., ty.Any],
+    obj: ty.Any,
+    out_dir: "os.PathLike[str] | str",
+    *args: ty.Any,
+    **kwargs: ty.Any,
+) -> ty.Any:
+    """A wrapper for extra methods that write a fileset to the `out_dir` passed as
+    their first argument (e.g. ``@extra(wrapper=staged_out_dir)``), which passes the
+    implementations a hidden staging directory alongside `out_dir` instead, and only
+    moves its contents into `out_dir` once they have returned. That way, only complete
+    output ends up in `out_dir`, and implementations are free to copy files there and
+    modify them in place.
+
+    The implementations need to return a fileset within the `out_dir` they are passed,
+    which is returned relocated to the requested `out_dir`. Any existing files in
+    `out_dir` are left in place, but a FileExistsError is raised, before anything is
+    moved, if the output would overwrite any of them.
+    """
+    out_dir = Path(out_dir).absolute()
+    out_dir.parent.mkdir(parents=True, exist_ok=True)
+    staging_dir = Path(
+        tempfile.mkdtemp(dir=out_dir.parent, prefix=f".staging-{out_dir.name}-")
+    )
+    try:
+        output = implementation(obj, staging_dir, *args, **kwargs)
+        try:
+            rel_fspaths = [
+                p.absolute().relative_to(staging_dir) for p in output.fspaths
+            ]
+        except ValueError:
+            raise RuntimeError(
+                f"The {getattr(implementation, '__name__', implementation)!r} "
+                f"implementation for {type(obj)} didn't write its output to the "
+                f"'out_dir' it was passed: {output.fspaths}"
+            ) from None
+        entries = sorted(staging_dir.iterdir())
+        if existing := [e.name for e in entries if (out_dir / e.name).exists()]:
+            raise FileExistsError(
+                f"Cannot write the output for {obj} into {out_dir} as {existing} "
+                "already exist(s) there"
+            )
+        out_dir.mkdir(exist_ok=True)
+        for entry in entries:
+            os.replace(entry, out_dir / entry.name)
+    finally:
+        shutil.rmtree(staging_dir, ignore_errors=True)
+    return type(output)([out_dir / p for p in rel_fspaths])
+
+
 def _import_base_extras_modules(datatype: ty.Type[DataType]) -> None:
     """Import the extras modules of the base classes of `datatype`, which may contain
     implementations it inherits (e.g. if it subclasses a format from another namespace

fileformats/core/fileset.py

@@ -578,22 +578,103 @@ def required_paths(self) -> ty.FrozenSet[Path]:
                     required.add(path)
         return frozenset(required)
 
-    def nested_filesets(self) -> ty.List["FileSet"]:
-        """Returns all nested filesets that are required for the format
+    def nested_filesets(self, recursive: bool = True) -> ty.List["FileSet"]:
+        """Returns all nested filesets that are required for the format, i.e. those
+        returned by its validated properties, including those within dicts and
+        sequences returned by them
+
+        Parameters
+        ----------
+        recursive : bool, optional
+            whether to also include the filesets nested within the nested filesets, by
+            default True
 
         Returns
         ------
         fileset : list[FileSet]
             a fileset that is nested within the broader fileset
         """
-        nested = []
+        nested: ty.List[FileSet] = []
+
+        def collect(value: ty.Any) -> None:
+            if isinstance(value, FileSet):
+                nested.append(value)
+                if recursive:
+                    nested.extend(value.nested_filesets())
+            elif isinstance(value, ty.Mapping):
+                for v in value.values():
+                    collect(v)
+            elif isinstance(value, (list, tuple, set, frozenset)):
+                for v in value:
+                    collect(v)
+
         for prop_name in sorted(self.validated_properties()):
-            prop = getattr(self, prop_name)
-            if isinstance(prop, FileSet):
-                nested.append(prop)
-                nested.extend(prop.nested_filesets())
+            collect(getattr(self, prop_name))
         return nested
 
+    def map_nested(
+        self,
+        out_dir: "os.PathLike[str] | str",
+        func: ty.Callable[["FileSet", Path], ty.Optional["FileSet"]],
+    ) -> Self:
+        """Writes a copy of the fileset to `out_dir`, in which the filesets nested
+        directly within it (see `nested_filesets`) are written by `func` and its
+        remaining files are copied unchanged, keeping the same layout relative to their
+        common parent directory. For example, to write converted or otherwise modified
+        versions of the nested filesets.
+
+        Nested filesets that point outside of the fileset (i.e. aren't made up of its
+        paths) aren't part of it, so are left out of the copy, as they are by `copy`
+
+        Parameters
+        ----------
+        out_dir : PathLike[str] | str
+            the directory to write the copy to
+        func : Callable[[FileSet, Path], FileSet | None]
+            called with each nested fileset and the directory it should be written to,
+            to be in the same position relative to the copy as the nested fileset is
+            to this fileset. It should write its counterpart there with the same
+            layout (e.g. names, so that the copy is still of this fileset's type) and
+            return it, or return None to have the nested fileset copied unchanged along
+            with the rest of the files. Nested filesets that are entirely within ones
+            already written by `func` aren't passed to it.
+
+        Returns
+        -------
+        Self
+            the copy of the fileset
+
+        Raises
+        ------
+        ValueError
+            if a nested fileset is made up of paths both within and outside the fileset
+        """
+        out_dir = Path(out_dir)
+        src_root = _root_dir(self)
+        own_paths = set(self.all_file_paths())
+        written: ty.Set[Path] = set()
+        for nested in self.nested_filesets(recursive=False):
+            nested_paths = set(nested.all_file_paths())
+            if outside := sorted(nested_paths - own_paths):
+                if outside == sorted(nested_paths):
+                    continue  # not part of the fileset
+                raise ValueError(
+                    f"{nested} is nested within {self}, but is made up of paths both "
+                    f"within it and outside it, {outside}"
+                )
+            if nested_paths <= written:
+                continue
+            dest_dir = out_dir / _root_dir(nested).relative_to(src_root)
+            if func(nested, dest_dir) is not None:
+                written.update(nested_paths)
+        for path in own_paths - written:
+            dest = out_dir / path.relative_to(src_root)
+            dest.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(path, dest)
+        # Reinstantiating the type checks that `func` preserved the layout it requires
+        # (e.g. didn't rename files it looks for)
+        return type(self)([out_dir / p.relative_to(src_root) for p in self.fspaths])
+
     def trim_paths(self) -> Self:
         """Trims paths in fspaths to only those that are "required" by the format class
         i.e. returned by a required property"""
@@ -2309,3 +2390,9 @@ def _destination_to_avoid(
     _formats_by_name: ty.Optional[ty.Dict[str, ty.Set[ty.Type["FileSet"]]]] = None
     _required_props: ty.Optional[ty.Tuple[str, ...]] = None
     _valid_class: ty.Optional[bool] = None
+
+
+def _root_dir(fileset: FileSet) -> Path:
+    """The directory containing all of the fileset's paths (the parent of a directory
+    fileset rather than the directory itself)"""
+    return Path(os.path.commonpath([p.parent for p in fileset.fspaths]))