Playing PCM data in python using SDL2 with SDL_mixer.

pcm.py

"""
Purpose:
  Play PCM data.

API overview:
  sound_init    Initialize sound system.
  sound_close   Close sound system.
  sound_play    Play PCM data.

Dependencies:
  py-sdl2 (https://github.com/py-sdl/py-sdl2/)
"""

import sys
import enum
import ctypes
import sdl2
import sdl2.sdlmixer

MIN_VOLUME = 0
MAX_VOLUME = 128

def clamp(value: int, min_value: int, max_value: int) -> int:
    return max(min_value, min(value, max_value))

class SampleFormat(enum.IntEnum):
    """
    Sample format of the audio data.

    Examples explaining the coding scheme.
    - S16:  Signed 16 bit integer sample data in native byte order.
    - U16L: Unsigned 16bit integer sample data i little endian byte order.
    - U16B: Unsigned 16bit integer sample data in big endian byte order.
    - F32:  32-bit floating point samples in native byte order.
    """
    S8   = sdl2.AUDIO_S8
    U8   = sdl2.AUDIO_U8
    S16  = sdl2.AUDIO_S16SYS
    S16L = sdl2.AUDIO_S16LSB
    S16B = sdl2.AUDIO_S16MSB
    U16  = sdl2.AUDIO_U16SYS
    U16L = sdl2.AUDIO_U16LSB
    U16B = sdl2.AUDIO_U16MSB
    S32  = sdl2.AUDIO_S32SYS
    S32L = sdl2.AUDIO_S32LSB
    S32B = sdl2.AUDIO_S32MSB
    F32L = sdl2.AUDIO_F32LSB
    F32B = sdl2.AUDIO_F32MSB
    F32  = sdl2.AUDIO_F32SYS

def sound_init(frequency: int = 44100, format: SampleFormat = SampleFormat.S16, channels: int = 2, chunksize = 1024) -> bool:
    """
    Initializes the sound subsystem and opens a playback device configured with the given arguments.
    Should be paired with a call to sound_close (even on failure).
    
    :param frequency:   Frequency to playback audio at (in Hz).
    :param format:      Sample format (data type and endianness).
    :param channels:    Number of channels (1 is mono, 2 is stereo, etc).
    :param chunksize:   Audio buffer size in sample frames (total samples divided by channel count).
    :returns:           True on success, False on failure.
    """
    
    # Could throw in a check for valid frequency (i.e. non-negative integer) here to give a better error message...
   
    if sdl2.SDL_InitSubSystem(sdl2.SDL_INIT_AUDIO) < 0:
        print(f"Failed to initialize audio subsystem: {sdl2.SDL_GetError()}", file = sys.stderr)
        return False

    if sdl2.sdlmixer.Mix_OpenAudio(frequency, format, channels, chunksize) < 0:
        print(f"Failed to open audio device: {sdl2.sdlmixer.Mix_GetError()}", file = sys.stderr)
        return False

    return True

def sound_close():
    """
    Closes the opened playback device and the sound subsystem.
    """
    
    sdl2.sdlmixer.Mix_CloseAudio()
    sdl2.SDL_QuitSubSystem(sdl2.SDL_INIT_AUDIO)
    sdl2.SDL_Quit()

def sound_play(data: bytes, volume: int = MAX_VOLUME):
    """
    Plays the sound with the specified volume.
    Blocks until playing is complete.
    See MIN_VOLUME and MAX_VOLUME for the valid volume range.

    :param data:   buffer containing the PCM samples to play.
    :param volume: Sets the volume used while playing the sound
    """
    
    try:
        buffer = (ctypes.c_ubyte * len(data)).from_buffer_copy(data)
    except TypeError as e:
        print(f"Failed to play sound: {e}", file = sys.stderr)
        return

    chunk = sdl2.sdlmixer.Mix_Chunk()
    chunk.allocated = 0
    chunk.abuf      = buffer
    chunk.alen      = len(data)
    chunk.volume    = clamp(volume, MIN_VOLUME, MAX_VOLUME)
    
    channel = sdl2.sdlmixer.Mix_PlayChannel(-1,  chunk, 0)
    if channel == -1:
        print(f"Failed to play sound: {sdl2.sdlmixer.Mix_GetError()}", file = sys.stderr)
    
    while sdl2.sdlmixer.Mix_Playing(channel):
        sdl2.SDL_Delay(50)

wav.py

# pcm.py example usage: Play wav file.

import sys
import wave
import pcm

def load_wave(path):
    try:
        with wave.open(path) as file:
            sample_frequency = file.getframerate()
            sample_width     = file.getsampwidth()
            channels         = file.getnchannels()
            frames           = file.getnframes()
            data             = file.readframes(frames)
    except Exception as e:
        print(f"Failed to load {path}: {e}", file = sys.stderr)
        return None
    
    if sample_width == 1:
        format = pcm.SampleFormat.U8
    elif sample_width == 2:
        format = pcm.SampleFormat.S16L
    else:
        print(f"Unsupported sample width: {sample_width} bytes", file = sys.stderr)
        return None
        
    return (data, sample_frequency, format, channels)

if __name__ == "__main__":
    # Load PCM data from wave file. 
    if len(sys.argv) != 2:
        print(f"Usage: python3 {__name__} <path_to_wav>")
        sys.exit()
    
    wav = load_wave(sys.argv[1])
    if not wav:
        sys.exit()

    data, sample_frequency, format, channels = wav
    
    # Play loaded PCM data.
    if pcm.sound_init(sample_frequency, format, channels):
        try:
            pcm.sound_play(data)
        except KeyboardInterrupt:
            print("\nKeyboard interrupt")
    pcm.sound_close()

wavegen.py

# pcm.py example usage: Playing the result from sine and square wave generation.

import math
import array
import pcm

def sine_wave(samples: array.array, sample_count: int, sample_rate: int, amplitude: int, frequency: int, phase_shift: int = 0, vertical_shift: int = 0) -> array.array:
    for i in range(sample_count):
        t = (i / sample_rate)
        sample = amplitude * math.sin(2 * math.pi * frequency * t + phase_shift) + vertical_shift
        sample = int(sample)
        samples.append(sample)

    return samples

def square_wave(samples: array.array, sample_count: int, sample_rate: int, amplitude: int, frequency: int) -> array.array:
    full_cycle = sample_rate / frequency
    half_cycle = full_cycle  / 2
    cycle = 0
    
    for i in range(sample_count):
        if cycle < half_cycle:
            samples.append(amplitude)
        else:
            samples.append(-amplitude)
        cycle = (cycle + 1) % full_cycle
    
    return samples

if __name__ == "__main__":
    # PCM playback configuration.
    volume        = int(pcm.MAX_VOLUME * 0.4)
    channels      = 1
    sample_rate   = 44100
    sample_format = pcm.SampleFormat.S16

    # Generate samples.
    amplitude    = 32767
    seconds      = 0.3
    sample_count = int(sample_rate * seconds)

    samples = array.array("h")

    sine_wave(samples, sample_count, sample_rate, amplitude, frequency = 200)
    sine_wave(samples, sample_count, sample_rate, amplitude, frequency = 250)
    sine_wave(samples, sample_count, sample_rate, amplitude, frequency = 300)
    sine_wave(samples, sample_count, sample_rate, amplitude, frequency = 350)
    
    amplitude = int(amplitude * 0.1) # Lower volume.

    square_wave(samples, sample_count, sample_rate, amplitude, frequency = 200)
    square_wave(samples, sample_count, sample_rate, amplitude, frequency = 250)
    square_wave(samples, sample_count, sample_rate, amplitude, frequency = 300)
    square_wave(samples, sample_count, sample_rate, amplitude, frequency = 350)
    
    # Play generated wave forms.    
    pcm.sound_init(sample_rate, sample_format, channels)
    pcm.sound_play(samples.tobytes(), volume)
    pcm.sound_close()