Update with code from sensAI v1.3.0

Author: opcode81

src/sensai/util/cache.py

@@ -9,9 +9,10 @@
 import threading
 import time
 from abc import abstractmethod, ABC
+from collections import OrderedDict
 from functools import wraps
 from pathlib import Path
-from typing import Any, Callable, Iterator, List, Optional, TypeVar, Generic, Union
+from typing import Any, Callable, Iterator, List, Optional, TypeVar, Generic, Union, Hashable
 
 from .hash import pickle_hash
 from .pickle import load_pickle, dump_pickle, setstate
@@ -20,6 +21,7 @@
 
 T = TypeVar("T")
 TKey = TypeVar("TKey")
+THashableKey = TypeVar("THashableKey", bound=Hashable)
 TValue = TypeVar("TValue")
 TData = TypeVar("TData")
 
@@ -788,3 +790,28 @@ def load(cls, path: Union[str, Path], backend="pickle"):
         if not isinstance(result, cls):
             raise Exception(f"Excepted instance of {cls}, instead got: {result.__class__.__name__}")
         return result
+
+
+class LRUCache(KeyValueCache[THashableKey, TValue], Generic[THashableKey, TValue]):
+    def __init__(self, capacity: int) -> None:
+        self._cache = OrderedDict()
+        self._capacity = capacity
+
+    def get(self, key: THashableKey) -> TValue:
+        if key not in self._cache:
+            return None
+        self._cache.move_to_end(key)
+        return self._cache[key]
+
+    def set(self, key: THashableKey, value: TValue):
+        if key in self._cache:
+            self._cache.move_to_end(key)
+        self._cache[key] = value
+        if len(self._cache) > self._capacity:
+            self._cache.popitem(last=False)
+
+    def __len__(self) -> int:
+        return len(self._cache)
+
+    def clear(self) -> None:
+        self._cache.clear()

src/sensai/util/helper.py

@@ -2,7 +2,7 @@
 This module contains various helper functions.
 """
 import math
-from typing import Any, Sequence, Union, TypeVar, List
+from typing import Any, Sequence, Union, TypeVar, List, Optional, Dict, Container, Iterable
 
 T = TypeVar("T")
 
@@ -58,6 +58,13 @@ def check_not_nan_dict(d: dict):
         raise ValueError(f"Got one or more NaN values: {invalid_keys}")
 
 
+def contains_any(container: Union[Container, Iterable], items: Sequence) -> bool:
+    for item in items:
+        if item in container:
+            return True
+    return False
+
+
 # noinspection PyUnusedLocal
 def mark_used(*args):
     """
@@ -83,4 +90,23 @@ def flatten_arguments(args: Sequence[Union[T, Sequence[T]]]) -> List[T]:
             result.extend(arg)
         else:
             result.append(arg)
-    return result
+    return result
+
+
+def kwarg_if_not_none(arg_name: str, arg_value: Any) -> Dict[str, Any]:
+    """
+    Supports the optional passing of a kwarg, returning a non-empty dictionary with the kwarg only
+    if the value is not None.
+
+    This can be helpful to improve backward compatibility for cases where a kwarg was added later
+    but old implementations weren't updated to have it, such that an exception will be raised only
+    if the kwarg is actually used.
+
+    :param arg_name: the argument name
+    :param arg_value: the value
+    :return: a dictionary containing the kwarg or, if the value is None, an empty dictionary
+    """
+    if arg_value is None:
+        return {}
+    else:
+        return {arg_name: arg_value}

src/sensai/util/io.py

@@ -52,6 +52,9 @@ def path(self, filename_suffix: str, extension_to_add=None, valid_other_extensio
             relevant if extensionToAdd is specified
         :return: the full path
         """
+        # replace forbidden characters
+        filename_suffix = filename_suffix.replace(">=", "gte").replace(">", "gt")
+
         if extension_to_add is not None:
             add_ext = True
             valid_extensions = set(valid_other_extensions) if valid_other_extensions is not None else set()
@@ -201,3 +204,19 @@ def _get_s3_object(self):
         session = boto3.session.Session(profile_name=os.getenv("AWS_PROFILE"))
         s3 = session.resource("s3")
         return s3.Bucket(self.bucket).Object(self.object)
+
+
+def create_path(root: str, *path_elems: str, is_dir: bool, make_dirs: bool = False) -> str:
+    path = os.path.join(root, *path_elems)
+    if make_dirs:
+        dir_path = path if is_dir else os.path.dirname(path)
+        os.makedirs(dir_path, exist_ok=True)
+    return path
+
+
+def create_file_path(root, *path_elems, make_dirs: bool = False) -> str:
+    return create_path(root, *path_elems, is_dir=False, make_dirs=make_dirs)
+
+
+def create_dir_path(root, *path_elems, make_dirs: bool = False) -> str:
+    return create_path(root, *path_elems, is_dir=True, make_dirs=make_dirs)

src/sensai/util/logging.py

@@ -2,10 +2,11 @@
 import logging as lg
 import sys
 import time
+from abc import ABC, abstractmethod
 from datetime import datetime
 from io import StringIO
 from logging import *
-from typing import List, Callable, Optional, TypeVar, TYPE_CHECKING
+from typing import List, Callable, Optional, TypeVar, TYPE_CHECKING, Generic
 
 from .time import format_duration
 
@@ -17,6 +18,7 @@
 
 LOG_DEFAULT_FORMAT = '%(levelname)-5s %(asctime)-15s %(name)s:%(funcName)s:%(lineno)d - %(message)s'
 T = TypeVar("T")
+THandler = TypeVar("THandler", bound=Handler)
 
 # Holds the log format that is configured by the user (using function `configure`), such
 # that it can be reused in other places
@@ -142,18 +144,18 @@ def datetime_tag() -> str:
 
 _fileLoggerPaths: List[str] = []
 _isAtExitReportFileLoggerRegistered = False
-_memoryLogStream: Optional[StringIO] = None
 
 
 def _at_exit_report_file_logger():
     for path in _fileLoggerPaths:
         print(f"A log file was saved to {path}")
 
 
-def add_file_logger(path, register_atexit=True):
+def add_file_logger(path, append=True, register_atexit=True) -> FileHandler:
     global _isAtExitReportFileLoggerRegistered
     log.info(f"Logging to {path} ...")
-    handler = FileHandler(path)
+    mode = "a" if append else "w"
+    handler = FileHandler(path, mode=mode)
     handler.setFormatter(Formatter(_logFormat))
     Logger.root.addHandler(handler)
     _fileLoggerPaths.append(path)
@@ -163,22 +165,24 @@ def add_file_logger(path, register_atexit=True):
     return handler
 
 
-def add_memory_logger() -> None:
+class MemoryStreamHandler(StreamHandler):
+    def __init__(self, stream: StringIO):
+        super().__init__(stream)
+
+    def get_log(self) -> str:
+        stream: StringIO = self.stream
+        return stream.getvalue()
+
+
+def add_memory_logger() -> MemoryStreamHandler:
     """
-    Enables in-memory logging (if it is not already enabled), i.e. all log statements are written to a memory buffer and can later be
-    read via function `get_memory_log()`
+    Adds an in-memory logger, i.e. all log statements are written to a memory buffer which can be retrieved
+    using the handler's `get_log` method.
     """
-    global _memoryLogStream
-    if _memoryLogStream is not None:
-        return
-    _memoryLogStream = StringIO()
-    handler = StreamHandler(_memoryLogStream)
+    handler = MemoryStreamHandler(StringIO())
     handler.setFormatter(Formatter(_logFormat))
     Logger.root.addHandler(handler)
-
-
-def get_memory_log():
-    return _memoryLogStream.getvalue()
+    return handler
 
 
 class StopWatch:
@@ -334,29 +338,62 @@ def __enter__(self):
         return self
 
 
-class FileLoggerContext:
+class LoggerContext(Generic[THandler], ABC):
     """
-    A context handler to be used in conjunction with Python's `with` statement which enables file-based logging.
+    Base class for context handlers to be used in conjunction with Python's `with` statement.
     """
-    def __init__(self, path: str, enabled=True):
+
+    def __init__(self, enabled=True):
         """
-        :param path: the path to the log file
         :param enabled: whether to actually perform any logging.
             This switch allows the with statement to be applied regardless of whether logging shall be enabled.
         """
         self.enabled = enabled
-        self.path = path
         self._log_handler = None
 
-    def __enter__(self):
+    @abstractmethod
+    def _create_log_handler(self) -> THandler:
+        pass
+
+    def __enter__(self) -> Optional[THandler]:
         if self.enabled:
-            self._log_handler = add_file_logger(self.path, register_atexit=False)
+            self._log_handler = self._create_log_handler()
+        return self._log_handler
 
     def __exit__(self, exc_type, exc_value, traceback):
         if self._log_handler is not None:
             remove_log_handler(self._log_handler)
 
 
+class FileLoggerContext(LoggerContext[FileHandler]):
+    """
+    A context handler to be used in conjunction with Python's `with` statement which enables file-based logging.
+    """
+
+    def __init__(self, path: str, append=True, enabled=True):
+        """
+        :param path: the path to the log file
+        :param append: whether to append in case the file already exists; if False, always create a new file.
+        :param enabled: whether to actually perform any logging.
+            This switch allows the with statement to be applied regardless of whether logging shall be enabled.
+        """
+        self.path = path
+        self.append = append
+        super().__init__(enabled=enabled)
+
+    def _create_log_handler(self) -> FileHandler:
+        return add_file_logger(self.path, append=self.append, register_atexit=False)
+
+
+class MemoryLoggerContext(LoggerContext[MemoryStreamHandler]):
+    """
+    A context handler to be used in conjunction with Python's `with` statement which enables in-memory logging.
+    """
+
+    def _create_log_handler(self) -> MemoryStreamHandler:
+        return add_memory_logger()
+
+
 class LoggingDisabledContext:
     """
     A context manager that will temporarily disable logging