junifer 0.0.7.dev332__py3-none-any.whl → 0.0.8__py3-none-any.whl

junifer/_version.py

@@ -1,5 +1,6 @@
-# file generated by setuptools-scm
+# file generated by vcs-versioning
 # don't change, don't track in version control
+from __future__ import annotations
 
 __all__ = [
     "__version__",
@@ -10,25 +11,14 @@ __all__ = [
     "commit_id",
 ]
 
-TYPE_CHECKING = False
-if TYPE_CHECKING:
-    from typing import Tuple
-    from typing import Union
-
-    VERSION_TUPLE = Tuple[Union[int, str], ...]
-    COMMIT_ID = Union[str, None]
-else:
-    VERSION_TUPLE = object
-    COMMIT_ID = object
-
 version: str
 __version__: str
-__version_tuple__: VERSION_TUPLE
-version_tuple: VERSION_TUPLE
-commit_id: COMMIT_ID
-__commit_id__: COMMIT_ID
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
 
-__version__ = version = '0.0.7.dev332'
-__version_tuple__ = version_tuple = (0, 0, 7, 'dev332')
+__version__ = version = '0.0.8'
+__version_tuple__ = version_tuple = (0, 0, 8)
 
-__commit_id__ = commit_id = None
+__commit_id__ = commit_id = 'g2b4aae01d'

junifer/api/__init__.pyi

@@ -1,4 +1,19 @@
-__all__ = ["decorators", "collect", "queue", "run", "reset", "list_elements"]
+__all__ = [
+    "decorators",
+    "collect",
+    "queue",
+    "run",
+    "reset",
+    "list_elements",
+    "parse_yaml",
+]
 
 from . import decorators
-from .functions import collect, list_elements, reset, run, queue
+from .functions import (
+    collect,
+    list_elements,
+    parse_yaml,
+    reset,
+    run,
+    queue,
+)

junifer/api/functions.py

@@ -5,10 +5,15 @@
 #          Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
+import atexit
+import importlib
+import importlib.util
 import os
 import shutil
+import sys
 from pathlib import Path
-from typing import Optional, Union
+
+import structlog
 
 from ..api.queue_context import GnuParallelLocalAdapter, HTCondorAdapter
 from ..datagrabber import BaseDataGrabber
@@ -27,10 +32,20 @@ from ..typing import (
     PreprocessorLike,
     StorageLike,
 )
-from ..utils import logger, raise_error, warn_with_log, yaml
+from ..utils import raise_error, warn_with_log, yaml
+
 
+__all__ = [
+    "collect",
+    "list_elements",
+    "parse_yaml",
+    "queue",
+    "reset",
+    "run",
+]
 
-__all__ = ["collect", "list_elements", "queue", "reset", "run"]
+_log = structlog.get_logger("junifer")
+logger = _log.bind(pkg="api")
 
 
 def _get_datagrabber(datagrabber_config: dict) -> DataGrabberLike:
@@ -121,13 +136,52 @@ def _get_storage(storage_config: dict) -> StorageLike:
     )
 
 
+def _get_workdir_params(workdir: str | Path | dict | None) -> dict:
+    """Get the parameters of the working directory manager.
+
+    Parameters
+    ----------
+    workdir : str or pathlib.Path or dict or None
+        The working directory as a path, or as a dictionary with the keys
+        ``path`` and ``cleanup`` (default True). If None, the default of
+        :class:`.WorkDirManager` is used.
+
+    Returns
+    -------
+    dict
+        The parameters for :class:`.WorkDirManager`.
+
+    """
+    if isinstance(workdir, dict):
+        return {
+            "workdir": Path(workdir["path"]),
+            "cleanup": workdir.get("cleanup", True),
+        }
+    if workdir is None:
+        return {"workdir": None, "cleanup": True}
+    return {"workdir": Path(workdir), "cleanup": True}
+
+
+def _setup_workdir(workdir: str | Path | dict | None) -> None:
+    """Set up the working directory manager.
+
+    Parameters
+    ----------
+    workdir : str or pathlib.Path or dict or None
+        The working directory (see :func:`_get_workdir_params`).
+
+    """
+    WorkDirManager(**_get_workdir_params(workdir))
+    atexit.register(WorkDirManager()._cleanup)
+
+
 def run(
-    workdir: Union[str, Path, dict],
+    workdir: str | Path | dict,
     datagrabber: dict,
     markers: list[dict],
     storage: dict,
-    preprocessors: Optional[list[dict]] = None,
-    elements: Optional[Elements] = None,
+    preprocessors: list[dict] | None = None,
+    elements: Elements | None = None,
 ) -> None:
     """Run the pipeline on the selected element.
 
@@ -165,24 +219,15 @@ def run(
         If invalid element selectors are found.
 
     """
-    # Conditional to handle workdir config
-    if isinstance(workdir, (str, Path)):
-        if isinstance(workdir, str):
-            workdir = {"workdir": Path(workdir), "cleanup": True}
-        else:
-            workdir = {"workdir": workdir, "cleanup": True}
-    elif isinstance(workdir, dict):
-        workdir["workdir"] = workdir.pop("path")
-
     # Initiate working directory manager with correct variation
-    if not workdir["cleanup"]:
+    if not _get_workdir_params(workdir)["cleanup"]:
         if elements is None or len(elements) > 1:
             raise_error(
                 "Cannot disable `workdir.cleanup` as "
                 f"{len(elements) if elements is not None else 'all'} "
                 "elements will be processed"
             )
-    WorkDirManager(**workdir)
+    _setup_workdir(workdir)
 
     # Get datagrabber to use
     datagrabber_object = _get_datagrabber(datagrabber.copy())
@@ -216,28 +261,15 @@ def run(
     # Fit elements
     with datagrabber_object:
         if elements is not None:
-            # Keep track of valid selectors
-            valid_elements = []
-            for t_element in datagrabber_object.filter(elements):
-                valid_elements.append(t_element)
-                mc.fit(datagrabber_object[t_element])
-            # Compute invalid selectors
-            invalid_elements = set(elements) - set(valid_elements)
-            # Report if invalid selectors are found
-            if invalid_elements:
-                raise_error(
-                    msg=(
-                        "The following element selectors are invalid:\n"
-                        f"{invalid_elements}"
-                    ),
-                    klass=RuntimeError,
-                )
+            # Only get the elements of the dataset for partial selectors
+            to_fit = datagrabber_object.select_elements(elements)
         else:
-            for t_element in datagrabber_object:
-                mc.fit(datagrabber_object[t_element])
+            to_fit = datagrabber_object.get_elements()
+        for t_element in to_fit:
+            mc.fit(datagrabber_object[t_element])
 
 
-def collect(storage: dict) -> None:
+def collect(storage: dict, workdir: str | Path | dict | None = None) -> None:
     """Collect and store data.
 
     Parameters
@@ -246,8 +278,12 @@ def collect(storage: dict) -> None:
         Storage to use. Must have a key ``kind`` with the kind of
         storage to use. All other keys are passed to the storage
         constructor.
+    workdir : str or pathlib.Path or dict or None, optional
+        The working directory, as in :func:`.run`. If None, the default of
+        :class:`.WorkDirManager` is used (default None).
 
     """
+    _setup_workdir(workdir)
     logger.info(f"Collecting data using {storage['kind']}")
     logger.debug(f"\tStorage params: {storage}")
     if "single_output" not in storage:
@@ -263,8 +299,8 @@ def queue(
     kind: str,
     jobname: str = "junifer_job",
     overwrite: bool = False,
-    elements: Optional[Elements] = None,
-    **kwargs: Union[str, int, bool, dict, tuple, list],
+    elements: Elements | None = None,
+    **kwargs: str | int | bool | dict | tuple | list,
 ) -> None:
     """Queue a job to be executed later.
 
@@ -291,6 +327,7 @@ def queue(
         if the ``jobdir`` exists and ``overwrite = False``.
 
     """
+    _setup_workdir(config.get("workdir"))
     valid_kind = ["HTCondor", "GNUParallelLocal"]
     if kind not in valid_kind:
         raise_error(
@@ -394,6 +431,7 @@ def reset(config: dict) -> None:
         The configuration to be used for resetting.
 
     """
+    _setup_workdir(config.get("workdir"))
     # Fetch storage
     storage = config["storage"]
     storage_uri = Path(storage["uri"])
@@ -425,7 +463,8 @@ def reset(config: dict) -> None:
 
 def list_elements(
     datagrabber: dict,
-    elements: Optional[Elements] = None,
+    elements: Elements | None = None,
+    workdir: str | Path | dict | None = None,
 ) -> str:
     """List elements of the datagrabber filtered using `elements`.
 
@@ -438,8 +477,13 @@ def list_elements(
     elements : list or None, optional
         Element(s) to filter using. Will be used to index the DataGrabber
         (default None).
+    workdir : str or pathlib.Path or dict or None, optional
+        Directory where the DataGrabber will be used, as in :func:`.run`.
+        If None, the default of :class:`.WorkDirManager` is used
+        (default None).
 
     """
+    _setup_workdir(workdir)
     # Get datagrabber to use
     datagrabber_object = _get_datagrabber(datagrabber)
 
@@ -462,3 +506,142 @@ def list_elements(
         elements_to_list.append(str_element)
 
     return "\n".join(elements_to_list)
+
+
+def parse_yaml(filepath: str | Path) -> dict:  # noqa: C901
+    """Parse YAML.
+
+    Parameters
+    ----------
+    filepath : str or pathlib.Path
+        The filepath to read from.
+
+    Returns
+    -------
+    dict
+        The contents represented as dictionary.
+
+    """
+    # Convert str to Path
+    if not isinstance(filepath, Path):
+        filepath = Path(filepath)
+
+    logger.info(f"Parsing yaml file: {filepath.absolute()!s}")
+    # Filepath existence check
+    if not filepath.exists():
+        raise_error(f"File does not exist: {filepath.absolute()!s}")
+    # Filepath reading
+    contents = yaml.load(filepath)
+    if "elements" in contents:
+        if contents["elements"] is None:
+            raise_error(
+                "The elements key was defined but its content is empty. "
+                "Please define the elements to operate on or remove the key."
+            )
+    # load modules
+    if "with" in contents:
+        to_load = contents["with"]
+        # Convert load modules to list
+        if not isinstance(to_load, list):
+            to_load = [to_load]
+        # Initialize list to have absolute paths for custom modules
+        final_to_load = []
+        for t_module in to_load:
+            if t_module.endswith(".py"):
+                logger.debug(f"Importing file: {t_module}")
+                # This resolves both absolute and relative paths
+                file_path = filepath.parent / t_module
+                if not file_path.exists():
+                    raise_error(
+                        f"File in 'with' section does not exist: {file_path}"
+                    )
+                # Add the parent directory to the sys.path so that the
+                # any imports from this module work correctly
+                t_path = str(file_path.parent)
+                if t_path not in sys.path:
+                    sys.path.append(str(file_path.parent))
+
+                spec = importlib.util.spec_from_file_location(
+                    t_module, file_path
+                )
+                module = importlib.util.module_from_spec(spec)  # type: ignore
+                sys.modules[t_module] = module
+                spec.loader.exec_module(module)  # type: ignore
+
+                # Add absolute path to final list
+                final_to_load.append(str(file_path.resolve()))
+
+                # Check if the module has junifer_module_deps function
+                if hasattr(module, "junifer_module_deps"):
+                    logger.debug(
+                        f"Module {t_module} has junifer_module_deps function"
+                    )
+                    # Get the dependencies
+                    deps = module.junifer_module_deps()
+                    # Add the dependencies to the final list
+                    for dep in deps:
+                        if dep not in final_to_load:
+                            final_to_load.append(
+                                str((file_path.parent / dep).resolve())
+                            )
+            else:
+                logger.info(f"Importing module: {t_module}")
+                importlib.import_module(t_module)
+                # Add module to final list
+                final_to_load.append(t_module)
+
+        # Replace modules to be loaded so that custom modules will take the
+        # absolute path. This was not the case as found in #224. Similar thing
+        # is done with the storage URI below.
+        contents["with"] = final_to_load
+
+    # Compute path for the URI parameter in storage files that are relative
+    # This is a tricky thing that appeared in #127. The problem is that
+    # the path in the URI parameter is relative to YAML file, not to the
+    # current working directory. If we leave it as is in the contents
+    # dict, then it will be used later in the ``build`` function as is,
+    # which will be computed relative to the current working directory.
+    # The solution is to compute the absolute path and replace the
+    # relative path in the contents dict with the absolute path.
+
+    # Check if the storage file is defined
+    if "storage" in contents:
+        if "uri" in contents["storage"]:
+            # Check if the storage file is relative
+            uri_path = Path(contents["storage"]["uri"])
+            if not uri_path.is_absolute():
+                # Compute the absolute path
+                contents["storage"]["uri"] = str(
+                    (filepath.parent / uri_path).resolve()
+                )
+
+    # Compute the absolute path of the data directory of the datagrabber (and
+    # of the datagrabbers of a MultipleDataGrabber) if it is relative; same
+    # motivation as above. Otherwise, it would be relative to the current
+    # working directory, which is the job directory when queueing.
+    if "datagrabber" in contents:
+        datagrabbers = [
+            contents["datagrabber"],
+            *contents["datagrabber"].get("datagrabbers", []),
+        ]
+        for datagrabber in datagrabbers:
+            if "datadir" in datagrabber:
+                datadir = Path(datagrabber["datadir"])
+                if not datadir.is_absolute():
+                    datagrabber["datadir"] = str(
+                        (filepath.parent / datadir).resolve()
+                    )
+
+    # Allow relative path if queue env kind is venv; same motivation as above
+    if "queue" in contents:
+        if "env" in contents["queue"]:
+            if "venv" == contents["queue"]["env"]["kind"]:
+                # Check if the env name is relative
+                venv_path = Path(contents["queue"]["env"]["name"])
+                if not venv_path.is_absolute():
+                    # Compute the absolute path
+                    contents["queue"]["env"]["name"] = str(
+                        (filepath.parent / venv_path).resolve()
+                    )
+
+    return contents

junifer/api/queue_context/__init__.pyi

@@ -1,5 +1,13 @@
-__all__ = ["QueueContextAdapter", "HTCondorAdapter", "GnuParallelLocalAdapter"]
+__all__ = [
+    "QueueContextAdapter",
+    "EnvKind",
+    "EnvShell",
+    "QueueContextEnv",
+    "HTCondorAdapter",
+    "HTCondorCollect",
+    "GnuParallelLocalAdapter",
+]
 
-from .queue_context_adapter import QueueContextAdapter
-from .htcondor_adapter import HTCondorAdapter
+from .queue_context_adapter import QueueContextAdapter, EnvKind, EnvShell, QueueContextEnv
+from .htcondor_adapter import HTCondorAdapter, HTCondorCollect
 from .gnu_parallel_local_adapter import GnuParallelLocalAdapter