Update with code from sensAI 1.5.0

Author: opcode81

src/sensai/util/cache.py

@@ -8,14 +8,14 @@
 import sqlite3
 import threading
 import time
-from abc import abstractmethod, ABC
+from abc import ABC, abstractmethod
 from collections import OrderedDict
 from functools import wraps
 from pathlib import Path
-from typing import Any, Callable, Iterator, List, Optional, TypeVar, Generic, Union, Hashable
+from typing import Any, Callable, Generic, Hashable, Iterator, List, Optional, TypeVar, Union
 
 from .hash import pickle_hash
-from .pickle import load_pickle, dump_pickle, setstate
+from .pickle import dump_pickle, load_pickle, setstate
 
 log = logging.getLogger(__name__)
 
@@ -443,22 +443,25 @@ class KeyType(enum.Enum):
         INTEGER = ("LONG", )
 
     def __init__(self, path, table_name="cache", deferred_commit_delay_secs=1.0, key_type: KeyType = KeyType.STRING,
-            max_key_length=255):
+            max_key_length=255, max_entries_for_commit: int = 500):
         """
         :param path: the path to the file that is to hold the SQLite database
         :param table_name: the name of the table to create in the database
         :param deferred_commit_delay_secs: the time frame during which no new data must be added for a pending transaction to be committed
         :param key_type: the type to use for keys; for complex keys (i.e. tuples), use STRING (conversions to string are automatic)
         :param max_key_length: the maximum key length for the case where the key_type can be parametrised (e.g. STRING)
         """
+        os.makedirs(os.path.dirname(path), exist_ok=True)
         self.path = path
         self.conn = SqliteConnectionManager.open_connection(path)
         self.table_name = table_name
         self.max_key_length = 255
         self.key_type = key_type
         self._update_hook = DelayedUpdateHook(self._commit, deferred_commit_delay_secs)
         self._num_entries_to_be_committed = 0
+        self._num_entries_total_changed = 0
         self._conn_mutex = threading.Lock()
+        self._max_entries_for_commit = max_entries_for_commit
 
         cursor = self.conn.cursor()
         cursor.execute(f"SELECT name FROM sqlite_master WHERE type='table';")
@@ -484,9 +487,10 @@ def _key_db_value(self, key):
     def _commit(self):
         self._conn_mutex.acquire()
         try:
-            log.info(f"Committing {self._num_entries_to_be_committed} cache entries to the SQLite database {self.path}")
-            self.conn.commit()
-            self._num_entries_to_be_committed = 0
+            if self._num_entries_to_be_committed > 0:
+                log.info(f"Committing {self._num_entries_to_be_committed} cache entries ({self._num_entries_total_changed} total in session) to the SQLite database {self.path}")
+                self.conn.commit()
+                self._num_entries_to_be_committed = 0
         finally:
             self._conn_mutex.release()
 
@@ -502,11 +506,15 @@ def set(self, key: TKey, value: TValue):
             else:
                 cursor.execute(f"UPDATE {self.table_name} SET cache_value=? WHERE cache_key=?", (pickle.dumps(value), key))
             self._num_entries_to_be_committed += 1
+            self._num_entries_total_changed += 1
             cursor.close()
         finally:
             self._conn_mutex.release()
 
-        self._update_hook.handle_update()
+        if self._num_entries_to_be_committed >= self._max_entries_for_commit:
+            self._commit()
+        else:
+            self._update_hook.handle_update()
 
     def _execute(self, cursor, *query):
         try:
@@ -553,6 +561,12 @@ def iter_items(self):
         finally:
             self._conn_mutex.release()
 
+    def finalise(self):
+        """
+        Commits any not yet commited entries to the database
+        """
+        self._commit()
+
 
 class SqlitePersistentList(PersistentList):
     def __init__(self, path):
@@ -774,6 +788,15 @@ def save(self, path: Union[str, Path], backend="pickle"):
         :param path:
         :param backend: pickle, cloudpickle, or joblib
         """
+        if not hasattr(self, "__getstate__"):
+            log.warning(
+                f"You are persisting an object {self} without __getstate__. "
+                f"This may lead to irrecoverable problems with backwards compatibility! "
+                f"It is highly recommended that you implement __getstate__ for everything you persist (especially for stateless classes, since"
+                f"adding state later on means they won't be able to load and even implementing __setstate__ won't help then). "
+                f"The easiest and recommended way to do this is to also inherit from `sensai.util.pickle.PersistableObject` whenever"
+                f"you inherit from `PickleLoadSaveMixin`"
+            )
         dump_pickle(self, path, backend=backend)
 
     @classmethod

src/sensai/util/cache_mysql.py

@@ -9,32 +9,74 @@
 
 
 class MySQLPersistentKeyValueCache(PersistentKeyValueCache):
+    """
+    Can cache arbitrary values in a MySQL database.
+    The keys are always strings at the database level, i.e. if a key is not a string, it is converted to a string using str().
+    """
 
     class ValueType(enum.Enum):
-        DOUBLE = ("DOUBLE", False)  # (SQL data type, isCachedValuePickled)
+        """
+        The value type to use within the MySQL database.
+        Note that the binary BLOB types can be used for all Python types that can be pickled, so the lack
+        of specific types (e.g. for strings) is not a problem.
+        """
+        # enum values are (SQL data type, isCachedValuePickled)
+        DOUBLE = ("DOUBLE", False)
         BLOB = ("BLOB", True)
+        """
+        for Python data types whose pickled representation is up to 64 KB
+        """
+        MEDIUMBLOB = ("MEDIUMBLOB", True)
+        """
+        for Python data types whose pickled representation is up to 16 MB
+        """
 
-    def __init__(self, host, db, user, pw, value_type: ValueType, table_name="cache", deferred_commit_delay_secs=1.0, in_memory=False):
-        import MySQLdb
-        self.conn = MySQLdb.connect(host=host, database=db, user=user, password=pw)
+    def __init__(self, host: str, db: str, user: str, pw: str, value_type: ValueType, table_name="cache",
+        connect_params: dict | None = None, in_memory=False, max_key_length: int = 255, port=3306):
+        """
+        :param host:
+        :param db:
+        :param user:
+        :param pw:
+        :param value_type: the type of value to store in the cache
+        :param table_name:
+        :param connect_params: additional parameters to pass to the pymysql.connect() function (e.g. ssl, etc.)
+        :param in_memory:
+        :param max_key_length: maximal length of the cache key string (keys are always strings) stored in the DB
+            (i.e. the MySQL type is VARCHAR[max_key_length])
+        :param port: the MySQL server port to connect to
+        """
+        import pymysql
+        if connect_params is None:
+            connect_params = {}
+        self._connect = lambda: pymysql.connect(host=host, database=db, user=user, password=pw, port=port, autocommit=True,
+            **connect_params)
+        self._conn = self._connect()
         self.table_name = table_name
-        self.max_key_length = 255
-        self._update_hook = DelayedUpdateHook(self._commit, deferred_commit_delay_secs)
-        self._num_entries_to_be_committed = 0
+        self.max_key_length = max_key_length
 
         cache_value_sql_type, self.is_cache_value_pickled = value_type.value
 
-        cursor = self.conn.cursor()
+        cursor = self._conn.cursor()
         cursor.execute(f"SHOW TABLES;")
         if table_name not in [r[0] for r in cursor.fetchall()]:
+            log.debug(f"Creating table {table_name}")
             cursor.execute(f"CREATE TABLE {table_name} (cache_key VARCHAR({self.max_key_length}) PRIMARY KEY, "
                            f"cache_value {cache_value_sql_type});")
         cursor.close()
 
         self._in_memory_df = None if not in_memory else self._load_table_to_data_frame()
 
+    def _cursor(self):
+        try:
+            self._conn.ping(reconnect=True)
+        except Exception as e:
+            log.error(f"Error while pinging MySQL server: {e}; Reconnecting ...")
+            self._conn = self._connect()
+        return self._conn.cursor()
+
     def _load_table_to_data_frame(self):
-        df = pd.read_sql(f"SELECT * FROM {self.table_name};", con=self.conn, index_col="cache_key")
+        df = pd.read_sql(f"SELECT * FROM {self.table_name};", con=self._conn, index_col="cache_key")
         if self.is_cache_value_pickled:
             df["cache_value"] = df["cache_value"].apply(pickle.loads)
         return df
@@ -43,16 +85,25 @@ def set(self, key, value):
         key = str(key)
         if len(key) > self.max_key_length:
             raise ValueError(f"Key too long, maximal key length is {self.max_key_length}")
-        cursor = self.conn.cursor()
+        cursor = self._cursor()
         cursor.execute(f"SELECT COUNT(*) FROM {self.table_name} WHERE cache_key=%s", (key,))
         stored_value = pickle.dumps(value) if self.is_cache_value_pickled else value
         if cursor.fetchone()[0] == 0:
-            cursor.execute(f"INSERT INTO {self.table_name} (cache_key, cache_value) VALUES (%s, %s)",
-                (key, stored_value))
+            from pymysql.err import IntegrityError
+            try:
+                cursor.execute(f"INSERT INTO {self.table_name} (cache_key, cache_value) VALUES (%s, %s)",
+                    (key, stored_value))
+            except IntegrityError as e:
+                if e.args[0] == 1062:  # Duplicate entry
+                    # This can only happen when the user is inserting the same value almost simultaneously (race condition)
+                    args = list(e.args)
+                    args[1] = f"{args[1]}; The duplicate entry is due to quasi-simultaneous insertions for the same key; " \
+                              "Check your application logic!"
+                    raise IntegrityError(*args)
+                else:
+                    raise
         else:
             cursor.execute(f"UPDATE {self.table_name} SET cache_value=%s WHERE cache_key=%s", (stored_value, key))
-        self._num_entries_to_be_committed += 1
-        self._update_hook.handle_update()
         cursor.close()
         if self._in_memory_df is not None:
             self._in_memory_df["cache_value"][str(key)] = value
@@ -64,9 +115,10 @@ def get(self, key):
         return value
 
     def _get_from_table(self, key):
-        cursor = self.conn.cursor()
+        cursor = self._cursor()
         cursor.execute(f"SELECT cache_value FROM {self.table_name} WHERE cache_key=%s", (str(key),))
         row = cursor.fetchone()
+        cursor.close()
         if row is None:
             return None
         stored_value = row[0]
@@ -81,8 +133,3 @@ def _get_from_in_memory_df(self, key):
         except Exception as e:
             log.debug(f"Unable to load value for key {str(key)} from in-memory dataframe: {e}")
             return None
-
-    def _commit(self):
-        log.info(f"Committing {self._num_entries_to_be_committed} cache entries to the database")
-        self.conn.commit()
-        self._num_entries_to_be_committed = 0

src/sensai/util/git.py

@@ -22,7 +22,13 @@ def is_clean(self) -> bool:
                     self.has_untracked_files)
 
 
-def git_status() -> Optional[GitStatus]:
+def git_status(log_error: bool = True) -> Optional[GitStatus]:
+    """
+    Gets the git status of the current repository.
+
+    :param log_error: whether to log an error if the git status cannot be determined
+    :return: the git status, or None if it cannot be determined
+    """
     try:
         commit_hash = subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode('ascii').strip()
         unstaged = bool(subprocess.check_output(['git', 'diff', '--name-only']).decode('ascii').strip())
@@ -35,5 +41,6 @@ def git_status() -> Optional[GitStatus]:
             has_untracked_files=untracked
         )
     except Exception as e:
-        log.error("Error determining Git status", exc_info=e)
+        if log_error:
+            log.error("Error determining Git status", exc_info=e)
         return None

src/sensai/util/io.py

@@ -99,6 +99,26 @@ def write_data_frame_csv_file(self, filename_suffix: str, df: "pd.DataFrame", in
             df.to_csv(p, index=index, header=header)
         return p
 
+    def write_data_frame_excel_file(self, filename_suffix: str, df: "pd.DataFrame", index=True, header=True, sheet_name="Sheet1"):
+        """
+        Saves the given data frame to an Excel spreadsheet.
+
+        Requires that an appropriate engine be installed, e.g. openpyxl.
+
+        :param filename_suffix: the filename suffix (or full basename for the case where no prefix was specified at construction),
+            to which the extension .xlsx will be added by default if no extension is given
+        :param df: the data frame to save
+        :param index: whether to save the index
+        :param header: whether to save the header row
+        :param sheet_name: the name of the spreadsheet
+        :return: the path to the saved file
+        """
+        p = self.path(filename_suffix, extension_to_add="xlsx")
+        if self.enabled:
+            self.log.info(f"Saving data frame Excel file {p}")
+            df.to_excel(p, index=index, header=header, sheet_name=sheet_name)
+        return p
+
     def write_figure(self, filename_suffix: str, fig: "plt.Figure", close_figure: Optional[bool] = None):
         """
         :param filename_suffix: the filename suffix, which may or may not include a file extension, valid extensions being {"png", "jpg"}
@@ -220,3 +240,34 @@ def create_file_path(root, *path_elems, make_dirs: bool = False) -> str:
 
 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)
+
+
+def filename_compatible(fn: str, replacement: str = "-") -> str:
+    """
+    Converts the given string into a string that can be used in a filename
+
+    :param fn: original filename, which may or may not be compatible with common filenames
+    :param replacement: default character to use as a replacement
+    :return: adapted filename
+    """
+    # Replace prohibited characters
+    # Windows prohibits: \ / : * ? " < > |
+    # Unix/Linux prohibits: /
+    # macOS prohibits: / and :
+    replacements = {
+        ">": "gt",
+        "<": "lt",
+        "/": replacement,
+        "\\": replacement,
+        ":": replacement,
+        "*": replacement,
+        "?": replacement,
+        "\"": replacement,
+        "|": replacement
+    }
+
+    result = fn
+    for char, repl in replacements.items():
+        result = result.replace(char, repl)
+
+    return result