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

hpcflow/sdk/persistence/store_resource.py

@@ -2,14 +2,22 @@
 Models of data stores as resources.
 """
 
+from __future__ import annotations
 from abc import ABC, abstractmethod
 import copy
 import json
-from pathlib import Path
-from typing import Callable, Union
+from typing import Any, Callable, TYPE_CHECKING
 
 from hpcflow.sdk.core.utils import get_md5_hash
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from logging import Logger
+    from pathlib import Path
+    import zarr  # type: ignore
+    from fsspec import AbstractFileSystem  # type: ignore
+    from ..app import BaseApp
+
 
 class StoreResource(ABC):
     """Class to represent a persistent resource within which store data lives.
@@ -25,37 +33,37 @@ class StoreResource(ABC):
         The store name.
     """
 
-    def __init__(self, app, name: str) -> None:
-        self.app = app
+    def __init__(self, app: BaseApp, name: str) -> None:
+        self._app = app
         self.name = name
-        self.data = {"read": None, "update": None}
+        self.data: dict[str, Any] = {"read": None, "update": None}
         self.hash = None
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(name={self.name!r})"
 
     @property
-    def logger(self):
+    def logger(self) -> Logger:
         """
         The logger.
         """
-        return self.app.persistence_logger
+        return self._app.persistence_logger
 
     @abstractmethod
-    def _load(self):
+    def _load(self) -> Any:
         pass
 
     @abstractmethod
-    def _dump(self, data):
+    def _dump(self, data: dict | list):
         pass
 
-    def open(self, action):
+    def open(self, action: str):
         """
         Open the store.
 
         Parameters
         ----------
-        action: str
+        action:
             What we are opening the store for; typically either ``read`` or ``update``.
         """
         if action == "read":
@@ -80,17 +88,17 @@ class StoreResource(ABC):
         self.data[action] = data
 
         try:
-            self.hash = get_md5_hash(data)
+            self.hash = get_md5_hash(data)  # type: ignore
         except Exception:
             pass
 
-    def close(self, action):
+    def close(self, action: str):
         """
         Close the store for a particular action.
 
         Parameters
         ----------
-        action: str
+        action:
             What we are closing the store for.
             Should match a previous call to :py:meth:`close`.
         """
@@ -135,24 +143,31 @@ class JSONFileStoreResource(StoreResource):
         The filesystem that the JSON file resides within.
     """
 
-    def __init__(self, app, name: str, filename: str, path: Union[str, Path], fs):
+    def __init__(
+        self,
+        app: BaseApp,
+        name: str,
+        filename: str,
+        path: str | Path,
+        fs: AbstractFileSystem,
+    ):
         self.filename = filename
         self.path = path
         self.fs = fs
         super().__init__(app, name)
 
     @property
-    def _full_path(self):
+    def _full_path(self) -> str:
         return f"{self.path}/{self.filename}"
 
-    def _load(self):
+    def _load(self) -> Any:
         self.logger.debug(f"{self!r}: loading JSON from file.")
         with self.fs.open(self._full_path, mode="rt") as fp:
             return json.load(fp)
 
-    def _dump(self, data):
+    def _dump(self, data: Mapping | list):
         self.logger.debug(f"{self!r}: dumping JSON to file")
-        if "runs" in data:
+        if isinstance(data, dict) and "runs" in data:
             self.logger.debug(f"...runs: {data['runs']}")
         with self.fs.open(self._full_path, mode="wt") as fp:
             json.dump(data, fp, indent=2)
@@ -172,16 +187,16 @@ class ZarrAttrsStoreResource(StoreResource):
         How to actually perform an open on the underlying resource.
     """
 
-    def __init__(self, app, name: str, open_call: Callable):
+    def __init__(self, app: BaseApp, name: str, open_call: Callable[..., zarr.Group]):
         self.open_call = open_call
         super().__init__(app, name)
 
-    def _load(self):
+    def _load(self) -> Any:
         self.logger.debug(f"{self!r}: loading Zarr attributes.")
         item = self.open_call(mode="r")
         return copy.deepcopy(item.attrs.asdict())
 
-    def _dump(self, data):
+    def _dump(self, data: dict | list):
         self.logger.debug(f"{self!r}: dumping Zarr attributes.")
         item = self.open_call(mode="r+")
         item.attrs.put(data)

hpcflow/sdk/persistence/types.py

@@ -0,0 +1,307 @@
+"""
+Types used in type-checking the persistence subsystem.
+"""
+from __future__ import annotations
+from typing import Any, Generic, TypeVar, TYPE_CHECKING
+from typing_extensions import TypedDict, NotRequired, TypeAlias
+
+if TYPE_CHECKING:
+    from .base import StoreTask, StoreElement, StoreElementIter, StoreEAR, StoreParameter
+    from ..core.json_like import JSONDocument
+    from ..core.parameters import ParameterValue
+    from ..core.types import IterableParam
+    from ..typing import DataIndex, ParamSource
+
+#: Bound type variable: :class:`StoreTask`.
+AnySTask = TypeVar("AnySTask", bound="StoreTask")
+#: Bound type variable: :class:`StoreElement`.
+AnySElement = TypeVar("AnySElement", bound="StoreElement")
+#: Bound type variable: :class:`StoreElementITer`.
+AnySElementIter = TypeVar("AnySElementIter", bound="StoreElementIter")
+#: Bound type variable: :class:`StoreEAR`.
+AnySEAR = TypeVar("AnySEAR", bound="StoreEAR")
+#: Bound type variable: :class:`StoreParameter`.
+AnySParameter = TypeVar("AnySParameter", bound="StoreParameter")
+#: Type of possible stored parameters.
+ParameterTypes: TypeAlias = (
+    "ParameterValue | list | tuple | set | dict | int | float | str | None | Any"
+)
+
+
+class File(TypedDict):
+    """
+    Descriptor for file metadata.
+    """
+
+    #: Whether to store the contents.
+    store_contents: bool
+    #: The path to the file.
+    path: str
+
+
+class FileDescriptor(TypedDict):
+    """
+    Descriptor for file metadata.
+    """
+
+    #: Whether this is an input file.
+    is_input: bool
+    #: Whether to store the contents.
+    store_contents: bool
+    #: Where the file will go.
+    dst_path: str
+    #: The path to the file.
+    path: str | None
+    #: Whether to delete the file after processing.
+    clean_up: bool
+    # The contents of the file.
+    contents: NotRequired[str]
+
+
+class LoopDescriptor(TypedDict):
+    """
+    Descriptor for loop metadata.
+    """
+
+    #: The parameters iterated over by the loop.
+    iterable_parameters: dict[str, IterableParam]
+    #: The template data from which the loop was created.
+    loop_template: NotRequired[dict[str, Any]]
+    #: The number of iterations generated by a loop.
+    #: Note that the type is really ``list[tuple[tuple[int, ...], int]]``
+    #: but the persistence implementations don't handle tuples usefully.
+    num_added_iterations: list[list[list[int] | int]]
+    #: The parents of the loop.
+    parents: list[str]
+
+
+# TODO: This type looks familiar...
+class StoreCreationInfo(TypedDict):
+    """
+    Information about the creation of the persistence store.
+    """
+
+    #: Information about the application.
+    app_info: dict[str, Any]
+    #: When the persistence store was created.
+    create_time: str
+    #: The unique identifier for for the store/workflow.
+    id: str
+
+
+class ElemMeta(TypedDict):
+    """
+    The kwargs supported for a StoreElement.
+    """
+
+    #: The ID of the element.
+    id_: int
+    #: The index of the element.
+    index: int
+    #: The index of the element in its element set.
+    es_idx: int
+    #: The indices of the element in the sequences that contain it.
+    seq_idx: dict[str, int]
+    #: The indices of the element's sources.
+    src_idx: dict[str, int]
+    #: The task associated with the element.
+    task_ID: int
+    #: The iteration IDs.
+    iteration_IDs: list[int]
+
+
+class IterMeta(TypedDict):
+    """
+    The kwargs supported for a StoreElementIter.
+    """
+
+    #: The index of the iteration.
+    data_idx: DataIndex
+    #: The EARs associated with the iteration.
+    EAR_IDs: dict[int, list[int]]
+    #: Whether the EARs have been initialised.
+    EARs_initialised: bool
+    #: The ID of the element.
+    element_ID: int
+    #: The loops containing the iteration.
+    loop_idx: dict[str, int]
+    #: The schema parameters being iterated over.
+    schema_parameters: list[str]
+
+
+class RunMeta(TypedDict):
+    """
+    The kwargs supported for StoreEAR.
+    """
+
+    #: The ID of the EAR.
+    id_: int
+    #: The ID of the element iteration containing the EAR.
+    elem_iter_ID: int
+    #: The index of the action that generated the EAR.
+    action_idx: int
+    #: The commands that the EAR will run.
+    commands_idx: list[int]
+    #: The data handled by the EAR.
+    data_idx: DataIndex
+    #: Metadata about the EAR.
+    metadata: Metadata | None
+    #: When the EAR ended, if known.
+    end_time: NotRequired[str | None]
+    #: The exit code of the EAR, if known.
+    exit_code: int | None
+    #: When the EAR started, if known.
+    start_time: NotRequired[str | None]
+    #: Working directory snapshot at start.
+    snapshot_start: dict[str, Any] | None
+    #: Working directory snapshot at end.
+    snapshot_end: dict[str, Any] | None
+    #: The index of the EAR in the submissions.
+    submission_idx: int | None
+    #: Where the EAR is set to run.
+    run_hostname: str | None
+    #: Whether the EAR succeeded, if known.
+    success: bool | None
+    #: Whether the EAR was skipped.
+    skip: bool
+
+
+class TaskMeta(TypedDict):
+    """
+    Information about a task.
+    """
+
+    #: The ID of the task.
+    id_: int
+    #: The index of the task in the workflow.
+    index: int
+    #: The elements in the task.
+    element_IDs: list[int]
+
+
+class TemplateMeta(TypedDict):  # FIXME: Incomplete, see WorkflowTemplate
+    """
+    Metadata about a workflow template.
+    """
+
+    #: Descriptors for loops.
+    loops: list[dict]
+    #: Descriptors for tasks.
+    tasks: list[dict]
+
+
+class Metadata(TypedDict):
+    """
+    Workflow metadata.
+    """
+
+    #: Information about the store's creation.
+    creation_info: NotRequired[StoreCreationInfo]
+    #: Elements in the workflow.
+    elements: NotRequired[list[ElemMeta]]
+    #: Iterations in the workflow.
+    iters: NotRequired[list[IterMeta]]
+    #: Loops in the workflow.
+    loops: NotRequired[list[LoopDescriptor]]
+    #: The name of the workflow.
+    name: NotRequired[str]
+    #: The number of added tasks.
+    num_added_tasks: NotRequired[int]
+    #: The replacement workflow, if any.
+    replaced_workflow: NotRequired[str]
+    #: Element Action Runs in the workflow.
+    runs: NotRequired[list[RunMeta]]
+    #: Tasks in the workflow.
+    tasks: NotRequired[list[TaskMeta]]
+    #: The template that generated the workflow.
+    template: NotRequired[TemplateMeta]
+    #: Custom template components used.
+    template_components: NotRequired[dict[str, Any]]
+    #: Format for timestamps.
+    ts_fmt: NotRequired[str]
+    #: Format for timestamps used in naming.
+    ts_name_fmt: NotRequired[str]
+
+
+class TypeLookup(TypedDict, total=False):
+    """
+    Information for looking up the type of a parameter.
+
+    Note
+    ----
+    Not a total typed dictionary.
+    """
+
+    #: Tuples involving the parameter.
+    tuples: list[list[int]]
+    #: Sets involving the parameter.
+    sets: list[list[int]]
+    #: Arrays involving the parameter.
+    arrays: list[list[list[int] | int]]
+    #: Masked arrays involving the parameter.
+    masked_arrays: list[list[int | list[int]]]
+
+
+class EncodedStoreParameter(TypedDict):
+    """
+    The encoding of a :class:`StoreParameter`.
+    """
+
+    #: The parameter data.
+    data: Any
+    #: Information for looking up the type.
+    type_lookup: TypeLookup
+
+
+class PersistenceCache(
+    TypedDict, Generic[AnySTask, AnySElement, AnySElementIter, AnySEAR, AnySParameter]
+):
+    """
+    Cache used internally by the persistence engine.
+    """
+
+    #: Tasks.
+    tasks: dict[int, AnySTask]
+    #: Elements.
+    elements: dict[int, AnySElement]
+    #: Element iterations.
+    element_iters: dict[int, AnySElementIter]
+    #: Element action runs.
+    EARs: dict[int, AnySEAR]
+    #: Parameter sources.
+    param_sources: dict[int, ParamSource]
+    #: Number of tasks.
+    num_tasks: int | None
+    #: Parameters.
+    parameters: dict[int, AnySParameter]
+    #: Number of element action runs.
+    num_EARs: int | None
+
+
+class ZarrAttrsDict(TypedDict):
+    """
+    Zarr workflow attributes descriptor.
+    """
+
+    #: Workflow name.
+    name: str
+    #: Timestamp format.
+    ts_fmt: str
+    #: Timestamp format for names.
+    ts_name_fmt: str
+    #: Information about the creation of the workflow and persistent store.
+    creation_info: StoreCreationInfo
+    #: The template used to build the workflow.
+    template: TemplateMeta
+    #: Custom components used to build the workflow.
+    template_components: dict[str, Any]
+    #: Number of tasks added.
+    num_added_tasks: int
+    #: Tasks in the workflow.
+    tasks: list[dict[str, Any]]
+    #: Loops in the workflow.
+    loops: list[dict[str, Any]]
+    #: Submissions by the workflow.
+    submissions: list[JSONDocument]
+    #: Replacement workflow, if any.
+    replaced_workflow: NotRequired[str]

hpcflow/sdk/persistence/utils.py

@@ -2,12 +2,22 @@
 Miscellaneous persistence-related helpers.
 """
 
+from __future__ import annotations
 from getpass import getpass
+from typing import TYPE_CHECKING
 
 from hpcflow.sdk.core.errors import WorkflowNotFoundError
 
+if TYPE_CHECKING:
+    from typing import Callable, TypeVar
+    from fsspec import AbstractFileSystem  # type: ignore
 
-def ask_pw_on_auth_exc(f, *args, add_pw_to=None, **kwargs):
+    T = TypeVar("T")
+
+
+def ask_pw_on_auth_exc(
+    f: Callable[..., T], *args, add_pw_to: str | None = None, **kwargs
+) -> tuple[T, str | None]:
     """
     Run the given function on the given arguments and add a password if the function
     fails with an SSHException.
@@ -24,19 +34,14 @@ def ask_pw_on_auth_exc(f, *args, add_pw_to=None, **kwargs):
         if not add_pw_to:
             kwargs["password"] = pw
         else:
-            kwargs[add_pw_to]["password"] = pw
+            kwargs[add_pw_to] = {**kwargs[add_pw_to], "password": pw}
 
         out = f(*args, **kwargs)
 
-        if not add_pw_to:
-            del kwargs["password"]
-        else:
-            del kwargs[add_pw_to]["password"]
-
     return out, pw
 
 
-def infer_store(path: str, fs) -> str:
+def infer_store(path: str, fs: AbstractFileSystem) -> str:
     """Identify the store type using the path and file system parsed by fsspec.
 
     Parameters
@@ -63,8 +68,6 @@ def infer_store(path: str, fs) -> str:
         elif fs.glob(f"{path}/metadata.json"):
             store_fmt = "json"
         else:
-            raise WorkflowNotFoundError(
-                f"Cannot infer a store format at path {path!r} with file system {fs!r}."
-            )
+            raise WorkflowNotFoundError(path, fs)
 
     return store_fmt