drb-client-discord-bot/app/internal/NoiseGatev2.py

import audioop
import math
import pyaudio
import discord
from internal.logger import create_logger

LOGGER = create_logger(__name__)

# noinspection PyUnresolvedReferences
class AudioStream:
    def __init__(self, _channels: int = 2, _sample_rate: int = 48000, _frames_per_buffer: int = 960,
                 _input_device_index: int = None, _output_device_index: int = None, _input: bool = True,
                 _output: bool = True, _init_on_startup: bool = True):
        # NOTE: frames_per_buffer changed to 960 to match Discord's 20ms frame size
        self.paInstance_kwargs = {
            'format': pyaudio.paInt16,
            'channels': _channels,
            'rate': _sample_rate,
            'input': _input,
            'output': _output,
            'frames_per_buffer': _frames_per_buffer
        }

        if _input_device_index is not None:
            if _input:
                self.paInstance_kwargs['input_device_index'] = _input_device_index
            else:
                LOGGER.warning("[AudioStream.__init__]:\tInput was not enabled."
                               " Reinitialize with '_input=True'")

        if _output_device_index is not None:
            if _output:
                self.paInstance_kwargs['output_device_index'] = _output_device_index
            else:
                LOGGER.warning("[AudioStream.__init__]:\tOutput was not enabled."
                               " Reinitialize with '_output=True'")

        if _init_on_startup:
            LOGGER.info("Creating PyAudio instance")
            self.paInstance = pyaudio.PyAudio()
            self.stream = None

        if _output_device_index is not None or _input_device_index is not None:
            if _init_on_startup:
                LOGGER.info("Init stream")
                self.init_stream()

    def init_stream(self, _new_output_device_index: int = None, _new_input_device_index: int = None):
        if _new_input_device_index is not None:
            if self.paInstance_kwargs['input']:
                self.paInstance_kwargs['input_device_index'] = _new_input_device_index
            else:
                LOGGER.warning("[AudioStream.init_stream]:\tInput was not enabled when initialized."
                               " Reinitialize with '_input=True'")

        if _new_output_device_index is not None:
            if self.paInstance_kwargs['output']:
                self.paInstance_kwargs['output_device_index'] = _new_output_device_index
            else:
                LOGGER.warning("[AudioStream.init_stream]:\tOutput was not enabled when initialized."
                               " Reinitialize with '_output=True'")

        self.close_if_open()
        self.stream = self.paInstance.open(**self.paInstance_kwargs)

    def close_if_open(self):
        if self.stream and self.stream.is_active():
            self.stream.stop_stream()
            self.stream.close()
            LOGGER.debug("[ReopenStream.close_if_open]:\t Stream was open; It was closed.")

    def list_devices(self, _display_input_devices: bool = True, _display_output_devices: bool = True):
        LOGGER.info('Getting a list of the devices connected')
        info = self.paInstance.get_host_api_info_by_index(0)
        numdevices = info.get('deviceCount')

        devices = {'Input': {}, 'Output': {}}
        for i in range(0, numdevices):
            device_info = self.paInstance.get_device_info_by_host_api_device_index(0, i)
            if (device_info.get('maxInputChannels')) > 0:
                input_device = device_info.get('name')
                devices['Input'][i] = input_device
                if _display_input_devices:
                    LOGGER.debug(f"Input Device id {i} - {input_device}")

            if (device_info.get('maxOutputChannels')) > 0:
                output_device = device_info.get('name')
                devices['Output'][i] = output_device
                if _display_output_devices:
                    LOGGER.debug(f"Output Device id {i} - {output_device}")
        return devices

# noinspection PyUnresolvedReferences
class NoiseGate(AudioStream):
    def __init__(self, _voice_connection, _noise_gate_threshold: int, **kwargs):
        super(NoiseGate, self).__init__(_init_on_startup=True, **kwargs)
        self.voice_connection = _voice_connection
        self.THRESHOLD = _noise_gate_threshold
        self.NGStream = NoiseGateStream(self)

    def run(self) -> None:
        LOGGER.debug("Starting stream")
        self.stream.start_stream()
        self.core()

    def core(self, error=None):
        if error:
            LOGGER.warning(f"Audio stream stopped unexpectedly with error: {error}")
            return # Avoid recursion on error

        if self.voice_connection.is_connected() and not self.voice_connection.is_playing():
            LOGGER.debug("Playing stream to discord")
            # The 'after' callback can be prone to loops, simplified here
            self.voice_connection.play(self.NGStream, after=lambda e: self.core(e))

    async def close(self):
        LOGGER.debug("Closing NoiseGate resources...")
        if self.voice_connection and self.voice_connection.is_connected():
            self.voice_connection.stop() # Stop sending audio
        
        self.close_if_open() # Close PyAudio stream
        
        if self.paInstance:
            self.paInstance.terminate()
        
        LOGGER.debug("NoiseGate resources closed.")

# noinspection PyUnresolvedReferences
class NoiseGateStream(discord.AudioSource):
    def __init__(self, noise_gate_instance: NoiseGate):
        super(NoiseGateStream, self).__init__()
        self.noise_gate = noise_gate_instance
        self.NG_fadeout = 12 # Equivalent to 240ms of audio frames (240 / 20ms)
        self.NG_fadeout_count = 0
        self.process_set_count = 0

    def read(self):
        try:
            # Check connection status via the parent NoiseGate instance
            if not self.noise_gate.voice_connection.is_connected():
                return b''

            # Read from the PyAudio stream, also via the parent instance
            curr_buffer = self.noise_gate.stream.read(960, exception_on_overflow=False)

            if not curr_buffer:
                return b''

            buffer_rms = audioop.rms(curr_buffer, 2)
            
            if buffer_rms > 0:
                buffer_decibel = 20 * math.log10(buffer_rms)

                if self.process_set_count % 10 == 0:
                    log_msg = f"[{'Open' if buffer_decibel >= self.noise_gate.THRESHOLD else 'Closed'}]"
                    LOGGER.debug(f"[NoiseGate {log_msg}] {buffer_decibel:.2f} dB")

                if buffer_decibel >= self.noise_gate.THRESHOLD:
                    self.NG_fadeout_count = self.NG_fadeout
                    self.process_set_count += 1
                    return bytes(curr_buffer)
                
                elif self.NG_fadeout_count > 0:
                    self.NG_fadeout_count -= 1
                    self.process_set_count += 1
                    return bytes(curr_buffer)

            return b'' # Return silence if below threshold and not in fadeout

        except IOError as e:
            LOGGER.error(f"PyAudio IOError in read(): {e}")
            return b'' # Return silence to not break the Discord audio pump
        except Exception as e:
            LOGGER.error(f"Unhandled exception in NoiseGateStream.read: {e}", exc_info=True)
            return b''