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

hpcflow/sdk/core/zarr_io.py

@@ -2,82 +2,71 @@
 Utilities for working with Zarr.
 """
 
-from typing import Any, Dict, Union
+from __future__ import annotations
+from typing import Any
+from typing_extensions import Self
 
-import zarr
+import zarr  # type: ignore
 import numpy as np
 
 from hpcflow.sdk.core.utils import get_in_container, get_relative_path, set_in_container
 
 
-PRIMITIVES = (
+#: The basic types that Zarr can handle directly with no special action.
+PRIMITIVES: tuple[type, ...] = (
     int,
     float,
     str,
     type(None),
 )
 
-
-def _zarr_encode(obj, zarr_group, path=None, encoded=None):
-    path = path or []
-    encoded = encoded or []
-
-    if len(path) > 50:
-        raise RuntimeError("I'm in too deep!")
-
-    if isinstance(obj, ZarrEncodable):
-        obj = obj.to_dict()
-        out, encoded = _zarr_encode(
-            obj, zarr_group=zarr_group, path=path, encoded=encoded
-        )
-
-    elif isinstance(obj, (list, tuple, set)):
-        out = []
-        for idx, item in enumerate(obj):
-            item, encoded = _zarr_encode(item, zarr_group, path + [idx], encoded)
-            out.append(item)
-        if isinstance(obj, tuple):
-            out = tuple(out)
-        elif isinstance(obj, set):
-            out = set(out)
-
-    elif isinstance(obj, dict):
-        out = {}
-        for dct_key, dct_val in obj.items():
-            dct_val, encoded = _zarr_encode(
-                dct_val, zarr_group, path + [dct_key], encoded
-            )
-            out.update({dct_key: dct_val})
-
-    elif isinstance(obj, PRIMITIVES):
-        out = obj
-
-    elif isinstance(obj, np.ndarray):
-        names = [int(i) for i in zarr_group.keys()]
-        if not names:
-            new_name = "0"
-        else:
-            new_name = str(max(names) + 1)
-
-        zarr_group.create_dataset(name=new_name, data=obj)
-        encoded.append(
-            {
-                "path": path,
-                "dataset": new_name,
-            }
-        )
-        out = None
-
-    return out, encoded
+#: Maximum nesting depth for encoding.
+MAX_DEPTH = 50
 
 
-def zarr_encode(data, zarr_group, is_pending_add, is_set):
+def zarr_encode(data, zarr_group: zarr.Group, is_pending_add: bool, is_set: bool):
     """
     Encode data into a zarr group.
     """
-    data, encoded = _zarr_encode(data, zarr_group)
+
+    encoded: list[dict] = []
+
+    def encode(obj: Any, path: list) -> Any:
+        if len(path) > MAX_DEPTH:
+            raise RuntimeError("I'm in too deep!")
+
+        if isinstance(obj, ZarrEncodable):
+            return encode(obj.to_dict(), path)
+        elif isinstance(obj, (list, tuple, set)):
+            out = (encode(item, [*path, idx]) for idx, item in enumerate(obj))
+            if isinstance(obj, tuple):
+                return tuple(out)
+            elif isinstance(obj, set):
+                return set(out)
+            else:
+                return list(out)
+        elif isinstance(obj, dict):
+            return {
+                dct_key: encode(dct_val, [*path, dct_key])
+                for dct_key, dct_val in obj.items()
+            }
+        elif isinstance(obj, PRIMITIVES):
+            return obj
+        elif isinstance(obj, np.ndarray):
+            new_name = str(max((int(i) + 1 for i in zarr_group.keys()), default=0))
+            zarr_group.create_dataset(name=new_name, data=obj)
+            encoded.append(
+                {
+                    "path": path,
+                    "dataset": new_name,
+                }
+            )
+            return None
+        else:
+            raise ValueError(f"unserializable type: {type(obj)}")
+
+    zarr_group.attrs["data"] = encode(data, [])
     zarr_group.attrs["encoded"] = encoded
-    zarr_group.attrs["data"] = data
     zarr_group.attrs["is_set"] = is_set
     if is_pending_add:
         zarr_group.attrs["is_pending_add"] = is_pending_add
@@ -88,100 +77,68 @@ def zarr_encode(data, zarr_group, is_pending_add, is_set):
 
 def _zarr_encode_NEW(
     obj: Any,
-    base_arr: zarr.Array,
     root_group: zarr.Group,
     arr_path: str,
-    path=None,
-    arr_lookup=None,
-):
+) -> tuple[Any, list[list]]:
     """
     Save arbitrarily-nested Python-primitive, `ZarrEncodable` and numpy array objects into
     Zarr.
 
     Parameters
     ----------
-    obj
-    base_arr
-        Zarr object array into which (the top-level) `obj` is saved using the MsgPack
-        encoder.
-    root_group
+    obj:
+        Object to encode.
+    root_group:
         Parent Zarr group into which new Zarr arrays will be added (at `arr_path`).
-    arr_path
+    arr_path:
         Path relative to `root_group` into which new Zarr arrays will be added.
 
     Returns
     -------
-    (data, arr_lookup)
-
-
+    data
+        The encoded data.
+    arr_lookup
+        How to look up where to rebuild Numpy arrays.
     """
 
-    path = path or []
-    arr_lookup = arr_lookup or []
-
-    if len(path) > 50:
-        raise RuntimeError("I'm in too deep!")
-
-    if isinstance(obj, ZarrEncodable):
-        data, arr_lookup = _zarr_encode_NEW(
-            obj=obj.to_dict(),
-            base_arr=base_arr,
-            root_group=root_group,
-            arr_path=arr_path,
-            path=path,
-        )
-
-    elif isinstance(obj, (list, tuple, set)):
-        data = []
-        for idx, item in enumerate(obj):
-            item, arr_lookup = _zarr_encode_NEW(
-                obj=item,
-                base_arr=base_arr,
-                root_group=root_group,
-                arr_path=arr_path,
-                path=path + [idx],
-            )
-            data.append(item)
-        if isinstance(obj, tuple):
-            data = tuple(data)
-        elif isinstance(obj, set):
-            data = set(data)
-
-    elif isinstance(obj, dict):
-        data = {}
-        for dct_key, dct_val in obj.items():
-            dct_val, arr_lookup = _zarr_encode_NEW(
-                obj=dct_val,
-                base_arr=base_arr,
-                root_group=root_group,
-                arr_path=arr_path,
-                path=path + [dct_key],
-            )
-            data[dct_key] = dct_val
-
-    elif isinstance(obj, PRIMITIVES):
-        data = obj
-
-    elif isinstance(obj, np.ndarray):
-        # Might need to generate new group:
-        param_arr_group = root_group.require_group(arr_path)
-        names = [int(i) for i in param_arr_group.keys()]
-        if not names:
-            new_idx = 0
+    arr_lookup: list[list] = []
+
+    def encode(obj: Any, path: list) -> Any:
+        if len(path) > MAX_DEPTH:
+            raise RuntimeError("I'm in too deep!")
+
+        if isinstance(obj, ZarrEncodable):
+            return encode(obj.to_dict(), path)
+        elif isinstance(obj, (list, tuple, set)):
+            items = (encode(item, [*path, idx]) for idx, item in enumerate(obj))
+            if isinstance(obj, tuple):
+                return tuple(items)
+            elif isinstance(obj, set):
+                return set(items)
+            else:
+                return list(items)
+        elif isinstance(obj, dict):
+            return {key: encode(val, [*path, key]) for key, val in obj.items()}
+        elif isinstance(obj, PRIMITIVES):
+            return obj
+        elif isinstance(obj, np.ndarray):
+            # Might need to generate new group:
+            param_arr_group = root_group.require_group(arr_path)
+            new_idx = max((int(i) + 1 for i in param_arr_group.keys()), default=0)
+            param_arr_group.create_dataset(name=f"arr_{new_idx}", data=obj)
+            arr_lookup.append([path, new_idx])
+            return None
         else:
-            new_idx = max(names) + 1
-        param_arr_group.create_dataset(name=f"arr_{new_idx}", data=obj)
-        arr_lookup.append([path, new_idx])
-        data = None
+            raise ValueError(f"unserializable type: {type(obj)}")
 
-    return data, arr_lookup
+    return encode(obj, []), arr_lookup
 
 
 def zarr_decode(
-    param_data: Union[None, Dict],
+    param_data: None | dict,
     arr_group: zarr.Group,
-    path=None,
-    dataset_copy=False,
+    path: list | None = None,
+    dataset_copy: bool = False,
 ):
     """
     Decode data from a zarr group.
@@ -220,7 +177,7 @@ class ZarrEncodable:
 
     _typ = None
 
-    def to_dict(self):
+    def to_dict(self) -> dict[str, Any]:
         """
         Convert this object to a dict.
         """
@@ -228,18 +185,22 @@ class ZarrEncodable:
             return dict(self.__dict__)
         elif hasattr(self, "__slots__"):
             return {k: getattr(self, k) for k in self.__slots__}
+        else:
+            # Should be unreachable
+            return {}
 
-    def to_zarr(self, zarr_group):
+    def to_zarr(self, zarr_group: zarr.Group):
         """
         Save this object into the given zarr group.
         """
-        data = self.to_dict()
-        zarr_encode(data, zarr_group)
+        zarr_encode(self.to_dict(), zarr_group, is_pending_add=False, is_set=False)
 
     @classmethod
-    def from_zarr(cls, zarr_group, dataset_copy=False):
+    def from_zarr(cls, zarr_group: zarr.Group, dataset_copy: bool = False) -> Self:
         """
         Read an instance of this class from the given zarr group.
         """
-        data = zarr_decode(zarr_group, dataset_copy=dataset_copy)
+        # FIXME: Do the read of the data!
+        param_data = None
+        data = zarr_decode(param_data, zarr_group, dataset_copy=dataset_copy)
         return cls(**data)

hpcflow/sdk/demo/cli.py

@@ -1,5 +1,10 @@
+"""
+CLI components for demonstration code.
+"""
+from __future__ import annotations
 from pathlib import Path
 from random import randint
+from typing import TYPE_CHECKING
 import click
 
 from hpcflow.sdk.core.utils import get_process_stamp
@@ -22,8 +27,13 @@ from hpcflow.sdk.cli_common import (
     make_status_opt,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+    from typing import Literal
+    from ..app import BaseApp
 
-def get_demo_software_CLI(app):
+
+def get_demo_software_CLI(app: BaseApp):
     """Generate the CLI to provide an example software."""
 
     @click.group()
@@ -35,7 +45,9 @@ def get_demo_software_CLI(app):
     @click.option("--infile2", "-i2", type=click.Path(exists=True), required=True)
     @click.option("--value", "-v")
     @click.option("--out", "-o")
-    def demo_do_something(infile1, infile2, value=None, out=None):
+    def demo_do_something(
+        infile1: Path, infile2: Path, value: str | None = None, out: str | None = None
+    ):
         click.echo("trying to do something")
 
         with Path(infile1).open("r") as handle:
@@ -47,7 +59,7 @@ def get_demo_software_CLI(app):
             out = "outfile.txt"
         out_path = Path(out)
         with out_path.open("a") as handle:
-            handle.write("{}\n".format(randint(0, 1e6)))
+            handle.write("{}\n".format(randint(0, int(1e6))))
             handle.write(
                 "{} Generated by `doSomething --infile1 {} --infile2 {}`.\n".format(
                     get_process_stamp(), infile1, infile2
@@ -69,10 +81,10 @@ def get_demo_software_CLI(app):
     return demo_software
 
 
-def get_demo_workflow_CLI(app):
+def get_demo_workflow_CLI(app: BaseApp):
     """Generate the CLI to provide access to builtin demo workflows."""
 
-    def list_callback(ctx, param, value):
+    def list_callback(ctx: click.Context, param, value: bool):
         if not value or ctx.resilient_parsing:
             return
         # TODO: format with Rich with a one-line description
@@ -105,16 +117,16 @@ def get_demo_workflow_CLI(app):
     @variables_option
     @make_status_opt
     def make_demo_workflow(
-        workflow_name,
-        format,
-        path,
-        name,
-        overwrite,
-        store,
-        ts_fmt=None,
-        ts_name_fmt=None,
-        variables=None,
-        status=True,
+        workflow_name: str,
+        format: Literal["json", "yaml"] | None,
+        path: Path | None,
+        name: str | None,
+        overwrite: bool,
+        store: str,
+        ts_fmt: str | None = None,
+        ts_name_fmt: str | None = None,
+        variables: Iterable[tuple[str, str]] = (),
+        status: bool = True,
     ):
         wk = app.make_demo_workflow(
             workflow_name=workflow_name,
@@ -148,22 +160,22 @@ def get_demo_workflow_CLI(app):
     @cancel_opt
     @submit_status_opt
     def make_and_submit_demo_workflow(
-        workflow_name,
-        format,
-        path,
-        name,
-        overwrite,
-        store,
-        ts_fmt=None,
-        ts_name_fmt=None,
-        variables=None,
-        js_parallelism=None,
-        wait=False,
-        add_to_known=True,
-        print_idx=False,
-        tasks=None,
-        cancel=False,
-        status=True,
+        workflow_name: str,
+        format: Literal["json", "yaml"] | None,
+        path: Path | None,
+        name: str | None,
+        overwrite: bool,
+        store: str,
+        ts_fmt: str | None = None,
+        ts_name_fmt: str | None = None,
+        variables: Iterable[tuple[str, str]] = (),
+        js_parallelism: bool | None = None,
+        wait: bool = False,
+        add_to_known: bool = True,
+        print_idx: bool = False,
+        tasks: list[int] | None = None,
+        cancel: bool = False,
+        status: bool = True,
     ):
         out = app.make_and_submit_demo_workflow(
             workflow_name=workflow_name,
@@ -184,20 +196,21 @@ def get_demo_workflow_CLI(app):
             status=status,
         )
         if print_idx:
+            assert isinstance(out, tuple)
             click.echo(out[1])
 
     @demo_workflow.command("copy")
     @click.argument("workflow_name")
     @click.argument("destination")
     @click.option("--doc/--no-doc", default=True)
-    def copy_demo_workflow(workflow_name, destination, doc):
+    def copy_demo_workflow(workflow_name: str, destination: str, doc: bool):
         app.copy_demo_workflow(name=workflow_name, dst=destination, doc=doc)
 
     @demo_workflow.command("show")
     @click.argument("workflow_name")
     @click.option("--syntax/--no-syntax", default=True)
     @click.option("--doc/--no-doc", default=True)
-    def show_demo_workflow(workflow_name, syntax, doc):
+    def show_demo_workflow(workflow_name: str, syntax: bool, doc: bool):
         app.show_demo_workflow(workflow_name, syntax=syntax, doc=doc)
 
     return demo_workflow

hpcflow/sdk/helper/cli.py

@@ -1,13 +1,12 @@
 """
 Common Click command line options related to the helper.
 """
-
-from datetime import timedelta
+from __future__ import annotations
+from typing import TYPE_CHECKING
 
 import click
-import psutil
 
-from .helper import (
+from hpcflow.sdk.helper.helper import (
     DEFAULT_TIMEOUT,
     DEFAULT_TIMEOUT_CHECK,
     DEFAULT_WATCH_INTERVAL,
@@ -22,7 +21,10 @@ from .helper import (
     get_helper_PID,
     get_helper_uptime,
 )
-from ..cli_common import _add_doc_from_help
+from hpcflow.sdk.cli_common import _add_doc_from_help
+
+if TYPE_CHECKING:
+    from ..app import BaseApp
 
 #: Helper option: ``--timeout``
 timeout_option = click.option(
@@ -54,7 +56,7 @@ watch_interval_option = click.option(
 _add_doc_from_help(timeout_option, timeout_check_interval_option, watch_interval_option)
 
 
-def get_helper_CLI(app):
+def get_helper_CLI(app: BaseApp):
     """Generate the CLI to provide some server-like functionality."""
 
     @click.group()
@@ -65,7 +67,7 @@ def get_helper_CLI(app):
     @timeout_option
     @timeout_check_interval_option
     @watch_interval_option
-    def start(timeout, timeout_check_interval, watch_interval):
+    def start(timeout: float, timeout_check_interval: float, watch_interval: float):
         """Start the helper process."""
         start_helper(app, timeout, timeout_check_interval, watch_interval)
 
@@ -78,7 +80,7 @@ def get_helper_CLI(app):
     @timeout_option
     @timeout_check_interval_option
     @watch_interval_option
-    def run(timeout, timeout_check_interval, watch_interval):
+    def run(timeout: float, timeout_check_interval: float, watch_interval: float):
         """Run the helper functionality."""
         run_helper(app, timeout, timeout_check_interval, watch_interval)
 
@@ -86,13 +88,13 @@ def get_helper_CLI(app):
     @timeout_option
     @timeout_check_interval_option
     @watch_interval_option
-    def restart(timeout, timeout_check_interval, watch_interval):
+    def restart(timeout: float, timeout_check_interval: float, watch_interval: float):
         """Restart (or start) the helper process."""
         restart_helper(app, timeout, timeout_check_interval, watch_interval)
 
     @helper.command()
     @click.option("-f", "--file", is_flag=True)
-    def pid(file):
+    def pid(file: bool):
         """Get the process ID of the running helper, if running."""
         pid_info = get_helper_PID(app)
         if pid_info:
@@ -103,32 +105,32 @@ def get_helper_CLI(app):
                 click.echo(pid)
 
     @helper.command()
-    def clear():
+    def clear() -> None:
         """Remove the PID file (and kill the helper process if it exists). This should not
         normally be needed."""
         clear_helper(app)
 
     @helper.command()
-    def uptime():
+    def uptime() -> None:
         """Get the uptime of the helper process, if it is running."""
         out = get_helper_uptime(app)
         if out:
             click.echo(out)
 
     @helper.command()
-    def log_path():
+    def log_path() -> None:
         """Get the path to the helper log file (may not exist)."""
         click.echo(get_helper_log_path(app))
 
     @helper.command()
-    def watch_list_path():
+    def watch_list_path() -> None:
         """Get the path to the workflow watch list file (may not exist)."""
         click.echo(get_watcher_file_path(app))
 
     @helper.command()
-    def watch_list():
+    def watch_list() -> None:
         """Get the list of workflows currently being watched."""
-        for wk in get_helper_watch_list(app) or []:
+        for wk in get_helper_watch_list(app) or ():
             click.echo(str(wk["path"]))
 
     return helper