sonusai 1.0.16__cp311-abi3-macosx_10_12_x86_64.whl

sonusai/mixture/data_io.py

@@ -0,0 +1,173 @@
+from typing import Any
+
+
+def _get_hdf5_name(location: str, index: str) -> str:
+    from os.path import join
+
+    return join(location, index + ".h5")
+
+
+def _get_pickle_name(location: str, index: str, item: str) -> str:
+    from os.path import join
+
+    return join(location, index, item + ".pkl")
+
+
+def read_hdf5_data(location: str, index: str, items: list[str] | str) -> dict[str, Any]:
+    """Read mixture, target, or noise data from an HDF5 file
+
+    :param location: Location of the file
+    :param index: Mixture, target, or noise index
+    :param items: String(s) of data to retrieve
+    :return: Dictionary of name: data
+    """
+    from os.path import exists
+    from typing import Any
+
+    import h5py
+    import numpy as np
+
+    def _get_dataset(file: h5py.File, d_name: str) -> Any:
+        if d_name in file:
+            data = np.array(file[d_name])
+            if data.size == 1:
+                item = data.item()
+                if isinstance(item, bytes):
+                    return item.decode("utf-8")
+                return item
+            return data
+        return None
+
+    if not isinstance(items, list):
+        items = [items]
+
+    h5_name = _get_hdf5_name(location, index)
+    if exists(h5_name):
+        try:
+            with h5py.File(h5_name, "r") as f:
+                result = {item: _get_dataset(f, item) for item in items}
+        except Exception as e:
+            raise OSError(f"Error reading {h5_name}: {e}") from e
+    else:
+        result = {item: None for item in items}
+
+    return result
+
+
+def write_hdf5_data(location: str, index: str, items: dict[str, Any]) -> None:
+    """Write mixture, target, or noise data to an HDF5 file
+
+    :param location: Location of the file
+    :param index: Mixture, target, or noise index
+    :param items: Dictionary of name: data
+    """
+    import h5py
+
+    h5_name = _get_hdf5_name(location, index)
+    with h5py.File(h5_name, "a") as f:
+        for name, data in items.items():
+            if name in f:
+                del f[name]
+            f.create_dataset(name=name, data=data)
+
+
+def read_pickle_data(location: str, index: str, items: list[str] | str) -> dict[str, Any]:
+    """Read mixture, target, or noise data from a pickle file
+
+    :param location: Location of the file
+    :param index: Mixture, target, or noise index
+    :param items: String(s) of data to retrieve
+    :return: Dictionary of name: data
+    """
+    import pickle
+    from os.path import exists
+    from typing import Any
+
+    if not isinstance(items, list):
+        items = [items]
+
+    result: dict[str, Any] = {}
+    for item in items:
+        pkl_name = _get_pickle_name(location, index, item)
+        if exists(pkl_name):
+            with open(pkl_name, "rb") as f:
+                result[item] = pickle.load(f)  # noqa: S301
+        else:
+            result[item] = None
+
+    return result
+
+
+def write_pickle_data(location: str, index: str, items: dict[str, Any]) -> None:
+    """Write mixture, target, or noise data to a pickle file
+
+    :param location: Location of the file
+    :param index: Mixture, target, or noise index
+    :param items: Dictionary of name: data
+    """
+    import pickle
+    from os import makedirs
+    from os.path import join
+
+    directory = join(location, index)
+    makedirs(directory, exist_ok=True)
+    for name, data in items.items():
+        pkl_name = _get_pickle_name(location, index, name)
+        with open(pkl_name, "wb") as f:
+            f.write(pickle.dumps(data))
+
+
+def clear_pickle_data(location: str, index: str, items: list[str] | str) -> None:
+    """Clear mixture, target, or noise data pickle file
+
+    :param location: Location of the file
+    :param index: Mixture, target, or noise index
+    :param items: String(s) of data to retrieve
+    """
+    from pathlib import Path
+
+    if not isinstance(items, list):
+        items = [items]
+
+    for item in items:
+        Path(_get_pickle_name(location, index, item)).unlink(missing_ok=True)
+
+
+def read_cached_data(location: str, name: str, index: str, items: list[str] | str) -> dict[str, Any]:
+    """Read cached data from a file
+
+    :param location: Location of the mixture database
+    :param name: Data name ('mixture', 'target', or 'noise')
+    :param index: Data index (mixture, target, or noise ID)
+    :param items: String(s) of data to retrieve
+    :return: Data (or tuple of data)
+    """
+    from os.path import join
+
+    return read_pickle_data(join(location, name), index, items)
+
+
+def write_cached_data(location: str, name: str, index: str, items: dict[str, Any]) -> None:
+    """Write data to a file
+
+    :param location: Location of the mixture database
+    :param name: Data name ('mixture', 'target', or 'noise')
+    :param index: Data index (mixture, target, or noise ID)
+    :param items: Dictionary of name: data
+    """
+    from os.path import join
+
+    write_pickle_data(join(location, name), index, items)
+
+
+def clear_cached_data(location: str, name: str, index: str, items: list[str] | str) -> None:
+    """Remove cached data file(s)
+
+    :param location: Location of the mixture database
+    :param name: Data name ('mixture', 'target', or 'noise')
+    :param index: Data index (mixture, target, or noise ID)
+    :param items: String(s) of data to clear
+    """
+    from os.path import join
+
+    clear_pickle_data(join(location, name), index, items)

sonusai/mixture/db.py

@@ -0,0 +1,169 @@
+import contextlib
+import sqlite3
+from os import remove
+from os.path import exists
+from os.path import join
+from os.path import normpath
+from sqlite3 import Connection
+from sqlite3 import Cursor
+from typing import Any
+
+from sonusai import logger_db
+
+
+def db_file(location: str, test: bool = False) -> str:
+    from .constants import MIXDB_NAME
+    from .constants import TEST_MIXDB_NAME
+
+    name = TEST_MIXDB_NAME if test else MIXDB_NAME
+    return normpath(join(location, name))
+
+
+class SQLiteDatabase:
+    """A context manager for SQLite database connections with configurable behavior."""
+
+    # Constants for database configuration
+    READONLY_MODE = "?mode=ro"
+    CONNECTION_TIMEOUT = 20
+
+    def __init__(
+        self,
+        location: str,
+        create: bool = False,
+        readonly: bool = True,
+        test: bool = False,
+        verbose: bool = False,
+    ) -> None:
+        """Initialize SQLite database connection manager.
+
+        :param location: Path to the database file.
+        :param create: If True, create a new database file, overwriting any existing one.
+        :param readonly: If True, open the database in read-only mode.
+        :param test: If True, use the test database path.
+        :param verbose: If True, enable SQL statement logging.
+        """
+        self.location = location
+        self.create = create
+        self.readonly = readonly and not create
+        self.test = test
+        self.verbose = verbose
+        self.con: Connection | None = None
+        self.cur: Cursor | None = None
+        self.db_path = db_file(location, test)
+
+    def __enter__(self) -> Cursor:
+        """Enter the context manager, establishing the database connection.
+
+        :return: A database cursor for executing queries.
+        :raises sqlite3.Error: If the connection fails.
+        """
+        try:
+            self._establish_connection()
+        except Exception:
+            self._close_resources()
+            raise
+
+        if self.cur:
+            self.cur.execute("BEGIN TRANSACTION")
+            return self.cur
+        raise sqlite3.Error("Failed to connect to database")
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Any | None,
+    ) -> None:
+        """Exit the context manager, committing changes if appropriate and closing resources.
+
+        :param exc_type: The exception type, if any.
+        :param exc_val: The exception value, if any.
+        :param exc_tb: The exception traceback, if any.
+        """
+        if self.con:
+            if not self.readonly:
+                if exc_type is None:
+                    # Commit only on successful exit
+                    self.con.commit()
+                else:
+                    # Rollback on exception
+                    self.con.rollback()
+            self._close_resources()
+
+    def _close_resources(self) -> None:
+        """Safely close database cursor and connection resources."""
+        if self.cur:
+            with contextlib.suppress(sqlite3.Error):
+                self.cur.close()
+            self.cur = None
+
+        if self.con:
+            with contextlib.suppress(sqlite3.Error):
+                self.con.close()
+            self.con = None
+
+    def _establish_connection(self) -> None:
+        """Establish a connection to the SQLite database.
+
+        :raises OSError: If the database file doesn't exist and create=False.
+        :raises sqlite3.Error: If connection to the database fails.
+        """
+        self._prepare_db_file(self.db_path)
+        uri = self._build_connection_uri(self.db_path)
+
+        try:
+            self.con = sqlite3.connect(
+                f"file:{uri}",
+                uri=True,
+                timeout=self.CONNECTION_TIMEOUT,
+                isolation_level=None,
+            )
+
+            if self.verbose and self.con:
+                self.con.set_trace_callback(logger_db.debug)
+
+            if not self.readonly:
+                self.con.execute("PRAGMA journal_mode=wal")
+                self.con.execute("PRAGMA synchronous=off")
+                self.con.execute("PRAGMA locking_mode=exclusive")
+                self.con.execute("PRAGMA cache_size=10000")
+                self.con.execute("PRAGMA temp_store=memory")
+                self.con.execute("PRAGMA mmap_size=268435456")  # 256MB
+            else:
+                self.con.execute("PRAGMA synchronous=normal")
+                self.con.execute("PRAGMA cache_size=10000")
+                self.con.execute("PRAGMA temp_store=memory")
+                self.con.execute("PRAGMA mmap_size=268435456")  # 256MB
+
+            self.con.commit()
+
+        except sqlite3.Error as e:
+            raise sqlite3.Error(f"Failed to connect to database: {e}") from e
+
+        self.cur = self.con.cursor()
+
+    def _prepare_db_file(self, db_path: str) -> None:
+        """Prepare the database file based on creation settings.
+
+        :param db_path: Path to the database file.
+        :raises OSError: If the database file doesn't exist and create=False.
+        """
+        if self.create and exists(db_path):
+            remove(db_path)
+
+        if not self.create and not exists(db_path):
+            raise OSError(f"Could not find mixture database in {self.location}")
+
+    def _build_connection_uri(self, db_path: str) -> str:
+        """Build the SQLite connection URI with appropriate options.
+
+        :param db_path: Path to the database file.
+        :return: A properly formatted SQLite URI with appropriate options.
+        """
+        uri = db_path
+
+        # Add readonly mode if needed
+        if self.readonly:
+            uri += self.READONLY_MODE
+
+        return uri

sonusai/mixture/db_datatypes.py

@@ -0,0 +1,92 @@
+from collections import namedtuple
+
+SourceFileRecord = namedtuple(
+    "SourceFileRecord",
+    [
+        "id",
+        "category",
+        "class_indices",
+        "level_type",
+        "name",
+        "samples",
+        "speaker_id",
+    ],
+)
+
+TopRecord = namedtuple(
+    "TopRecord",
+    [
+        "id",
+        "class_balancing",
+        "feature",
+        "mixid_width",
+        "num_classes",
+        "seed",
+        "speaker_metadata_tiers",
+        "textgrid_metadata_tiers",
+        "version",
+    ],
+)
+
+ClassLabelRecord = namedtuple(
+    "ClassLabelRecord",
+    [
+        "id",
+        "label",
+    ],
+)
+
+ClassWeightsThresholdRecord = namedtuple(
+    "ClassWeightsThresholdRecord",
+    [
+        "id",
+        "threshold",
+    ],
+)
+
+ImpulseResponseFileRecord = namedtuple(
+    "ImpulseResponseFileRecord",
+    [
+        "id",
+        "delay",
+        "name",
+    ],
+)
+
+SpectralMaskRecord = namedtuple(
+    "SpectralMaskRecord",
+    [
+        "id",
+        "f_max_width",
+        "f_num",
+        "t_max_percent",
+        "t_max_width",
+        "t_num",
+    ],
+)
+
+SourceRecord = namedtuple(
+    "SourceRecord",
+    [
+        "id",
+        "effects",
+        "file_id",
+        "pre_tempo",
+        "repeat",
+        "snr",
+        "snr_gain",
+        "snr_random",
+        "start",
+    ],
+)
+
+MixtureRecord = namedtuple(
+    "MixtureRecord",
+    [
+        "id",
+        "name",
+        "samples",
+        "spectral_mask_id",
+        "spectral_mask_seed",
+    ],
+)