Address code review feedback: improve error handling, validation, and test coverage

Co-authored-by: wannaphong <8536487+wannaphong@users.noreply.github.com>

Author: Copilot

pythaitts/__init__.py

@@ -55,7 +55,7 @@ def load_pretrained(self,version):
             from pythaitts.pretrained.vachana_tts import VachanaTTS
             self.model = VachanaTTS()
         else:
-            raise NotImplemented(
+            raise NotImplementedError(
                 "PyThaiTTS doesn't support %s pretrained." % self.pretrained
             )
 

pythaitts/pretrained/vachana_tts.py

@@ -8,9 +8,15 @@
 See more: https://github.com/VYNCX/VachanaTTS2
 """
 import tempfile
+import wave
+import numpy as np
+import os
 
 
 class VachanaTTS:
+    # Supported voice options
+    SUPPORTED_VOICES = ["th_f_1", "th_m_1", "th_f_2", "th_m_2"]
+    
     def __init__(self) -> None:
         """
         Initialize VachanaTTS model.
@@ -35,6 +41,12 @@ def __call__(self, text: str, speaker_idx: str = "th_f_1", return_type: str = "f
         :param kwargs: Additional parameters (volume, speed, noise_scale, noise_w_scale)
         :return: File path if return_type is "file", otherwise audio waveform data
         """
+        # Validate speaker_idx
+        if speaker_idx not in self.SUPPORTED_VOICES:
+            raise ValueError(
+                f"Unsupported voice '{speaker_idx}'. Supported voices are: {', '.join(self.SUPPORTED_VOICES)}"
+            )
+        
         # Extract additional parameters with defaults
         volume = kwargs.get('volume', 1.0)
         speed = kwargs.get('speed', 1.0)
@@ -43,39 +55,43 @@ def __call__(self, text: str, speaker_idx: str = "th_f_1", return_type: str = "f
 
         if return_type == "waveform":
             # For waveform return, we need to generate to a temp file then read it
-            import wave
-            import numpy as np
-            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
-                temp_filename = fp.name
-            
-            # Generate the audio file
-            self.tts_func(
-                text,
-                voice=speaker_idx,
-                output=temp_filename,
-                volume=volume,
-                speed=speed,
-                noise_scale=noise_scale,
-                noise_w_scale=noise_w_scale
-            )
-            
-            # Read the waveform from the file
-            with wave.open(temp_filename, 'rb') as wav_file:
-                n_frames = wav_file.getnframes()
-                audio_data = wav_file.readframes(n_frames)
-                # Convert bytes to numpy array
-                import struct
-                sample_width = wav_file.getsampwidth()
-                if sample_width == 2:
-                    waveform = np.frombuffer(audio_data, dtype=np.int16)
-                else:
-                    waveform = np.frombuffer(audio_data, dtype=np.int8)
-            
-            # Clean up temp file
-            import os
-            os.unlink(temp_filename)
-            
-            return waveform
+            temp_filename = None
+            try:
+                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
+                    temp_filename = fp.name
+                
+                # Generate the audio file
+                self.tts_func(
+                    text,
+                    voice=speaker_idx,
+                    output=temp_filename,
+                    volume=volume,
+                    speed=speed,
+                    noise_scale=noise_scale,
+                    noise_w_scale=noise_w_scale
+                )
+                
+                # Read the waveform from the file
+                with wave.open(temp_filename, 'rb') as wav_file:
+                    n_frames = wav_file.getnframes()
+                    audio_data = wav_file.readframes(n_frames)
+                    sample_width = wav_file.getsampwidth()
+                    
+                    # Convert bytes to numpy array based on sample width
+                    if sample_width == 1:
+                        waveform = np.frombuffer(audio_data, dtype=np.int8)
+                    elif sample_width == 2:
+                        waveform = np.frombuffer(audio_data, dtype=np.int16)
+                    elif sample_width == 4:
+                        waveform = np.frombuffer(audio_data, dtype=np.int32)
+                    else:
+                        raise ValueError(f"Unsupported sample width: {sample_width} bytes")
+                
+                return waveform
+            finally:
+                # Clean up temp file
+                if temp_filename and os.path.exists(temp_filename):
+                    os.unlink(temp_filename)
         else:
             # File output
             if filename is None:

tests/test_vachana.py

@@ -4,6 +4,7 @@
 """
 import unittest
 from unittest.mock import Mock, patch, MagicMock
+import numpy as np
 from pythaitts import TTS
 
 
@@ -67,6 +68,46 @@ def test_vachana_with_preprocessing(self, mock_vachana_class):
         self.assertIn("ห้า", processed_text)
         self.assertIn("คนคน", processed_text)
 
+    @patch('pythaitts.pretrained.vachana_tts.VachanaTTS')
+    def test_vachana_all_supported_voices(self, mock_vachana_class):
+        """Test that all supported voices work correctly"""
+        # Setup mock
+        mock_instance = Mock()
+        mock_instance.return_value = "/tmp/output.wav"
+        mock_vachana_class.return_value = mock_instance
+        
+        # Create TTS instance
+        tts = TTS(pretrained="vachana")
+        
+        # Test all supported voices
+        supported_voices = ["th_f_1", "th_m_1", "th_f_2", "th_m_2"]
+        for voice in supported_voices:
+            mock_instance.reset_mock()
+            result = tts.tts("สวัสดี", speaker_idx=voice)
+            
+            # Verify the voice was passed correctly
+            call_args = mock_instance.call_args
+            self.assertEqual(call_args.kwargs['speaker_idx'], voice)
+
+    @patch('pythaitts.pretrained.vachana_tts.VachanaTTS')
+    def test_vachana_waveform_return(self, mock_vachana_class):
+        """Test waveform return type functionality"""
+        # Setup mock
+        mock_instance = Mock()
+        mock_waveform = np.array([0.1, 0.2, 0.3, 0.4])
+        mock_instance.return_value = mock_waveform
+        mock_vachana_class.return_value = mock_instance
+        
+        # Create TTS instance
+        tts = TTS(pretrained="vachana")
+        
+        # Call tts method with waveform return type
+        result = tts.tts("สวัสดี", speaker_idx="th_f_1", return_type="waveform")
+        
+        # Verify the return type was set correctly
+        call_args = mock_instance.call_args
+        self.assertEqual(call_args.kwargs['return_type'], "waveform")
+
 
 if __name__ == '__main__':
     unittest.main()