diff --git a/src/osekit/core/audio_file.py b/src/osekit/core/audio_file.py index f35cad54..10b0b765 100644 --- a/src/osekit/core/audio_file.py +++ b/src/osekit/core/audio_file.py @@ -5,6 +5,9 @@ import typing from typing import TYPE_CHECKING +from osekit.config import TIMESTAMP_FORMATS_EXPORTED_FILES +from osekit.utils.timestamp import strptime_from_text + if TYPE_CHECKING: from os import PathLike from pathlib import Path @@ -30,6 +33,7 @@ def __init__( begin: Timestamp | None = None, strptime_format: str | list[str] | None = None, timezone: str | pytz.timezone | None = None, + **kwargs: dict, ) -> None: """Initialize an ``AudioFile`` object with a path and a begin timestamp. @@ -55,7 +59,8 @@ def __init__( If different from a timezone parsed from the filename, the timestamps' timezone will be converted from the parsed timezone to the specified timezone. - + kwargs: dict + Audio file info that might bypass the afm.info() call on deserialization. """ super().__init__( path=path, @@ -63,13 +68,25 @@ def __init__( strptime_format=strptime_format, timezone=timezone, ) - sample_rate, frames, channels = afm.info(path) - duration = frames / sample_rate + sample_rate, channels, end = self._get_info(path=path, kwargs=kwargs) self.sample_rate = sample_rate self.channels = channels - self.end = self.begin + Timedelta(seconds=duration) + self.end = end self._check_validity() + def _get_info(self, path: Path, kwargs: dict) -> tuple[int, int, Timestamp]: + keys = ["sample_rate", "channels", "end"] + if all(key in kwargs for key in keys): + end = strptime_from_text( + text=kwargs["end"], + datetime_template=TIMESTAMP_FORMATS_EXPORTED_FILES, + ) + return kwargs["sample_rate"], kwargs["channels"], end + sample_rate, frames, channels = afm.info(path=path) + duration = frames / sample_rate + end = self.begin + Timedelta(seconds=duration) + return sample_rate, channels, end + def _check_validity(self) -> None: """Raise an error if the audio file is not valid.""" if not self.duration: @@ -169,3 +186,17 @@ def stream(self, chunk_size: int) -> np.ndarray: if data.ndim == 1: return data[:, None] # 2D array to match the format of multichannel audio return data + + def to_dict(self) -> dict: + """Serialize an ``AudioFile`` to a dictionary. + + Returns + ------- + dict: + The serialized dictionary representing the ``AudioFile``. + + """ + return super().to_dict() | { + "channels": self.channels, + "sample_rate": self.sample_rate, + } diff --git a/src/osekit/core/base_file.py b/src/osekit/core/base_file.py index f3df8fc4..5e9fe399 100644 --- a/src/osekit/core/base_file.py +++ b/src/osekit/core/base_file.py @@ -45,6 +45,7 @@ def __init__( end: Timestamp | None = None, strptime_format: str | list[str] | None = None, timezone: str | pytz.timezone | None = None, + **kwargs: dict, ) -> None: """Initialize a File object with a path and timestamps. @@ -145,13 +146,15 @@ def from_dict(cls: type[Self], serialized: dict) -> type[Self]: The deserialized File object. """ - path = serialized["path"] + path = serialized.pop("path") + begin = strptime_from_text( + text=serialized.pop("begin"), + datetime_template=TIMESTAMP_FORMATS_EXPORTED_FILES, + ) return cls( path=path, - begin=strptime_from_text( - text=serialized["begin"], - datetime_template=TIMESTAMP_FORMATS_EXPORTED_FILES, - ), + begin=begin, + **serialized, ) def __hash__(self) -> int: diff --git a/src/osekit/core/spectro_file.py b/src/osekit/core/spectro_file.py index 041e6db8..6ea96335 100644 --- a/src/osekit/core/spectro_file.py +++ b/src/osekit/core/spectro_file.py @@ -36,6 +36,7 @@ def __init__( begin: Timestamp | None = None, strptime_format: str | list[str] | None = None, timezone: str | pytz.timezone | None = None, + **kwargs: dict, ) -> None: """Initialize a ``SpectroFile`` object from a ``path`` and begin timestamp. diff --git a/tests/test_serialization.py b/tests/test_serialization.py index c9ddb846..6d2ba6df 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -1,6 +1,7 @@ from __future__ import annotations from pathlib import Path, PureWindowsPath +from typing import Any import numpy as np import pytest @@ -8,6 +9,8 @@ from scipy.signal import ShortTimeFFT from scipy.signal.windows import hamming, hann +import osekit.core +from osekit.audio_backend.audio_file_manager import AudioFileManager from osekit.config import ( TIMESTAMP_FORMAT_EXPORTED_FILES_LOCALIZED, TIMESTAMP_FORMAT_EXPORTED_FILES_UNLOCALIZED, @@ -15,6 +18,7 @@ ) from osekit.core.audio_data import AudioData from osekit.core.audio_dataset import AudioDataset +import osekit.core.audio_file from osekit.core.audio_file import AudioFile from osekit.core.frequency_scale import Scale, ScalePart from osekit.core.instrument import Instrument @@ -29,6 +33,60 @@ from osekit.utils.audio import Normalization +def test_audio_file_from_dict_depends_on_available_info( + audio_files: tuple[list[AudioFile], Any], monkeypatch: pytest.MonkeyPatch +) -> None: + audio_files, _ = audio_files + af = audio_files[0] + minimum = { + "begin": af.begin.strftime(TIMESTAMP_FORMAT_EXPORTED_FILES_LOCALIZED), + "end": af.end.strftime(TIMESTAMP_FORMAT_EXPORTED_FILES_LOCALIZED), + "path": af.path, + } + full = minimum | { + "sample_rate": af.sample_rate, + "channels": af.channels, + } + + afm_calls = [0] + + info_original = AudioFileManager.info + afm = osekit.core.audio_file.afm + + def patch_afm_info(*args, **kwargs) -> tuple[int, int, int]: + afm_calls[0] += 1 + return info_original(self=afm, *args, **kwargs) + + monkeypatch.setattr(afm, "info", patch_afm_info) + + # AudioFile deserialization with minimum info should call afm info to read metadata + AudioFile.from_dict(minimum) + assert afm_calls[0] == 1 + + # AudioFile deserialization with full info should not call afm info + AudioFile.from_dict(full) + assert afm_calls[0] == 1 + + +def test_audio_file_to_dict_should_contain_afm_info( + audio_files: tuple[list[AudioFile], Any], monkeypatch: pytest.MonkeyPatch +) -> None: + audio_files, _ = audio_files + af = audio_files[0] + + minimum = { + "begin": af.begin.strftime(TIMESTAMP_FORMAT_EXPORTED_FILES_LOCALIZED), + "end": af.end.strftime(TIMESTAMP_FORMAT_EXPORTED_FILES_LOCALIZED), + "path": af.path, + } + full = minimum | { + "sample_rate": af.sample_rate, + "channels": af.channels, + } + + assert all(key in af.to_dict() for key in full) + + @pytest.mark.parametrize( ("audio_files", "begin", "end", "sample_rate", "normalization"), [