hpcflow-new2 0.2.0a190__py3-none-any.whl → 0.2.0a200__py3-none-any.whl

hpcflow/sdk/core/types.py

@@ -1,6 +1,7 @@
 """
 Types to support the core SDK.
 """
+
 from __future__ import annotations
 from typing import Any, Literal, Protocol, TYPE_CHECKING
 from typing_extensions import NotRequired, TypeAlias, TypedDict
@@ -205,6 +206,9 @@ class SchemaInputKwargs(TypedDict):
     multiple: bool
     #: The labels.
     labels: dict[str, LabelInfo] | None
+    #: The number or proportional of permitted unset parameter data found when reolving
+    #: this input from upstream outputs.
+    allow_failed_dependencies: int | float | bool | None
 
 
 class RuleArgs(TypedDict):
@@ -333,11 +337,51 @@ class WorkflowTemplateTaskData(TypedDict):
     #: The schema, if known.
     schema: NotRequired[Any | list[Any]]
     #: The element sets, if known.
-    element_sets: NotRequired[list["WorkflowTemplateTaskData"]]
+    element_sets: NotRequired[list[WorkflowTemplateElementSetData]]
     #: The output labels, if known.
     output_labels: NotRequired[list[str]]
 
 
+class WorkflowTemplateElementSetData(TypedDict):
+    """
+    Descriptor for element set data within a workflow template parametrisation.
+    """
+
+    #: Inputs to the set of elements.
+    inputs: NotRequired[list[dict[str, Any]]]
+    #: Input files to the set of elements.
+    input_files: NotRequired[list[dict[str, Any]]]
+    #: Description of how to repeat the set of elements.
+    repeats: NotRequired[int | list[RepeatsDescriptor]]
+    #: Groupings in the set of elements.
+    groups: NotRequired[list[dict[str, Any]]]
+    #: Resources to use for the set of elements.
+    resources: NotRequired[dict[str, Any]]
+    #: Input value sequences to parameterise over.
+    sequences: NotRequired[list[dict[str, Any]]]
+    #: Input value multi-path sequences to parameterise over.
+    multi_path_sequences: NotRequired[list[dict[str, Any]]]
+    #: Input source descriptors.
+    input_sources: NotRequired[dict[str, list]]
+    #: How to handle nesting of iterations.
+    nesting_order: NotRequired[dict[str, float]]
+    #: Which environment preset to use.
+    env_preset: NotRequired[str]
+    #: Environment descriptors to use.
+    environments: NotRequired[dict[str, dict[str, Any]]]
+    #: List of global element iteration indices from which inputs for
+    #: the new elements associated with this element set may be sourced.
+    #: If ``None``, all iterations are valid.
+    sourceable_elem_iters: NotRequired[list[int]]
+    #: Whether to allow sources to come from distinct element sub-sets.
+    allow_non_coincident_task_sources: NotRequired[bool]
+    #: Whether this initialisation is the first for this data (i.e. not a
+    #: reconstruction from persistent workflow data), in which case, we merge
+    #: ``environments`` into ``resources`` using the "any" scope, and merge any multi-
+    #: path sequences into the sequences list.
+    is_creation: NotRequired[bool]
+
+
 class Pending(TypedDict):
     """
     Pending update information. Internal use only.
@@ -385,3 +429,12 @@ class ResourcePersistingWorkflow(Protocol):
         """
         Check if all the parameters exist.
         """
+
+
+BlockActionKey: TypeAlias = "tuple[int | str, int | str, int | str]"
+"""
+The type of indices that locate an action within a submission. The indices represent,
+respectively, the jobscript index, the jobscript-block index, and the block-action index.
+Usually, these are integers, but in the case of strings, they will correspond to shell
+environment variables.
+"""

hpcflow/sdk/core/utils.py

@@ -4,8 +4,12 @@ Miscellaneous utilities.
 
 from __future__ import annotations
 from collections import Counter
+from asyncio import events
+import contextvars
+import contextlib
 import copy
 import enum
+import functools
 import hashlib
 from itertools import accumulate, islice
 from importlib import resources
@@ -20,7 +24,8 @@ import string
 import subprocess
 from datetime import datetime, timedelta, timezone
 import sys
-from typing import cast, overload, TypeVar, TYPE_CHECKING
+import traceback
+from typing import Literal, cast, overload, TypeVar, TYPE_CHECKING
 import fsspec  # type: ignore
 import numpy as np
 
@@ -33,12 +38,13 @@ from hpcflow.sdk.core.errors import (
     MissingVariableSubstitutionError,
 )
 from hpcflow.sdk.log import TimeIt
+from hpcflow.sdk.utils.deferred_file import DeferredFileWriter
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence
     from contextlib import AbstractContextManager
     from types import ModuleType
-    from typing import Any, IO
+    from typing import Any, IO, Iterator
     from typing_extensions import TypeAlias
     from numpy.typing import NDArray
     from ..typing import PathLike
@@ -124,7 +130,9 @@ def check_valid_py_identifier(name: str) -> str:
 
 
 @overload
-def group_by_dict_key_values(lst: list[dict[T, T2]], key: T) -> list[list[dict[T, T2]]]:
+def group_by_dict_key_values(  # type: ignore[overload-overlap]
+    lst: list[dict[T, T2]], key: T
+) -> list[list[dict[T, T2]]]:
     ...
 
 
@@ -326,7 +334,7 @@ def get_relative_path(path1: Sequence[T], path2: Sequence[T]) -> Sequence[T]:
 
 
 def search_dir_files_by_regex(
-    pattern: str | re.Pattern[str], directory: str = "."
+    pattern: str | re.Pattern[str], directory: str | os.PathLike = "."
 ) -> list[str]:
     """Search recursively for files in a directory by a regex pattern and return matching
     file paths, relative to the given directory."""
@@ -406,20 +414,44 @@ def substitute_string_vars(string: str, variables: dict[str, str]):
 
 @TimeIt.decorator
 def read_YAML_str(
-    yaml_str: str, typ="safe", variables: dict[str, str] | None = None
+    yaml_str: str, typ="safe", variables: dict[str, str] | Literal[False] | None = None
 ) -> Any:
-    """Load a YAML string. This will produce basic objects."""
-    if variables is not None and "<<var:" in yaml_str:
-        yaml_str = substitute_string_vars(yaml_str, variables=variables)
+    """Load a YAML string. This will produce basic objects.
+
+    Parameters
+    ----------
+    yaml_str:
+        The YAML string to parse.
+    typ:
+        Load type passed to the YAML library.
+    variables:
+        String variables to substitute in `yaml_str`. Substitutions will be attempted if
+        the file looks to contain variable references (like "<<var:name>>"). If set to
+        `False`, no substitutions will occur.
+    """
+    if variables is not False and "<<var:" in yaml_str:
+        yaml_str = substitute_string_vars(yaml_str, variables=variables or {})
     yaml = YAML(typ=typ)
     return yaml.load(yaml_str)
 
 
 @TimeIt.decorator
 def read_YAML_file(
-    path: PathLike, typ="safe", variables: dict[str, str] | None = None
+    path: PathLike, typ="safe", variables: dict[str, str] | Literal[False] | None = None
 ) -> Any:
-    """Load a YAML file. This will produce basic objects."""
+    """Load a YAML file. This will produce basic objects.
+
+    Parameters
+    ----------
+    path:
+        Path to the YAML file to parse.
+    typ:
+        Load type passed to the YAML library.
+    variables:
+        String variables to substitute in the file given by `path`. Substitutions will be
+        attempted if the file looks to contain variable references (like "<<var:name>>").
+        If set to `False`, no substitutions will occur.
+    """
     with fsspec.open(path, "rt") as f:
         yaml_str: str = f.read()
     return read_YAML_str(yaml_str, typ=typ, variables=variables)
@@ -432,15 +464,37 @@ def write_YAML_file(obj, path: str | Path, typ: str = "safe") -> None:
         yaml.dump(obj, fp)
 
 
-def read_JSON_string(json_str: str, variables: dict[str, str] | None = None) -> Any:
-    """Load a JSON string. This will produce basic objects."""
-    if variables is not None and "<<var:" in json_str:
-        json_str = substitute_string_vars(json_str, variables=variables)
+def read_JSON_string(
+    json_str: str, variables: dict[str, str] | Literal[False] | None = None
+) -> Any:
+    """Load a JSON string. This will produce basic objects.
+
+    Parameters
+    ----------
+    json_str:
+        The JSON string to parse.
+    variables:
+        String variables to substitute in `json_str`. Substitutions will be attempted if
+        the file looks to contain variable references (like "<<var:name>>"). If set to
+        `False`, no substitutions will occur.
+    """
+    if variables is not False and "<<var:" in json_str:
+        json_str = substitute_string_vars(json_str, variables=variables or {})
     return json.loads(json_str)
 
 
-def read_JSON_file(path, variables: dict[str, str] | None = None) -> Any:
-    """Load a JSON file. This will produce basic objects."""
+def read_JSON_file(path, variables: dict[str, str] | Literal[False] | None = None) -> Any:
+    """Load a JSON file. This will produce basic objects.
+
+    Parameters
+    ----------
+    path:
+        Path to the JSON file to parse.
+    variables:
+        String variables to substitute in the file given by `path`. Substitutions will be
+        attempted if the file looks to contain variable references (like "<<var:name>>").
+        If set to `False`, no substitutions will occur.
+    """
     with fsspec.open(path, "rt") as f:
         json_str: str = f.read()
     return read_JSON_string(json_str, variables=variables)
@@ -771,7 +825,12 @@ class JSONLikeDirSnapShot(DirectorySnapshot):
         See :py:meth:`to_json_like`.
     """
 
-    def __init__(self, root_path: str | None = None, data: dict[str, list] | None = None):
+    def __init__(
+        self,
+        root_path: str | None = None,
+        data: dict[str, list] | None = None,
+        use_strings: bool = False,
+    ):
         """
         Create an empty snapshot or load from JSON-like data.
         """
@@ -786,6 +845,7 @@ class JSONLikeDirSnapShot(DirectorySnapshot):
             for name, item in data.items():
                 # add root path
                 full_name = str(PurePath(root_path) / PurePath(name))
+                item = [int(i) for i in item] if use_strings else item
                 stat_dat, inode_key = item[:-2], item[-2:]
                 self._stat_info[full_name] = os.stat_result(stat_dat)
                 self._inode_to_path[tuple(inode_key)] = full_name
@@ -794,7 +854,7 @@ class JSONLikeDirSnapShot(DirectorySnapshot):
         """Take the snapshot."""
         super().__init__(*args, **kwargs)
 
-    def to_json_like(self) -> dict[str, Any]:
+    def to_json_like(self, use_strings: bool = False) -> dict[str, Any]:
         """Export to a dict that is JSON-compatible and can be later reloaded.
 
         The last two integers in `data` for each path are the keys in
@@ -807,13 +867,16 @@ class JSONLikeDirSnapShot(DirectorySnapshot):
         # store efficiently:
         inode_invert = {v: k for k, v in self._inode_to_path.items()}
         data: dict[str, list] = {
-            str(PurePath(k).relative_to(root_path)): [*v, *inode_invert[k]]
+            str(PurePath(k).relative_to(root_path)): [
+                str(i) if use_strings else i for i in [*v, *inode_invert[k]]
+            ]
             for k, v in self._stat_info.items()
         }
 
         return {
             "root_path": root_path,
             "data": data,
+            "use_strings": use_strings,
         }
 
 
@@ -1083,3 +1146,57 @@ def get_file_context(
     except AttributeError:
         # < python 3.9
         return resources.path(package, src or "")
+
+
+@contextlib.contextmanager
+def redirect_std_to_file(
+    file,
+    mode: Literal["w", "a"] = "a",
+    ignore: Callable[[BaseException], Literal[True] | int] | None = None,
+) -> Iterator[None]:
+    """Temporarily redirect both stdout and stderr to a file, and if an exception is
+    raised, catch it, print the traceback to that file, and exit.
+
+    File creation is deferred until an actual write is required.
+
+    Parameters
+    ----------
+    ignore
+        Callable to test if a given exception should be ignored. If an exception is
+        not ignored, its traceback will be printed to `file` and the program will
+        exit with exit code 1. The callable should accept one parameter, the
+        exception, and should return True if that exception should be ignored, or
+        an integer representing the exit code to exit the program with if that
+        exception should not be ignored.  By default, no exceptions are ignored.
+
+    """
+    ignore = ignore or (lambda _: 1)
+    with DeferredFileWriter(file, mode=mode) as fp:
+        with contextlib.redirect_stdout(fp):
+            with contextlib.redirect_stderr(fp):
+                try:
+                    yield
+                except BaseException as exc:
+                    ignore_ret = ignore(exc)
+                    if ignore_ret is not True:
+                        traceback.print_exc()
+                        sys.exit(ignore_ret)
+
+
+async def to_thread(func, /, *args, **kwargs):
+    """Copied from https://github.com/python/cpython/blob/4b4227b907a262446b9d276c274feda2590a4e6e/Lib/asyncio/threads.py
+    to support Python 3.8, which does not have `asyncio.to_thread`.
+
+    Asynchronously run function *func* in a separate thread.
+
+    Any *args and **kwargs supplied for this function are directly passed
+    to *func*. Also, the current :class:`contextvars.Context` is propagated,
+    allowing context variables from the main thread to be accessed in the
+    separate thread.
+
+    Return a coroutine that can be awaited to get the eventual result of *func*.
+    """
+    loop = events.get_running_loop()
+    ctx = contextvars.copy_context()
+    func_call = functools.partial(ctx.run, func, *args, **kwargs)
+    return await loop.run_in_executor(None, func_call)