import pyaudio
import wave

class SoundsWapper:
    def __init__(self, chunk=1024, format=pyaudio.paInt16, channels=1, rate=44100):
        self.CHUNK = chunk
        self.FORMAT = format
        self.CHANNELS = channels
        self.RATE = rate
        self.p = pyaudio.PyAudio()

    def record(self, seconds, filename='output.wav'):
        stream = self.p.open(format=self.FORMAT,
                             channels=self.CHANNELS,
                             rate=self.RATE,
                             input=True,
                             frames_per_buffer=self.CHUNK)

        print("开始录音...")
        frames = []

        for i in range(0, int(self.RATE / self.CHUNK * seconds)):
            data = stream.read(self.CHUNK)
            frames.append(data)

        print("录音结束.")

        stream.stop_stream()
        stream.close()

        wf = wave.open(filename, 'wb')
        wf.setnchannels(self.CHANNELS)
        wf.setsampwidth(self.p.get_sample_size(self.FORMAT))
        wf.setframerate(self.RATE)
        wf.writeframes(b''.join(frames))
        wf.close()

    def play(self, filename='output.wav'):
        wf = wave.open(filename, 'rb')

        stream = self.p.open(format=self.p.get_format_from_width(wf.getsampwidth()),
                             channels=wf.getnchannels(),
                             rate=wf.getframerate(),
                             output=True)

        print("开始播放...")
        data = wf.readframes(self.CHUNK)

        while data:
            stream.write(data)
            data = wf.readframes(self.CHUNK)

        print("播放结束.")

        stream.stop_stream()
        stream.close()

    def close(self):
        self.p.terminate()

# 使用示例
if __name__ == "__main__":
    sw = SoundsWapper()

    # 录音5秒钟
    sw.record(5)

    # 播放录音
    sw.play()

    # 清理资源
    sw.close()