Refactor to make record function more readable

Signed-off-by: makaveli10 <suryanvineet47@gmail.com>

whisper_live/client.py

@@ -260,6 +260,7 @@ class Client:
         while time.time() - self.last_response_received < self.disconnect_if_no_response_for:
             continue
 
+
 class TranscriptionTeeClient:
     """
     Client for handling audio recording, streaming, and transcription tasks via one or more
@@ -476,6 +477,42 @@ class TranscriptionTeeClient:
 
         return process
 
+    def save_chunk(self, n_audio_file):
+        """
+        Saves the current audio frames to a WAV file in a separate thread.
+
+        Args:
+        n_audio_file (int): The index of the audio file which determines the filename.
+                            This helps in maintaining the order and uniqueness of each chunk.
+        """
+        t = threading.Thread(
+            target=self.write_audio_frames_to_file,
+            args=(self.frames[:], f"chunks/{n_audio_file}.wav",),
+        )
+        t.start()
+
+    def finalize_recording(self, n_audio_file):
+        """
+        Finalizes the recording process by saving any remaining audio frames,
+        closing the audio stream, and terminating the process.
+
+        Args:
+        n_audio_file (int): The file index to be used if there are remaining audio frames to be saved.
+                            This index is incremented before use if the last chunk is saved.
+        """
+        if self.save_output_recording and len(self.frames):
+            self.write_audio_frames_to_file(
+                self.frames[:], f"chunks/{n_audio_file}.wav"
+            )
+            n_audio_file += 1
+        self.stream.stop_stream()
+        self.stream.close()
+        self.p.terminate()
+        self.close_all_clients()
+        if self.save_output_recording:
+            self.write_output_recording(n_audio_file)
+        self.write_all_clients_srt()
+
     def record(self):
         """
         Record audio data from the input stream and save it to a WAV file.
@@ -508,31 +545,13 @@ class TranscriptionTeeClient:
                 # save frames if more than a minute
                 if len(self.frames) > 60 * self.rate:
                     if self.save_output_recording:
-                        t = threading.Thread(
-                            target=self.write_audio_frames_to_file,
-                            args=(
-                                self.frames[:],
-                                f"chunks/{n_audio_file}.wav",
-                            ),
-                        )
-                        t.start()
+                        self.save_chunk(n_audio_file)
                         n_audio_file += 1
                     self.frames = b""
             self.write_all_clients_srt()
 
         except KeyboardInterrupt:
-            if self.save_output_recording and len(self.frames):
-                self.write_audio_frames_to_file(
-                    self.frames[:], f"chunks/{n_audio_file}.wav"
-                )
-                n_audio_file += 1
-            self.stream.stop_stream()
-            self.stream.close()
-            self.p.terminate()
-            self.close_all_clients()
-            if self.save_output_recording:
-                self.write_output_recording(n_audio_file)
-            self.write_all_clients_srt()
+            self.finalize_recording(n_audio_file)
 
     def write_audio_frames_to_file(self, frames, file_name):
         """
@@ -608,6 +627,7 @@ class TranscriptionTeeClient:
         raw_data = np.frombuffer(buffer=audio_bytes, dtype=np.int16)
         return raw_data.astype(np.float32) / 32768.0
 
+
 class TranscriptionClient(TranscriptionTeeClient):
     """
     Client for handling audio transcription tasks via a single WebSocket connection.
@@ -633,8 +653,23 @@ class TranscriptionClient(TranscriptionTeeClient):
         transcription_client()
         ```
     """
-    def __init__(self, host, port, lang=None, translate=False, model="small", use_vad=True, save_output_recording=False, output_recording_filename="./output_recording.wav"):
+    def __init__(
+        self,
+        host,
+        port,
+        lang=None,
+        translate=False,
+        model="small",
+        use_vad=True,
+        save_output_recording=False,
+        output_recording_filename="./output_recording.wav"
+    ):
         self.client = Client(host, port, lang, translate, model, srt_file_path="output.srt", use_vad=use_vad)
         if save_output_recording and not output_recording_filename.endswith(".wav"):
             raise ValueError(f"Please provide a valid `output_recording_filename`: {output_recording_filename}")
-        TranscriptionTeeClient.__init__(self, [self.client], save_output_recording=save_output_recording, output_recording_filename=output_recording_filename)
+        TranscriptionTeeClient.__init__(
+            self,
+            [self.client],
+            save_output_recording=save_output_recording,
+            output_recording_filename=output_recording_filename
+        )