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

hpcflow/sdk/log.py

@@ -2,12 +2,40 @@
 Interface to the standard logger, and performance logging utility.
 """
 
+from __future__ import annotations
 from functools import wraps
 import logging
+import logging.handlers
 from pathlib import Path
 import time
 from collections import defaultdict
+from collections.abc import Callable, Sequence
 import statistics
+from dataclasses import dataclass
+from typing import ClassVar, TypeVar, TYPE_CHECKING
+from typing_extensions import ParamSpec
+
+if TYPE_CHECKING:
+    from .app import BaseApp
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+@dataclass
+class _Summary:
+    """
+    Summary of a particular node's execution time.
+    """
+
+    number: int
+    mean: float
+    stddev: float
+    min: float
+    max: float
+    sum: float
+    children: dict[tuple[str, ...], _Summary]
 
 
 class TimeIt:
@@ -16,29 +44,39 @@ class TimeIt:
     """
 
     #: Whether the instrumentation is active.
-    active = False
+    active: ClassVar = False
     #: Where to log to.
-    file_path = None
+    file_path: ClassVar[str | None] = None
     #: The details be tracked.
-    timers = defaultdict(list)
+    timers: ClassVar[dict[tuple[str, ...], list[float]]] = defaultdict(list)
     #: Traces of the stack.
-    trace = []
+    trace: ClassVar[list[str]] = []
     #: Trace indices.
-    trace_idx = []
+    trace_idx: ClassVar[list[int]] = []
     #: Preceding traces.
-    trace_prev = []
+    trace_prev: ClassVar[list[str]] = []
     #: Preceding trace indices.
-    trace_idx_prev = []
+    trace_idx_prev: ClassVar[list[int]] = []
+
+    def __enter__(self):
+        self.__class__.active = True
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        try:
+            self.__class__.summarise_string()
+        finally:
+            self.__class__.reset()
+            self.__class__.active = False
 
     @classmethod
-    def decorator(cls, func):
+    def decorator(cls, func: Callable[P, T]) -> Callable[P, T]:
         """
         Decorator for a method that is to have its execution time monitored.
         """
 
         @wraps(func)
-        def wrapper(*args, **kwargs):
-
+        def wrapper(*args, **kwargs) -> T:
             if not cls.active:
                 return func(*args, **kwargs)
 
@@ -68,48 +106,45 @@ class TimeIt:
         return wrapper
 
     @classmethod
-    def summarise(cls):
+    def _summarise(cls) -> dict[tuple[str, ...], _Summary]:
         """
         Produce a machine-readable summary of method execution time statistics.
         """
         stats = {
-            k: {
-                "number": len(v),
-                "mean": statistics.mean(v),
-                "stddev": statistics.pstdev(v),
-                "min": min(v),
-                "max": max(v),
-                "sum": sum(v),
-            }
+            k: _Summary(
+                len(v),
+                statistics.mean(v),
+                statistics.pstdev(v),
+                min(v),
+                max(v),
+                sum(v),
+                {},
+            )
             for k, v in cls.timers.items()
         }
 
         # make a graph
-        for k in stats:
-            stats[k]["children"] = {}
-
-        for key in sorted(stats.keys(), key=lambda x: len(x), reverse=True):
+        for key in sorted(stats, key=lambda x: len(x), reverse=True):
             if len(key) == 1:
                 continue
             value = stats.pop(key)
-            parent = key[:-1]
-            for other_key in stats.keys():
-                if other_key == parent:
-                    stats[other_key]["children"][key] = value
-                    break
+            parent_key = key[:-1]
+            if parent_key in stats:
+                stats[parent_key].children[key] = value
 
         return stats
 
     @classmethod
-    def summarise_string(cls):
+    def summarise_string(cls) -> None:
         """
         Produce a human-readable summary of method execution time statistics.
         """
 
-        def _format_nodes(node, depth=0, depth_final=None):
-            if depth_final is None:
-                depth_final = []
-            out = []
+        def _format_nodes(
+            node: dict[tuple[str, ...], _Summary],
+            depth: int = 0,
+            depth_final: Sequence[bool] = (),
+        ):
             for idx, (k, v) in enumerate(node.items()):
                 is_final_child = idx == len(node) - 1
                 angle = "└ " if is_final_child else "├ "
@@ -117,38 +152,40 @@ class TimeIt:
                 if depth > 0:
                     bars = "".join(f"{'│ ' if not i else '  '}" for i in depth_final)
                 k_str = bars + (angle if depth > 0 else "") + f"{k[depth]}"
-                number = v["number"]
-                min_str = f"{v['min']:10.6f}" if number > 1 else f"{f'-':^12s}"
-                max_str = f"{v['max']:10.6f}" if number > 1 else f"{f'-':^12s}"
-                stddev_str = f"({v['stddev']:8.6f})" if number > 1 else f"{f' ':^10s}"
+                min_str = f"{v.min:10.6f}" if v.number > 1 else f"{f'-':^12s}"
+                max_str = f"{v.max:10.6f}" if v.number > 1 else f"{f'-':^12s}"
+                stddev_str = f"({v.stddev:8.6f})" if v.number > 1 else f"{f' ':^10s}"
                 out.append(
-                    f"{k_str:.<80s} {v['sum']:12.6f} "
-                    f"{v['mean']:10.6f} {stddev_str} {number:8d} "
+                    f"{k_str:.<80s} {v.sum:12.6f} "
+                    f"{v.mean:10.6f} {stddev_str} {v.number:8d} "
                     f"{min_str} {max_str} "
                 )
-                depth_final_next = list(depth_final) + (
-                    [is_final_child] if depth > 0 else []
-                )
-                out.extend(
-                    _format_nodes(
-                        v["children"], depth=depth + 1, depth_final=depth_final_next
-                    )
-                )
-            return out
+                depth_final_next = list(depth_final)
+                if depth > 0:
+                    depth_final_next.append(is_final_child)
+                _format_nodes(v.children, depth + 1, depth_final_next)
 
-        summary = cls.summarise()
+        summary = cls._summarise()
 
         out = [
             f"{'function':^80s} {'sum /s':^12s} {'mean (stddev) /s':^20s} {'N':^8s} "
             f"{'min /s':^12s} {'max /s':^12s}"
         ]
-        out += _format_nodes(summary)
+        _format_nodes(summary)
         out_str = "\n".join(out)
         if cls.file_path:
             Path(cls.file_path).write_text(out_str, encoding="utf-8")
         else:
             print(out_str)
 
+    @classmethod
+    def reset(cls):
+        cls.timers = defaultdict(list)
+        cls.trace = []
+        cls.trace_idx = []
+        cls.trace_prev = []
+        cls.trace_idx_prev = []
+
 
 class AppLog:
     """
@@ -156,22 +193,23 @@ class AppLog:
     """
 
     #: Default logging level for the console.
-    DEFAULT_LOG_CONSOLE_LEVEL = "WARNING"
+    DEFAULT_LOG_CONSOLE_LEVEL: ClassVar = "WARNING"
     #: Default logging level for log files.
-    DEFAULT_LOG_FILE_LEVEL = "INFO"
+    DEFAULT_LOG_FILE_LEVEL: ClassVar = "WARNING"
 
-    def __init__(self, app, log_console_level=None):
+    def __init__(self, app: BaseApp, log_console_level: str | None = None) -> None:
         #: The application context.
-        self.app = app
+        self._app = app
         #: The base logger for the application.
         self.logger = logging.getLogger(app.package_name)
         self.logger.setLevel(logging.DEBUG)
         #: The handler for directing logging messages to the console.
-        self.console_handler = self._add_console_logger(
+        self.console_handler = self.__add_console_logger(
             level=log_console_level or AppLog.DEFAULT_LOG_CONSOLE_LEVEL
         )
+        self.file_handler: logging.FileHandler | None = None
 
-    def _add_console_logger(self, level, fmt=None):
+    def __add_console_logger(self, level: str, fmt: str | None = None) -> logging.Handler:
         fmt = fmt or "%(levelname)s %(name)s: %(message)s"
         handler = logging.StreamHandler()
         handler.setFormatter(logging.Formatter(fmt))
@@ -179,18 +217,30 @@ class AppLog:
         self.logger.addHandler(handler)
         return handler
 
-    def update_console_level(self, new_level):
+    def update_console_level(self, new_level: str | None = None) -> None:
         """
         Set the logging level for console messages.
         """
-        if new_level:
-            self.console_handler.setLevel(new_level.upper())
-
-    def add_file_logger(self, path, level=None, fmt=None, max_bytes=None):
+        new_level = new_level or AppLog.DEFAULT_LOG_CONSOLE_LEVEL
+        self.console_handler.setLevel(new_level.upper())
+
+    def update_file_level(self, new_level: str | None = None) -> None:
+        if self.file_handler:
+            new_level = new_level or AppLog.DEFAULT_LOG_FILE_LEVEL
+            self.file_handler.setLevel(new_level.upper())
+
+    def add_file_logger(
+        self,
+        path: str | Path,
+        level: str | None = None,
+        fmt: str | None = None,
+        max_bytes: int | None = None,
+    ) -> None:
         """
         Add a log file.
         """
-        fmt = fmt or f"%(asctime)s %(levelname)s %(name)s: %(message)s"
+        path = Path(path)
+        fmt = fmt or "%(asctime)s %(levelname)s %(name)s: %(message)s"
         level = level or AppLog.DEFAULT_LOG_FILE_LEVEL
         max_bytes = max_bytes or int(10e6)
 
@@ -202,12 +252,13 @@ class AppLog:
         handler.setFormatter(logging.Formatter(fmt))
         handler.setLevel(level.upper())
         self.logger.addHandler(handler)
-        return handler
-
-    def remove_file_handlers(self):
-        """Remove all file handlers."""
-        # TODO: store a `file_handlers` attribute as well as `console_handlers`
-        for hdlr in self.logger.handlers:
-            if isinstance(hdlr, logging.FileHandler):
-                self.logger.debug(f"Removing file handler from the AppLog: {hdlr!r}.")
-                self.logger.removeHandler(hdlr)
+        self.file_handler = handler
+
+    def remove_file_handler(self) -> None:
+        """Remove the file handler."""
+        if self.file_handler:
+            self.logger.debug(
+                f"Removing file handler from the AppLog: {self.file_handler!r}."
+            )
+            self.logger.removeHandler(self.file_handler)
+            self.file_handler = None

hpcflow/sdk/persistence/__init__.py

@@ -2,40 +2,17 @@
 Workflow persistence subsystem.
 """
 
-import copy
-from pathlib import Path
-import random
-import string
-import time
-from typing import Type, Union
+from __future__ import annotations
+from typing import TYPE_CHECKING
 
-from reretry import retry
+if TYPE_CHECKING:
+    from .base import PersistentStore
 
-from hpcflow.sdk.persistence.base import PersistentStore
-from hpcflow.sdk.persistence.json import JSONPersistentStore
-from hpcflow.sdk.persistence.zarr import ZarrPersistentStore, ZarrZipPersistentStore
 
-_ALL_STORE_CLS = {
-    "zarr": ZarrPersistentStore,
-    "zip": ZarrZipPersistentStore,
-    "json": JSONPersistentStore,
-    # "json-single": JSONPersistentStore,  # TODO
-}
-#: The name of the default persistence store.
-DEFAULT_STORE_FORMAT = "zarr"
-#: The persistence formats supported.
-ALL_STORE_FORMATS = tuple(_ALL_STORE_CLS.keys())
-#: The persistence formats supported for creation.
-ALL_CREATE_STORE_FORMATS = tuple(
-    k for k, v in _ALL_STORE_CLS.items() if v._features.create
-)
-
-
-def store_cls_from_str(store_format: str) -> Type[PersistentStore]:
+def store_cls_from_str(store_format: str) -> type[PersistentStore]:
     """
     Get the class that implements the persistence store from its name.
     """
-    try:
-        return _ALL_STORE_CLS[store_format]
-    except KeyError:
-        raise ValueError(f"Store format {store_format!r} not known.")
+    from .discovery import store_cls_from_str as impl
+
+    return impl(store_format)