RF: Bring in fMRIPrep's data loader

Author: mgxd

nibabies/cli/run.py

@@ -130,8 +130,6 @@ def main():
                 _copy_any(dseg_tsv, str(config.execution.nibabies_dir / "desc-aparcaseg_dseg.tsv"))
         # errno = 0
         finally:
-            from pkg_resources import resource_filename as pkgrf
-
             from ..reports.core import generate_reports
 
             # Generate reports phase

nibabies/cli/workflow.py

@@ -44,8 +44,6 @@ def build_workflow(config_file):
 
     # Called with reports only
     if config.execution.reports_only:
-        from pkg_resources import resource_filename as pkgrf
-
         build_logger.log(
             25,
             "Running --reports-only on participants %s",
@@ -132,14 +130,16 @@ def build_boilerplate(workflow):
         from shutil import copyfile
         from subprocess import CalledProcessError, TimeoutExpired, check_call
 
-        from pkg_resources import resource_filename as pkgrf
+        from nibabies.data import load as load_data
+
+        bib = load_data.readable("boilerplate.bib").read_text()
 
         # Generate HTML file resolving citations
         cmd = [
             "pandoc",
             "-s",
             "--bibliography",
-            pkgrf("nibabies", "data/boilerplate.bib"),
+            bib,
             "--citeproc",
             "--metadata",
             'pagetitle="nibabies citation boilerplate"',
@@ -159,7 +159,7 @@ def build_boilerplate(workflow):
             "pandoc",
             "-s",
             "--bibliography",
-            pkgrf("nibabies", "data/boilerplate.bib"),
+            bib,
             "--natbib",
             str(citation_files["md"]),
             "-o",
@@ -171,4 +171,4 @@ def build_boilerplate(workflow):
         except (FileNotFoundError, CalledProcessError, TimeoutExpired):
             config.loggers.cli.warning("Could not generate CITATION.tex file:\n%s", " ".join(cmd))
         else:
-            copyfile(pkgrf("nibabies", "data/boilerplate.bib"), citation_files["bib"])
+            copyfile(bib, citation_files["bib"])

nibabies/conftest.py

@@ -4,7 +4,8 @@
 from tempfile import TemporaryDirectory
 
 import pytest
-from pkg_resources import resource_filename
+
+from nibabies.data import load as load_data
 
 FILES = (
     "functional.nii",
@@ -34,5 +35,5 @@ def data_dir():
 @pytest.fixture(autouse=True)
 def set_namespace(doctest_namespace, data_dir):
     doctest_namespace["data_dir"] = data_dir
-    doctest_namespace["test_data"] = Path(resource_filename("nibabies", "tests/data"))
+    doctest_namespace["test_data"] = load_data.cached('../tests/data')
     doctest_namespace["Path"] = Path

nibabies/data/__init__.py

@@ -1,6 +1,25 @@
+"""Data file retrieval
+
+.. autofunction:: load
+
+.. automethod:: load.readable
+
+.. automethod:: load.as_path
+
+.. automethod:: load.cached
+
+.. autoclass:: Loader
+"""
+
+from __future__ import annotations
+
 import atexit
-from contextlib import ExitStack
+import os
+from contextlib import AbstractContextManager, ExitStack
+from functools import cached_property
 from pathlib import Path
+from types import ModuleType
+from typing import Union
 
 try:
     from functools import cache
@@ -10,16 +29,155 @@
 try:  # Prefer backport to leave consistency to dependency spec
     from importlib_resources import as_file, files
 except ImportError:
-    from importlib.resources import as_file, files
+    from importlib.resources import as_file, files  # type: ignore
+
+try:  # Prefer stdlib so Sphinx can link to authoritative documentation
+    from importlib.resources.abc import Traversable
+except ImportError:
+    from importlib_resources.abc import Traversable
+
+__all__ = ["load"]
+
+
+class Loader:
+    """A loader for package files relative to a module
+
+    This class wraps :mod:`importlib.resources` to provide a getter
+    function with an interpreter-lifetime scope. For typical packages
+    it simply passes through filesystem paths as :class:`~pathlib.Path`
+    objects. For zipped distributions, it will unpack the files into
+    a temporary directory that is cleaned up on interpreter exit.
+
+    This loader accepts a fully-qualified module name or a module
+    object.
+
+    Expected usage::
+
+        '''Data package
+
+        .. autofunction:: load_data
+
+        .. automethod:: load_data.readable
+
+        .. automethod:: load_data.as_path
+
+        .. automethod:: load_data.cached
+        '''
+
+        from nibabies.data import Loader
+
+        load_data = Loader(__package__)
+
+    :class:`~Loader` objects implement the :func:`callable` interface
+    and generate a docstring, and are intended to be treated and documented
+    as functions.
+
+    For greater flexibility and improved readability over the ``importlib.resources``
+    interface, explicit methods are provided to access resources.
+
+    +---------------+----------------+------------------+
+    | On-filesystem | Lifetime       | Method           |
+    +---------------+----------------+------------------+
+    | `True`        | Interpreter    | :meth:`cached`   |
+    +---------------+----------------+------------------+
+    | `True`        | `with` context | :meth:`as_path`  |
+    +---------------+----------------+------------------+
+    | `False`       | n/a            | :meth:`readable` |
+    +---------------+----------------+------------------+
+
+    It is also possible to use ``Loader`` directly::
+
+        from nibabies.data import Loader
+
+        Loader(other_package).readable('data/resource.ext').read_text()
+
+        with Loader(other_package).as_path('data') as pkgdata:
+            # Call function that requires full Path implementation
+            func(pkgdata)
+
+        # contrast to
+
+        from importlib_resources import files, as_file
+
+        files(other_package).joinpath('data/resource.ext').read_text()
+
+        with as_file(files(other_package) / 'data') as pkgdata:
+            func(pkgdata)
+
+    .. automethod:: readable
+
+    .. automethod:: as_path
+
+    .. automethod:: cached
+    """
+
+    def __init__(self, anchor: Union[str, ModuleType]):
+        self._anchor = anchor
+        self.files = files(anchor)
+        self.exit_stack = ExitStack()
+        atexit.register(self.exit_stack.close)
+        # Allow class to have a different docstring from instances
+        self.__doc__ = self._doc
+
+    @cached_property
+    def _doc(self):
+        """Construct docstring for instances
+
+        Lists the public top-level paths inside the location, where
+        non-public means has a `.` or `_` prefix or is a 'tests'
+        directory.
+        """
+        top_level = sorted(
+            os.path.relpath(p, self.files) + "/"[: p.is_dir()]
+            for p in self.files.iterdir()
+            if p.name[0] not in (".", "_") and p.name != "tests"
+        )
+        doclines = [
+            f"Load package files relative to ``{self._anchor}``.",
+            "",
+            "This package contains the following (top-level) files/directories:",
+            "",
+            *(f"* ``{path}``" for path in top_level),
+        ]
+
+        return "\n".join(doclines)
+
+    def readable(self, *segments) -> Traversable:
+        """Provide read access to a resource through a Path-like interface.
+
+        This file may or may not exist on the filesystem, and may be
+        efficiently used for read operations, including directory traversal.
+
+        This result is not cached or copied to the filesystem in cases where
+        that would be necessary.
+        """
+        return self.files.joinpath(*segments)
+
+    def as_path(self, *segments) -> AbstractContextManager[Path]:
+        """Ensure data is available as a :class:`~pathlib.Path`.
+
+        This method generates a context manager that yields a Path when
+        entered.
+
+        This result is not cached, and any temporary files that are created
+        are deleted when the context is exited.
+        """
+        return as_file(self.files.joinpath(*segments))
+
+    @cache
+    def cached(self, *segments) -> Path:
+        """Ensure data is available as a :class:`~pathlib.Path`.
 
-__all__ = ["load_resource"]
+        Any temporary files that are created remain available throughout
+        the duration of the program, and are deleted when Python exits.
 
-exit_stack = ExitStack()
-atexit.register(exit_stack.close)
+        Results are cached so that multiple calls do not unpack the same
+        data multiple times, but the cache is sensitive to the specific
+        argument(s) passed.
+        """
+        return self.exit_stack.enter_context(as_file(self.files.joinpath(*segments)))
 
-path = files(__package__)
+    __call__ = cached
 
 
-@cache
-def load_resource(fname: str) -> Path:
-    return exit_stack.enter_context(as_file(path.joinpath(fname)))
+load = Loader(__package__)

nibabies/reports/core.py

@@ -2,7 +2,7 @@
 
 from nireports.assembler.report import Report
 
-from nibabies.data import load_resource
+from nibabies.data import load as load_data
 
 
 def run_reports(
@@ -22,7 +22,7 @@ def run_reports(
         run_uuid,
         subject=subject,
         session=session,
-        bootstrap_file=load_resource('reports-spec.yml'),
+        bootstrap_file=load_data.readable('reports-spec.yml'),
         reportlets_dir=reportlets_dir,
     ).generate_report()
 

nibabies/tests/test_config.py

@@ -27,10 +27,10 @@
 
 import pytest
 from niworkflows.utils.spaces import format_reference
-from pkg_resources import resource_filename as pkgrf
 from toml import loads
 
-from .. import config
+from nibabies import config
+from nibabies.data import load as load_data
 
 
 def _reset_config():
@@ -58,8 +58,7 @@ def test_reset_config():
 
 def test_config_spaces():
     """Check that all necessary spaces are recorded in the config."""
-    filename = Path(pkgrf('nibabies', 'data/tests/config.toml'))
-    settings = loads(filename.read_text())
+    settings = loads(load_data.readable('tests/config.toml').read_text())
     for sectionname, configs in settings.items():
         if sectionname != 'environment':
             section = getattr(config, sectionname)

nibabies/utils/misc.py

@@ -8,7 +8,7 @@
 from typing import Union
 
 from nibabies import __version__
-from nibabies.data import load_resource
+from nibabies.data import load as load_data
 
 
 def fix_multi_source_name(in_files):
@@ -144,7 +144,9 @@ def save_fsLR_mcribs(mcribs_dir: str | Path) -> None:
     template_dir = Path(mcribs_dir) / 'templates_fsLR'
     template_dir.mkdir(exist_ok=True)
 
-    for src in load_resource('atlases').glob('*sphere.surf.gii'):
+    atlases = load_data.cached('atlases')
+
+    for src in atlases.glob('*sphere.surf.gii'):
         if not (dst := (template_dir / src.name)).exists():
             try:
                 shutil.copyfile(src, dst)

nibabies/workflows/anatomical/resampling.py

@@ -8,7 +8,7 @@
 from smriprep.workflows.surfaces import _collate, init_morph_grayords_wf
 
 from nibabies.config import DEFAULT_MEMORY_MIN_GB
-from nibabies.data import load_resource
+from nibabies.data import load as load_data
 from nibabies.interfaces.utils import CiftiSelect
 
 
@@ -59,7 +59,7 @@ def init_anat_fsLR_resampling_wf(
     select_surfaces = pe.Node(CiftiSelect(), name='select_surfaces')
 
     if mcribs:
-        atlases = load_resource('atlases')
+        atlases = load_data.cached('atlases')
         # use dHCP 32k fsLR instead
         select_surfaces.inputs.template_spheres = [
             str(atlases / 'tpl-dHCP_space-fsLR_hemi-L_den-32k_desc-week42_sphere.surf.gii'),
@@ -240,7 +240,7 @@ def init_mcribs_morph_grayords_wf(
         ],
     )
 
-    atlases = load_resource('atlases')
+    atlases = load_data.cached('atlases')
     resample.inputs.new_sphere = [  # 32k
         str(atlases / 'tpl-dHCP_space-fsLR_hemi-L_den-32k_desc-week42_sphere.surf.gii'),
         str(atlases / 'tpl-dHCP_space-fsLR_hemi-R_den-32k_desc-week42_sphere.surf.gii'),

nibabies/workflows/anatomical/surfaces.py

@@ -12,8 +12,8 @@
 from niworkflows.utils.connections import pop_file
 from smriprep.workflows.surfaces import init_gifti_surface_wf
 
-from ...config import DEFAULT_MEMORY_MIN_GB
-from ...data import load_resource
+from nibabies.config import DEFAULT_MEMORY_MIN_GB
+from nibabies.data import load as load_data
 
 SURFACE_INPUTS = [
     "subjects_dir",
@@ -239,7 +239,7 @@ def init_mcribs_sphere_reg_wf(*, name="mcribs_sphere_reg_wf"):
     fix_meta = pe.MapNode(FixGiftiMetadata(), iterfield="in_file", name="fix_meta")
 
     # load template files
-    atlases = load_resource('atlases')
+    atlases = load_data.cached('atlases')
 
     # SurfaceSphereProjectUnProject
     # project to 41k dHCP atlas sphere

nibabies/workflows/bold/resampling.py

@@ -24,7 +24,7 @@
 from niworkflows.interfaces.workbench import MetricDilate, MetricMask, MetricResample
 
 from nibabies.config import DEFAULT_MEMORY_MIN_GB
-from nibabies.data import load_resource
+from nibabies.data import load as load_data
 
 if ty.TYPE_CHECKING:
     from niworkflows.utils.spaces import SpatialReferences
@@ -589,7 +589,7 @@ def init_bold_fsLR_resampling_wf(
     # select white, midthickness and pial surfaces based on hemi
     select_surfaces = pe.Node(CiftiSelect(), name='select_surfaces')
     if mcribs:
-        atlases = load_resource('atlases')
+        atlases = load_data.cached('atlases')
         # use dHCP 32k fsLR instead
         select_surfaces.inputs.template_spheres = [
             str(atlases / 'tpl-dHCP_space-fsLR_hemi-L_den-32k_desc-week42_sphere.surf.gii'),