♻️ Improved dependency management in BatchFile

dtriand · dtriand · commit a0f068b0d987 · 2025-06-14T15:18:22.000+02:00
xtl.automate.batchfile:BatchFile
- The constructor now also accepts a list of dependencies that are
  required for execution of the script
- New .dependencies property that returns a list of DependencySettings
  from the requested dependencies
- New .get_preamble() method that get a list of commands from the
  compute site which are added on the top of the script, e.g. module
  commands

xtl.automate.sites:ComputeSite
- New .get_preamble() abstract method

xtl.automate.sites:LocalSite
- The ._default_shell attribute is now set to DefaultShell, not None

xtl.automate.sites:ModulesSite
- New compute site that handles loading of environment modules in
  batch files

xtl.config.settings:DependencySettings
- Added .provides attribute to list all executables provided by a
  certain dependency
- Added .modules attribute to list all modules that are required to
  provide a certain dependency
diff --git a/src/xtl/automate/batchfile.py b/src/xtl/automate/batchfile.py
@@ -1,19 +1,25 @@
 __all__ = ['BatchFile']
 
 from pathlib import Path, PurePosixPath
-import stat
-from typing import Any, Sequence, Optional
+from typing import Any, Iterable, TYPE_CHECKING
 
+from xtl import Logger
 from xtl.automate.sites import ComputeSite, LocalSite
 from xtl.automate.shells import Shell, DefaultShell, WslShell
 from xtl.common.os import FilePermissions
 from xtl.common.compatibility import OS_POSIX
 
+if TYPE_CHECKING:
+    from xtl.config.settings import DependencySettings
+
+logger = Logger(__name__)
+
 
 class BatchFile:
 
-    def __init__(self, filename: str | Path, compute_site: Optional[ComputeSite] = None,
-                 shell: Shell | WslShell = DefaultShell):
+    def __init__(self, filename: str | Path, compute_site: ComputeSite = None,
+                 shell: Shell | WslShell = DefaultShell,
+                 dependencies: str | Iterable[str] = None):
         """
         A class for programmatically creating batch files. Additional configuration can be done by passing a ComputeSite
         instance.
@@ -42,6 +48,14 @@ def __init__(self, filename: str | Path, compute_site: Optional[ComputeSite] = N
             raise TypeError(f"compute_site must be an instance of ComputeSite, not {type(compute_site)}")
         self._compute_site = compute_site
 
+        # Set dependencies
+        if dependencies is None:
+            self._dependencies = set()
+        elif isinstance(dependencies, str):
+            self._dependencies = {dependencies}
+        elif isinstance(dependencies, Iterable):
+            self._dependencies = set(dependencies)
+
         # List of lines of the batch file
         self._lines = []
 
@@ -83,6 +97,22 @@ def permissions(self, value: int | str | FilePermissions):
         """
         self._permissions = FilePermissions(value)
 
+    @property
+    def dependencies(self) -> list['DependencySettings']:
+        """
+        Returns a list of DependencySettings that are required by this batch file.
+        """
+        from xtl import settings
+        deps = list()
+
+        for dname, dep in settings.dependencies:
+            if dname in self._dependencies:
+                deps.append(dep)
+            else:
+                logger.warning('Dependency %(name)s not found in xtl.settings, '
+                               'skipping', {'name': dname})
+        return deps
+
     def get_execute_command(self, arguments: list = None, as_list: bool = False) -> str:
         if self._wsl_filename:
             return self.shell.get_batch_command(batch_file=self._wsl_filename, batch_arguments=arguments, as_list=as_list)
@@ -112,11 +142,11 @@ def add_command(self, command: str) -> None:
         """
         self._add_line(command)
 
-    def add_commands(self, *commands: str | Sequence[str]) -> None:
+    def add_commands(self, *commands: str | Iterable[str]) -> None:
         """
         Add multiple commands to the batch file all at once.
         """
-        if len(commands) == 1 and isinstance(commands[0], Sequence):  # unpack a list or tuple of commands
+        if len(commands) == 1 and isinstance(commands[0], Iterable):  # unpack a list or tuple of commands
             commands = commands[0]
         for command in commands:
             if not isinstance(command, str):
@@ -129,18 +159,24 @@ def add_comment(self, comment):
         """
         self._add_line(f"{self.shell.comment_char} {comment}")
 
-    def load_modules(self, modules: str | Sequence[str]):
+    def load_modules(self, modules: str | Iterable[str]):
         """
         Add command for loading one or more modules on the compute site.
         """
+        # TODO: Remove when refactoring AutoPROCJob
         self.add_command(self.compute_site.load_modules(modules))
 
     def purge_modules(self):
         """
         Add command for purging all loaded modules on the compute site.
         """
+        # TODO: Remove when refactoring AutoPROCJob
         self.add_command(self.compute_site.purge_modules())
 
+    def get_preamble(self) -> list[str]:
+        return self.compute_site.get_preamble(dependencies=self.dependencies,
+                                              shell=self.shell)
+
     def assign_variable(self, variable: str, value: Any):
         """
         Assign a value to a variable in the batch file.
@@ -161,9 +197,13 @@ def save(self, change_permissions: bool = True):
         self._filename.unlink(missing_ok=True)
 
         # Write contents to file
-        text = self.shell.shebang + self.shell.new_line_char if self.shell.shebang else ''
-        text += self.shell.new_line_char.join(self._lines)
-        self._filename.write_text(text, encoding='utf-8', newline=self.shell.new_line_char)
+        nl = self.shell.new_line_char
+        preamble = self.get_preamble()
+        text = self.shell.shebang + nl if self.shell.shebang else ''
+        if preamble:
+            text += nl.join(self.get_preamble()) + nl
+        text += nl.join(self._lines)
+        self._filename.write_text(text, encoding='utf-8', newline=nl)
 
         # Update permissions (user: read, write, execute; group: read, write)
         if change_permissions and OS_POSIX:
diff --git a/src/xtl/automate/sites.py b/src/xtl/automate/sites.py
@@ -1,9 +1,12 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Optional, Sequence
+from typing import Optional, Iterable, TYPE_CHECKING
 
 from xtl.automate.priority_system import PrioritySystemType, DefaultPrioritySystem, NicePrioritySystem
-from xtl.automate.shells import Shell, DefaultShell, BashShell
+from xtl.automate.shells import Shell, DefaultShell, BashShell, CmdShell, WslShell
+
+if TYPE_CHECKING:
+    from xtl.config.settings import DependencySettings
 
 
 @dataclass
@@ -15,7 +18,7 @@ class ComputeSite(ABC):
 
     _priority_system: Optional[PrioritySystemType] = None
     _default_shell: Optional[Shell] = None
-    _supported_shells: Sequence[Shell] | None = None
+    _supported_shells: Iterable[Shell] | None = None
 
     @property
     def priority_system(self):
@@ -38,7 +41,7 @@ def default_shell(self) -> Optional[Shell]:
         return self._default_shell
 
     @property
-    def supported_shells(self) -> Sequence[Shell] | None:
+    def supported_shells(self) -> Iterable[Shell] | None:
         """
         The supported shells for this compute site. This is used to validate the shell passed to the BatchFile.
         """
@@ -53,17 +56,25 @@ def is_valid_shell(self, shell: Shell) -> bool:
         return shell in self.supported_shells
 
     @abstractmethod
-    def load_modules(self, modules: str | Sequence[str]) -> str:
+    def load_modules(self, modules: str | Iterable[str]) -> str:
         """
         Generates a command for loading the specified modules on the compute site.
         """
+        #  TODO: Remove when refactoring AutoPROCJob
         pass
 
     @abstractmethod
     def purge_modules(self) -> str:
         """
         Generates a command for purging all loaded modules on the compute site.
         """
+        #  TODO: Remove when refactoring AutoPROCJob
+        pass
+
+    @abstractmethod
+    def get_preamble(self, dependencies: 'DependencySettings' |
+                                         Iterable['DependencySettings'],
+                     shell: Shell | WslShell = None) -> list[str]:
         pass
 
     def prepare_command(self, command: str) -> str:
@@ -80,13 +91,64 @@ def __init__(self):
         assumes that all required executables are available on PATH.
         """
         self._priority_system = DefaultPrioritySystem()
+        self._default_shell = DefaultShell
 
     def load_modules(self, modules) -> str:
         return ''
 
     def purge_modules(self) -> str:
         return ''
 
+    def get_preamble(self, dependencies: 'DependencySettings' |
+                     Iterable['DependencySettings'],
+                     shell: Shell | WslShell = None) -> list[str]:
+        # Returns an empty preamble for the local site, as it assumes that all required
+        # executables are available on PATH
+        return []
+
+
+class ModulesSite(ComputeSite):
+    def __init__(self):
+        self._default_shell = DefaultShell
+        self._priority_system = DefaultPrioritySystem()
+
+    def load_modules(self, modules: str | Iterable[str]) -> str: ...
+
+    def purge_modules(self) -> str: ...
+
+    @staticmethod
+    def _load_modules(modules: Iterable[str], shell: Shell = None) -> str:
+        cmd = 'module load ' + ' '.join(modules)
+        if shell is CmdShell:
+            # For Windows CMD, we need to use 'call' to execute the module commands
+            cmd = f'call {cmd}'
+        return cmd
+
+    @staticmethod
+    def _purge_modules(shell: Shell = None) -> str:
+        cmd = 'module purge'
+        if shell is CmdShell:
+            cmd = f'call {cmd}'
+        return cmd
+
+    def get_preamble(self, dependencies: 'DependencySettings' |
+                                         Iterable['DependencySettings'],
+                     shell: Shell | WslShell = None) -> list[str]:
+        if shell is None:
+            shell = self._default_shell
+        if not isinstance(dependencies, Iterable):
+            dependencies = [dependencies]
+
+        modules = []
+        for dep in dependencies:
+            if dep.modules:
+                modules.extend(dep.modules)
+
+        cmds = [self._purge_modules(shell=shell)]
+        if modules:
+            cmds.append(self._load_modules(modules, shell=shell))
+        return cmds
+
 
 class BiotixHPC(ComputeSite):
     def __init__(self):
@@ -99,13 +161,13 @@ def __init__(self):
         self._default_shell = BashShell
         self._supported_shells = [BashShell]
 
-    def load_modules(self, modules: str | Sequence[str]) -> str:
+    def load_modules(self, modules: str | Iterable[str]) -> str:
         mods = []
         if isinstance(modules, str):
             # Check for space-separated modules in a single string
             for mod in modules.split():
                 mods.append(mod)
-        elif isinstance(modules, Sequence):
+        elif isinstance(modules, Iterable):
             # Otherwise append each module in the list
             for i, mod in enumerate(modules):
                 if not isinstance(mod, str):
@@ -122,5 +184,11 @@ def load_modules(self, modules: str | Sequence[str]) -> str:
     def purge_modules(self) -> str:
         return 'module purge'
 
+    def get_preamble(self, dependencies: 'DependencySettings' |
+                                         Iterable['DependencySettings'],
+                     shell: Shell | WslShell = None) -> list[str]:
+        # TODO: Implement this
+        return []
+
 
-ComputeSiteType = LocalSite | BiotixHPC
+ComputeSiteType = LocalSite | ModulesSite | BiotixHPC
diff --git a/src/xtl/config/settings.py b/src/xtl/config/settings.py
@@ -5,8 +5,7 @@
 """
 
 from pathlib import Path
-from pprint import pprint
-from typing import ClassVar, Optional
+from typing import ClassVar, Optional, TYPE_CHECKING
 
 from pydantic import PrivateAttr
 
@@ -100,14 +99,25 @@ class DependencySettings(Settings):
     """
 
     # Model attributes
-    path: Optional[Path] = Option(default=None,
-                                  desc='Directory containing binaries',
-                                  validator=CastAsNoneIfEmpty())
-    # TODO: Enable when ModuleSite is implemented
-    modules: Optional[list[str]] = Option(default=None,
-                                          desc='Modules that provide the dependency',
-                                          validator=CastAsNoneIfEmpty(),
-                                          exclude=True)
+    provides: set[str] = \
+        Option(
+            default_factory=set,
+            desc='List of executables provided by this dependency',
+            exclude=True
+        )
+
+    path: Optional[Path] = \
+        Option(
+            default=None,
+            desc='Directory containing binaries',
+            validator=CastAsNoneIfEmpty()
+        )
+
+    modules: set[str] = \
+        Option(
+            default_factory=set,
+            desc='Modules that provide the dependency'
+        )
 
 
 class DependenciesSettings(Settings):
@@ -116,7 +126,12 @@ class DependenciesSettings(Settings):
     """
 
     # Model attributes
-    autoproc: DependencySettings = Option(default=DependencySettings())
+    autoproc: DependencySettings = \
+        Option(
+            default=DependencySettings(
+                provides={'process', 'process_wf'}
+            )
+        )
 
 
 class CLIAutoprocSettings(Settings):
diff --git a/tests/automate/test_sites.py b/tests/automate/test_sites.py
@@ -1,7 +1,7 @@
 import pytest
 
 from xtl.automate.priority_system import DefaultPrioritySystem, NicePrioritySystem
-from xtl.automate.shells import BashShell
+from xtl.automate.shells import BashShell, DefaultShell
 from xtl.automate.sites import ComputeSite, LocalSite, BiotixHPC
 
 
@@ -17,7 +17,7 @@ def test_priority_system(self):
 
     def test_default_shell(self):
         cs = LocalSite()
-        assert cs.default_shell is None
+        assert cs.default_shell is DefaultShell
         assert cs.supported_shells is None
 
     def test_load_modules(self):
