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

junifer/api/queue_context/gnu_parallel_local_adapter.py

@@ -6,11 +6,17 @@
 import shutil
 import textwrap
 from pathlib import Path
-from typing import Optional
+from typing import Any
 
 from ...typing import Elements
-from ...utils import logger, make_executable, raise_error, run_ext_cmd
-from .queue_context_adapter import QueueContextAdapter
+from ...utils import make_executable, raise_error, run_ext_cmd
+from .queue_context_adapter import (
+    EnvKind,
+    EnvShell,
+    QueueContextAdapter,
+    QueueContextEnv,
+    logger,
+)
 
 
 __all__ = ["GnuParallelLocalAdapter"]
@@ -27,14 +33,14 @@ class GnuParallelLocalAdapter(QueueContextAdapter):
         The path to the job directory.
     yaml_config_path : pathlib.Path
         The path to the YAML config file.
-    elements : list of str or tuple
+    elements : ``Elements``
         Element(s) to process. Will be used to index the DataGrabber.
     pre_run : str or None, optional
         Extra shell commands to source before the run (default None).
     pre_collect : str or None, optional
-        Extra bash commands to source before the collect (default None).
-    env : dict, optional
-        The Python environment configuration. If None, will run without a
+        Extra shell commands to source before the collect (default None).
+    env : :class:`.QueueContextEnv` or None, optional
+        The environment configuration. If None, will run without a
         virtual environment of any kind (default None).
     verbose : str, optional
         The level of verbosity (default "info").
@@ -44,12 +50,6 @@ class GnuParallelLocalAdapter(QueueContextAdapter):
     submit : bool, optional
         Whether to submit the jobs (default False).
 
-    Raises
-    ------
-    ValueError
-        If ``env.kind`` is invalid or
-        if ``env.shell`` is invalid.
-
     See Also
     --------
     QueueContextAdapter :
@@ -59,87 +59,44 @@ class GnuParallelLocalAdapter(QueueContextAdapter):
 
     """
 
-    def __init__(
-        self,
-        job_name: str,
-        job_dir: Path,
-        yaml_config_path: Path,
-        elements: Elements,
-        pre_run: Optional[str] = None,
-        pre_collect: Optional[str] = None,
-        env: Optional[dict[str, str]] = None,
-        verbose: str = "info",
-        verbose_datalad: Optional[str] = None,
-        submit: bool = False,
-    ) -> None:
-        """Initialize the class."""
-        self._job_name = job_name
-        self._job_dir = job_dir
-        self._yaml_config_path = yaml_config_path
-        self._elements = elements
-        self._pre_run = pre_run
-        self._pre_collect = pre_collect
-        self._check_env(env)
-        self._verbose = verbose
-        self._verbose_datalad = verbose_datalad
-        self._submit = submit
-
-        self._log_dir = self._job_dir / "logs"
-        self._pre_run_path = self._job_dir / "pre_run.sh"
-        self._pre_collect_path = self._job_dir / "pre_collect.sh"
-        self._run_path = self._job_dir / f"run_{self._job_name}.sh"
-        self._collect_path = self._job_dir / f"collect_{self._job_name}.sh"
-        self._run_joblog_path = self._job_dir / f"run_{self._job_name}_joblog"
-        self._elements_file_path = self._job_dir / "elements"
-
-    def _check_env(self, env: Optional[dict[str, str]]) -> None:
-        """Check value of env parameter on init.
-
-        Parameters
-        ----------
-        env : dict or None
-            The value of env parameter.
-
-        Raises
-        ------
-        ValueError
-            If ``env.kind`` is invalid.
-
-        """
-        # Set env related variables
-        if env is None:
-            env = {"kind": "local"}
-        # Check env kind
-        valid_env_kinds = ["conda", "venv", "local"]
-        if env["kind"] not in valid_env_kinds:
-            raise_error(
-                f"Invalid value for `env.kind`: {env['kind']}, "
-                f"must be one of {valid_env_kinds}"
+    job_name: str
+    job_dir: Path
+    yaml_config_path: Path
+    elements: Elements
+    pre_run: str | None = None
+    pre_collect: str | None = None
+    env: QueueContextEnv | None = None
+    verbose: str = "info"
+    verbose_datalad: str | None = None
+    submit: bool = False
+
+    def model_post_init(self, context: Any):  # noqa: D102
+        if self.env is None:
+            self.env = QueueContextEnv(
+                kind=EnvKind.Local, shell=EnvShell.Bash, name=""
             )
+        if self.env["kind"] == EnvKind.Local:
+            # No virtual environment
+            self._executable = "junifer"
+            self._arguments = ""
         else:
-            # Check shell
-            shell = env.get("shell", "bash")
-            valid_shells = ["bash", "zsh"]
-            if shell not in valid_shells:
-                raise_error(
-                    f"Invalid value for `env.shell`: {shell}, "
-                    f"must be one of {valid_shells}"
-                )
-            self._shell = shell
-            # Set variables
-            if env["kind"] == "local":
-                # No virtual environment
-                self._executable = "junifer"
-                self._arguments = ""
-            else:
-                self._executable = f"run_{env['kind']}.{self._shell}"
-                self._arguments = f"{env['name']} junifer"
-                self._exec_path = self._job_dir / self._executable
-
-    def elements(self) -> str:
+            if self.env["name"] is None:
+                raise_error("`env.name` is required")
+            self._executable = f"run_{self.env['kind']}.{self.env['shell']}"
+            self._arguments = f"{self.env['name']} junifer"
+            self._exec_path = self.job_dir / self._executable
+        self._log_dir = self.job_dir / "logs"
+        self._pre_run_path = self.job_dir / "pre_run.sh"
+        self._pre_collect_path = self.job_dir / "pre_collect.sh"
+        self._run_path = self.job_dir / f"run_{self.job_name}.sh"
+        self._collect_path = self.job_dir / f"collect_{self.job_name}.sh"
+        self._run_joblog_path = self.job_dir / f"run_{self.job_name}_joblog"
+        self._elements_file_path = self.job_dir / "elements"
+
+    def elements_to_run(self) -> str:
         """Return elements to run."""
         elements_to_run = []
-        for element in self._elements:
+        for element in self.elements:
             # Stringify elements if tuple for operation
             str_element = (
                 ",".join(element) if isinstance(element, tuple) else element
@@ -148,26 +105,26 @@ class GnuParallelLocalAdapter(QueueContextAdapter):
 
         return "\n".join(elements_to_run)
 
-    def pre_run(self) -> str:
+    def pre_run_cmds(self) -> str:
         """Return pre-run commands."""
         fixed = (
-            f"#!/usr/bin/env {self._shell}\n\n"
+            f"#!/usr/bin/env {self.env['shell']}\n\n"
             "# This script is auto-generated by junifer.\n\n"
             "# Force datalad to run in non-interactive mode\n"
             "DATALAD_UI_INTERACTIVE=false\n"
         )
-        var = self._pre_run or ""
+        var = self.pre_run or ""
         return fixed + "\n" + var
 
-    def run(self) -> str:
+    def run_cmds(self) -> str:
         """Return run commands."""
-        verbose_args = f"--verbose {self._verbose}"
-        if self._verbose_datalad:
+        verbose_args = f"--verbose {self.verbose}"
+        if self.verbose_datalad:
             verbose_args = (
-                f"{verbose_args} --verbose-datalad {self._verbose_datalad}"
+                f"{verbose_args} --verbose-datalad {self.verbose_datalad}"
             )
         return (
-            f"#!/usr/bin/env {self._shell}\n\n"
+            f"#!/usr/bin/env {self.env['shell']}\n\n"
             "# This script is auto-generated by junifer.\n\n"
             "# Run pre_run.sh\n"
             f"sh {self._pre_run_path.resolve()!s}\n\n"
@@ -177,38 +134,38 @@ class GnuParallelLocalAdapter(QueueContextAdapter):
             "--delay 60 "  # wait 1 min before next job is spawned
             f"--results {self._log_dir} "
             f"--arg-file {self._elements_file_path.resolve()!s} "
-            f"{self._job_dir.resolve()!s}/{self._executable} "
+            f"{self.job_dir.resolve()!s}/{self._executable} "
             f"{self._arguments} run "
-            f"{self._yaml_config_path.resolve()!s} "
+            f"{self.yaml_config_path.resolve()!s} "
             f"{verbose_args} "
             f"--element"
         )
 
-    def pre_collect(self) -> str:
+    def pre_collect_cmds(self) -> str:
         """Return pre-collect commands."""
         fixed = (
-            f"#!/usr/bin/env {self._shell}\n\n"
+            f"#!/usr/bin/env {self.env['shell']}\n\n"
             "# This script is auto-generated by junifer.\n"
         )
-        var = self._pre_collect or ""
+        var = self.pre_collect or ""
         return fixed + "\n" + var
 
-    def collect(self) -> str:
+    def collect_cmds(self) -> str:
         """Return collect commands."""
-        verbose_args = f"--verbose {self._verbose}"
-        if self._verbose_datalad:
+        verbose_args = f"--verbose {self.verbose}"
+        if self.verbose_datalad:
             verbose_args = (
-                f"{verbose_args} --verbose-datalad {self._verbose_datalad}"
+                f"{verbose_args} --verbose-datalad {self.verbose_datalad}"
             )
         return (
-            f"#!/usr/bin/env {self._shell}\n\n"
+            f"#!/usr/bin/env {self.env['shell']}\n\n"
             "# This script is auto-generated by junifer.\n\n"
             "# Run pre_collect.sh\n"
             f"sh {self._pre_collect_path.resolve()!s}\n\n"
             "# Run `junifer collect`\n"
-            f"{self._job_dir.resolve()!s}/{self._executable} "
+            f"{self.job_dir.resolve()!s}/{self._executable} "
             f"{self._arguments} collect "
-            f"{self._yaml_config_path.resolve()!s} "
+            f"{self.yaml_config_path.resolve()!s} "
             f"{verbose_args}"
         )
 
@@ -231,40 +188,44 @@ class GnuParallelLocalAdapter(QueueContextAdapter):
             f"{self._elements_file_path.resolve()!s}"
         )
         self._elements_file_path.touch()
-        self._elements_file_path.write_text(textwrap.dedent(self.elements()))
+        self._elements_file_path.write_text(
+            textwrap.dedent(self.elements_to_run())
+        )
         # Create pre run
         logger.info(
-            f"Writing {self._pre_run_path.name} to {self._job_dir.resolve()!s}"
+            f"Writing {self._pre_run_path.name} to {self.job_dir.resolve()!s}"
         )
         self._pre_run_path.touch()
-        self._pre_run_path.write_text(textwrap.dedent(self.pre_run()))
+        self._pre_run_path.write_text(textwrap.dedent(self.pre_run_cmds()))
         make_executable(self._pre_run_path)
         # Create run
         logger.info(
-            f"Writing {self._run_path.name} to {self._job_dir.resolve()!s}"
+            f"Writing {self._run_path.name} to {self.job_dir.resolve()!s}"
         )
         self._run_path.touch()
-        self._run_path.write_text(textwrap.dedent(self.run()))
+        self._run_path.write_text(textwrap.dedent(self.run_cmds()))
         make_executable(self._run_path)
         # Create pre collect
         logger.info(
             f"Writing {self._pre_collect_path.name} to "
-            f"{self._job_dir.resolve()!s}"
+            f"{self.job_dir.resolve()!s}"
         )
         self._pre_collect_path.touch()
-        self._pre_collect_path.write_text(textwrap.dedent(self.pre_collect()))
+        self._pre_collect_path.write_text(
+            textwrap.dedent(self.pre_collect_cmds())
+        )
         make_executable(self._pre_collect_path)
         # Create collect
         logger.info(
-            f"Writing {self._collect_path.name} to {self._job_dir.resolve()!s}"
+            f"Writing {self._collect_path.name} to {self.job_dir.resolve()!s}"
         )
         self._collect_path.touch()
-        self._collect_path.write_text(textwrap.dedent(self.collect()))
+        self._collect_path.write_text(textwrap.dedent(self.collect_cmds()))
         make_executable(self._collect_path)
         # Submit if required
         run_cmd = f"sh {self._run_path.resolve()!s}"
         collect_cmd = f"sh {self._collect_path.resolve()!s}"
-        if self._submit:
+        if self.submit:
             logger.info(
                 "Shell scripts created, the following will be run:\n"
                 f"{run_cmd}\n"