Parallel Processing¶

def tqdm_write_sink(message: str) -> None:
    """
    Custom sink for loguru that writes messages using tqdm.write().

    This ensures that log messages don't interfere with tqdm progress bars.
    """  # noqa: D401
    tqdm.write(message, end="")

def load_analyse_binaural(
    wav_file: Path,
    levels: dict[str, float] | list[float] | None,
    analysis_settings: AnalysisSettings,
    resample: int | None = None,
    *,
    parallel_mosqito: bool = True,
) -> pd.DataFrame:
    """
    Load and analyze a single binaural audio file.

    Parameters
    ----------
    wav_file
        Path to the WAV file.
    levels
        Dictionary with calibration levels for each channel.
    analysis_settings
        Analysis settings object.
    resample
        Sampling rate to resample the audio to before analysis.
    parallel_mosqito
        Whether to process MoSQITo metrics in parallel. Defaults to True.

    Returns
    -------
    :
        DataFrame with analysis results.

    """
    logger.info(f"Processing {wav_file}")
    try:
        b = Binaural.from_wav(wav_file, resample=resample)
        if levels is not None:
            if isinstance(levels, dict) and b.recording in levels:
                decibel = (levels[b.recording]["Left"], levels[b.recording]["Right"])
                b.calibrate_to(decibel, inplace=True)
            elif isinstance(levels, list | tuple):
                logger.debug(f"Calibrating {wav_file} to {levels} dB")
                b.calibrate_to(levels, inplace=True)
            else:
                logger.warning(f"No calibration levels found for {wav_file}")
        else:
            logger.warning(f"No calibration levels found for {wav_file}")
        return process_all_metrics(b, analysis_settings, parallel=parallel_mosqito)
    except Exception as e:
        logger.error(f"Error processing {wav_file}: {e!s}")
        raise

def parallel_process(
    wav_files: list[Path],
    results_df: pd.DataFrame,
    levels: dict,
    analysis_settings: AnalysisSettings,
    max_workers: int | None = None,
    resample: int | None = None,
    *,
    parallel_mosqito: bool = True,
) -> pd.DataFrame:
    """
    Process multiple binaural files in parallel.

    Parameters
    ----------
    wav_files
        List of WAV files to process.
    results_df
        Initial results DataFrame to update.
    levels
        Dictionary with calibration levels for each file.
    analysis_settings
        Analysis settings object.
    max_workers
        Maximum number of worker processes.
        If None, it will default to the number of processors on the machine.
    resample
        Sampling rate to resample the audio to before analysis.
    parallel_mosqito
        Whether to process MoSQITo metrics in parallel within each file.
        Defaults to True.

    Returns
    -------
    :
        Updated results DataFrame with analysis results for all files.

    """
    logger.info(f"Starting parallel processing of {len(wav_files)} files")

    # Add a handler that uses tqdm.write for output
    tqdm_handler_id = logger.add(tqdm_write_sink, format="{message}")

    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
        futures = []
        for wav_file in wav_files:
            future = executor.submit(
                load_analyse_binaural,
                wav_file,
                levels,
                analysis_settings,
                resample,
                parallel_mosqito=parallel_mosqito,
            )
            futures.append(future)

        with tqdm(total=len(futures), desc="Processing files") as pbar:
            for future in concurrent.futures.as_completed(futures):
                try:
                    result = future.result()
                    results_df = add_results(results_df, result)
                except Exception as e:  # noqa: BLE001
                    logger.error(f"Error processing file: {e!s}")
                finally:
                    pbar.update(1)

    # Remove the tqdm-compatible handler
    logger.remove(tqdm_handler_id)

    logger.info("Parallel processing completed")
    return results_df