AudioWriter

AudioWriter class for writing AudioTree objects to disk with manifest support.

class AudioWriter(directory: str | Path = '.', *, pattern: str = 'audio_{index:04d}.wav', include_timestamp: bool = False, compress_manifest: bool = True, write_audio: bool = True, subtype: str | None = None, pbar: Any | None = None, close_pbar: bool = False, show_progress: bool = False, progress_desc: str | None = None, exist_ok: bool = False, manifest_every: int = 1000)

Write AudioTree objects sequentially to disk with optional manifest generation.

The AudioWriter provides a stateful way to write multiple AudioTree objects, maintaining consistent naming and optionally generating manifest files that track all written audio files and their extras.

Parameters:

  • directory – Directory where audio files will be written

  • pattern – Filename pattern with {index} placeholder for sequential numbering

  • include_timestamp – Whether to record a timestamp manifest column. One timestamp is minted per write() call and shared by every entry in that batch, so it records when the batch was written and distinct timestamps count write() calls.

  • compress_manifest – Whether to compress NPZ manifest files (only applies to npz format)

  • write_audio – Whether to write audio files to disk (default True). When False, only manifest is generated with extras

  • subtype – soundfile subtype string (e.g. "PCM_16", "PCM_24", "FLOAT"), forwarded to soundfile.write and recorded in the manifest’s subtype column. When None (default) the writer picks the least destructive subtype the container supports — "FLOAT" for WAV/AIFF/CAF/…, "PCM_24" for FLAC, the container’s own default otherwise. That is deliberately not libsndfile’s default of PCM_16, which hard-clips the out-of-[-1, 1] samples that float model output routinely contains. Passing a fixed-point subtype explicitly is fine; the writer then warns (RuntimeWarning) whenever an item it clips actually exceeds the representable range.

  • pbar – Optional tqdm progress bar instance to update during writing

  • close_pbar – Whether to close the progress bar on exit (default False)

  • show_progress – Create an internal tqdm progress bar. Raises ImportError if tqdm is not installed — install it with audiotree[progress].

  • progress_desc – Description for internal progress bar (default “Writing audio”)

  • exist_ok – Whether to write into a directory that already holds a manifest. Defaults to False, which raises FileExistsError rather than overwriting an existing dataset (and catches two writers aimed at one directory). Pass True to deliberately overwrite or append.

  • manifest_every – Rewrite manifest.npz once this many entries have accumulated since the last write, so a run killed after 100k files leaves a manifest describing (nearly) all of them instead of none. Each rewrite costs one pass over every entry so far, hence the throttle; pass 0 to write the manifest only at close().

Example

Write a handful of (silent, one-second mono) AudioTree objects to a temporary directory. A manifest.npz is refreshed every manifest_every entries and finalized on context exit.

>>> import tempfile
>>> from pathlib import Path
>>> from audiotree import AudioTree
>>> import jax.numpy as jnp
>>> out_dir = tempfile.mkdtemp()
>>> audio_batches = [AudioTree.create(jnp.zeros((1, 44100)), 44100) for _ in range(3)]
>>> with AudioWriter(out_dir) as writer:
...     for audio in audio_batches:
...         _ = writer.write(audio)
>>> sorted(p.name for p in Path(out_dir).glob("*"))
['audio_0000.wav', 'audio_0001.wav', 'audio_0002.wav', 'manifest.npz']

Writing again to a directory that already holds a manifest raises, so a finished dataset is never silently overwritten:

>>> AudioWriter(out_dir)
Traceback (most recent call last):
    ...
FileExistsError: ... already contains a dataset (manifest.npz). ...

Pass an external progress bar with pbar=..., or have the writer create its own with show_progress=True. exist_ok=True opts in to reusing the directory:

>>> from tqdm import tqdm
>>> pbar = tqdm(total=len(audio_batches), desc="Processing")
>>> with AudioWriter(out_dir, pbar=pbar, exist_ok=True) as writer:
...     for audio in audio_batches:
...         _ = writer.write(audio)
>>> with AudioWriter(out_dir, show_progress=True, exist_ok=True) as writer:
...     for audio in audio_batches:
...         _ = writer.write(audio)
close()

Manually close the writer, save manifest, and close progress bar.

get_stats() → Dict

Get statistics about written files.

Returns:

Dictionary containing write statistics. When include_timestamp is set, total_batches counts write() calls: each call mints one timestamp shared by its batch’s entries, so distinct timestamps are distinct batches.

save_manifest() → Path | None

Write manifest.npz atomically (temp file + rename).

A reader – or a retry after a crash – sees either the previous manifest or the new one, never the half-written NPZ that a kill during savez would otherwise leave behind (which reads as a corrupt zip and blocks re-rendering with “already contains a dataset”).

Returns:

Path to the saved manifest file, or None if no data to save

write(tree: AudioTree, tags: Dict | None = None) → List[Path]

Write all items in an AudioTree batch to disk.

Parameters:

  • tree – AudioTree containing one or more audio items in batch dimension

  • tags – Optional dictionary of custom metadata to include in manifest

Returns:

List of Path objects for all written files

Raises:

ValueError – If AudioTree fields don’t match previously written trees