auto-editor 27.0.0__py3-none-any.whl → 27.1.0__py3-none-any.whl

auto_editor/wavfile.py

@@ -1,310 +0,0 @@
-from __future__ import annotations
-
-import io
-import struct
-import sys
-from typing import TYPE_CHECKING, Literal
-
-import numpy as np
-
-PCM = 0x0001
-IEEE_FLOAT = 0x0003
-EXTENSIBLE = 0xFFFE
-
-AudioData = np.memmap | np.ndarray
-Endian = Literal[">", "<"]  # Big Endian, Little Endian
-ByteOrd = Literal["big", "little"]
-
-if TYPE_CHECKING:
-    Reader = io.BufferedReader | io.BytesIO
-    Writer = io.BufferedWriter | io.BytesIO
-
-
-class WavError(Exception):
-    pass
-
-
-def _read_fmt_chunk(
-    fid: Reader, bytes_order: ByteOrd
-) -> tuple[int, int, int, int, int]:
-    size = int.from_bytes(fid.read(4), bytes_order)
-
-    if size < 16:
-        raise WavError("Binary structure of wave file is not compliant")
-
-    format_tag = int.from_bytes(fid.read(2), bytes_order)
-    channels = int.from_bytes(fid.read(2), bytes_order)
-    sr = int.from_bytes(fid.read(4), bytes_order)
-    fid.read(4)  # This is bitrate but we don't need it
-    block_align = int.from_bytes(fid.read(2), bytes_order)
-    bit_depth = int.from_bytes(fid.read(2), bytes_order)
-    bytes_read = 16
-
-    if format_tag == EXTENSIBLE and size >= 18:
-        ext_chunk_size = int.from_bytes(fid.read(2), bytes_order)
-        bytes_read += 2
-        if ext_chunk_size >= 22:
-            extensible_chunk_data = fid.read(22)
-            bytes_read += 22
-            raw_guid = extensible_chunk_data[6:22]
-
-            if bytes_order == "big":
-                tail = b"\x00\x00\x00\x10\x80\x00\x00\xaa\x00\x38\x9b\x71"
-            else:
-                tail = b"\x00\x00\x10\x00\x80\x00\x00\xaa\x00\x38\x9b\x71"
-            if raw_guid.endswith(tail):
-                format_tag = int.from_bytes(raw_guid[:4], bytes_order)
-        else:
-            raise WavError("Binary structure of wave file is not compliant")
-
-    if format_tag not in {PCM, IEEE_FLOAT}:
-        raise WavError(
-            f"Encountered unknown format tag: {format_tag:#06x}, while reading fmt chunk."
-        )
-
-    # move file pointer to next chunk
-    if size > bytes_read:
-        fid.read(size - bytes_read)
-
-    # fmt should always be 16, 18 or 40, but handle it just in case
-    _handle_pad_byte(fid, size)
-
-    return format_tag, channels, sr, block_align, bit_depth
-
-
-def _read_data_chunk(
-    fid: Reader,
-    format_tag: int,
-    channels: int,
-    bit_depth: int,
-    en: Endian,
-    block_align: int,
-    data_size: int | None,
-) -> AudioData:
-    bytes_order: ByteOrd = "big" if en == ">" else "little"
-    size = int.from_bytes(fid.read(4), bytes_order)
-
-    if data_size is not None:
-        # size is only 32-bits here, so get real size from header.
-        size = data_size
-
-    bytes_per_sample = block_align // channels
-
-    if bytes_per_sample in {3, 5, 6, 7}:
-        raise WavError(f"Unsupported bytes per sample: {bytes_per_sample}")
-
-    if format_tag == PCM:
-        if 1 <= bit_depth <= 8:
-            dtype = "u1"  # WAVs of 8-bit integer or less are unsigned
-        elif bit_depth <= 64:
-            dtype = f"{en}i{bytes_per_sample}"
-        else:
-            raise WavError(
-                f"Unsupported bit depth: the WAV file has {bit_depth}-bit integer data."
-            )
-    elif format_tag == IEEE_FLOAT:
-        if bit_depth in {32, 64}:
-            dtype = f"{en}f{bytes_per_sample}"
-        else:
-            raise WavError(
-                f"Unsupported bit depth: the WAV file has {bit_depth}-bit floating-point data."
-            )
-    else:
-        raise WavError(
-            f"Unknown wave file format: {format_tag:#06x}. Supported formats: PCM, IEEE_FLOAT"
-        )
-    if size % 2 == 0:
-        n_samples = size // block_align
-    else:
-        n_samples = (size - 1) // block_align
-
-    if isinstance(fid, io.BufferedReader):
-        data: AudioData = np.memmap(
-            fid, dtype=dtype, mode="c", offset=fid.tell(), shape=(n_samples, channels)
-        )
-        fid.seek(size, 1)
-    else:
-        bytes_per_sample = np.dtype(dtype).itemsize
-        buffer = fid.read(n_samples * channels * bytes_per_sample)
-        data = np.frombuffer(buffer, dtype=dtype).reshape((n_samples, channels))
-
-    _handle_pad_byte(fid, size)
-
-    return data
-
-
-def _skip_unknown_chunk(fid: Reader, en: Endian) -> None:
-    data = fid.read(4)
-
-    if len(data) == 4:
-        size = struct.unpack(f"{en}I", data)[0]
-        if size == 0:
-            raise WavError("Unknown chunk size is 0")
-        fid.seek(size, 1)
-        _handle_pad_byte(fid, size)
-    elif len(data) == 0:
-        pass  # It's okay, we've hit EOF
-    else:
-        raise WavError(
-            f"Unknown chunk size has wrong length, expected 4, got {len(data)}"
-        )
-
-
-def _read_rf64_chunk(fid: Reader) -> tuple[int, int, Endian]:
-    # https://tech.ebu.ch/docs/tech/tech3306v1_0.pdf
-    # https://www.itu.int/dms_pubrec/itu-r/rec/bs/R-REC-BS.2088-1-201910-I!!PDF-E.pdf
-
-    heading = fid.read(12)
-    if heading != b"\xff\xff\xff\xffWAVEds64":
-        raise WavError(f"Invalid heading for rf64 chunk: {heading!r}")
-
-    chunk_size = fid.read(4)
-
-    bw_size_low = fid.read(4)
-    bw_size_high = fid.read(4)
-
-    en: Endian = ">" if (bw_size_high > bw_size_low) else "<"
-    bytes_order: ByteOrd = "big" if en == ">" else "little"
-
-    data_size_low = fid.read(4)
-    data_size_high = fid.read(4)
-
-    file_size = int.from_bytes(bw_size_low + bw_size_high, "little")
-    data_size = int.from_bytes(data_size_low + data_size_high, "little")
-
-    int_chunk_size = int.from_bytes(chunk_size, bytes_order)
-    if type(int_chunk_size) is not int or int_chunk_size > 40:
-        raise WavError("Invalid chunk size in RF64 chunk")
-
-    fid.read(40 - int_chunk_size)
-
-    return data_size, file_size, en
-
-
-def _read_riff_chunk(sig: bytes, fid: Reader) -> tuple[None, int, Endian]:
-    en: Endian = "<" if sig == b"RIFF" else ">"
-    bytes_order: ByteOrd = "big" if en == ">" else "little"
-
-    file_size = int.from_bytes(fid.read(4), bytes_order) + 8
-
-    form = fid.read(4)
-    if form != b"WAVE":
-        raise WavError(f"Not a WAV file. RIFF form type is {form!r}.")
-
-    return None, file_size, en
-
-
-def _handle_pad_byte(fid: Reader, size: int) -> None:
-    if size % 2 == 1:
-        fid.seek(1, 1)
-
-
-def read(fid: Reader) -> tuple[int, AudioData]:
-    file_sig = fid.read(4)
-    if file_sig in {b"RIFF", b"RIFX"}:
-        data_size, file_size, en = _read_riff_chunk(file_sig, fid)
-    elif file_sig == b"RF64":
-        data_size, file_size, en = _read_rf64_chunk(fid)
-    else:
-        raise WavError(f"File format {file_sig!r} not supported.")
-
-    bytes_order: ByteOrd = "big" if en == ">" else "little"
-    fmt_chunk_received = False
-
-    while fid.tell() < file_size:
-        chunk_id = fid.read(4)
-
-        if not chunk_id:
-            raise WavError("Unexpected end of file.")
-        if len(chunk_id) < 4:
-            raise WavError(f"Incomplete chunk ID: {chunk_id!r}")
-
-        if chunk_id == b"fmt ":
-            fmt_chunk_received = True
-            format_tag, channels, sr, block_align, bit_depth = _read_fmt_chunk(
-                fid, bytes_order
-            )
-        elif chunk_id == b"data":
-            if not fmt_chunk_received:
-                raise WavError("No fmt chunk before data")
-
-            data = _read_data_chunk(
-                fid,
-                format_tag,
-                channels,
-                bit_depth,
-                en,
-                block_align,
-                data_size,
-            )
-
-            fid.seek(0)
-            return sr, data
-
-        elif chunk_id == b"\x00\x00\x00\x00":
-            raise WavError("Invalid chunk ID")
-        else:
-            _skip_unknown_chunk(fid, en)
-
-    raise WavError("Found no data")
-
-
-def write(fid: Writer, sr: int, arr: np.ndarray) -> None:
-    channels = 1 if arr.ndim == 1 else arr.shape[1]
-    bit_depth = arr.dtype.itemsize * 8
-    block_align = channels * (bit_depth // 8)
-    data_size = arr.nbytes
-    total_size = 44 + data_size  # Basic WAV header size + data size
-
-    if is_rf64 := total_size > 0xFFFFFFFF:
-        fid.write(b"RF64\xff\xff\xff\xffWAVE")
-        ds64_size = 28
-        ds64_chunk_data = (0).to_bytes(ds64_size, "little")  # placeholder values
-        fid.write(b"ds64" + struct.pack("<I", ds64_size) + ds64_chunk_data)
-    else:
-        fid.write(b"RIFF" + struct.pack("<I", total_size - 8) + b"WAVE")
-
-    dkind = arr.dtype.kind
-    format_tag = IEEE_FLOAT if dkind == "f" else PCM
-
-    fmt_chunk_data = struct.pack(
-        "<HHIIHH", format_tag, channels, sr, 0, block_align, bit_depth
-    )
-    fid.write(b"fmt " + struct.pack("<I", len(fmt_chunk_data)) + fmt_chunk_data)
-
-    # Data chunk
-    fid.write(b"data")
-    fid.write(struct.pack("<I", 0xFFFFFFFF if is_rf64 else data_size))
-
-    if arr.dtype.byteorder == ">" or (
-        arr.dtype.byteorder == "=" and sys.byteorder == "big"
-    ):
-        arr = arr.byteswap()
-    fid.write(arr.ravel().view("b").data)
-
-    if is_rf64:
-        end_position = fid.tell()
-        fid.seek(16)  # Position at the start of 'ds64' chunk size
-
-        file_size = end_position - 20
-        fid.write(struct.pack("<I", ds64_size))
-        fid.write(file_size.to_bytes(8, "little") + data_size.to_bytes(8, "little"))
-
-        fid.seek(end_position)
-
-
-def main() -> None:
-    data = np.random.rand(48000, 2)
-    with open("test.wav", "wb") as file:
-        write(file, 48_000, data)
-
-    with open("test.wav", "rb") as file:
-        read_sr, read_data = read(file)
-
-    assert read_sr == 48_000
-    assert np.array_equal(data, read_data)
-    print("success")
-
-
-if __name__ == "__main__":
-    main()