Convert full_get_segment_* and similar to use an internal helper in…

…stead of duplicating code
tazz4843 · Feb 19, 2025 · 9a96b0e · 9a96b0e
1 parent 099faf4
commit 9a96b0e
Showing 1 changed file with 29 additions and 33 deletions.
diff --git a/src/whisper_state.rs b/src/whisper_state.rs
@@ -346,68 +346,64 @@ impl WhisperState {
         Ok(unsafe { whisper_rs_sys::whisper_full_get_segment_t1_from_state(self.ptr, segment) })
     }
 
-    /// Get the text of the specified segment.
+    fn full_get_segment_raw(&self, segment: c_int) -> Result<&CStr, WhisperError> {
+        let ret =
+            unsafe { whisper_rs_sys::whisper_full_get_segment_text_from_state(self.ptr, segment) };
+        if ret.is_null() {
+            return Err(WhisperError::NullPointer);
+        }
+        unsafe { Ok(CStr::from_ptr(ret)) }
+    }
+
+    /// Get the raw bytes of the specified segment.
     ///
     /// # Arguments
     /// * segment: Segment index.
     ///
     /// # Returns
-    /// Ok(String) on success, Err(WhisperError) on failure.
+    /// `Ok(Vec<u8>)` on success, with the returned bytes or
+    /// `Err(WhisperError::NullPointer)` on failure (this is the only possible error)
     ///
     /// # C++ equivalent
     /// `const char * whisper_full_get_segment_text(struct whisper_context * ctx, int i_segment)`
-    pub fn full_get_segment_text(&self, segment: c_int) -> Result<String, WhisperError> {
-        let ret =
-            unsafe { whisper_rs_sys::whisper_full_get_segment_text_from_state(self.ptr, segment) };
-        if ret.is_null() {
-            return Err(WhisperError::NullPointer);
-        }
-        let c_str = unsafe { CStr::from_ptr(ret) };
-        let r_str = c_str.to_str()?;
-        Ok(r_str.to_string())
+    pub fn full_get_segment_bytes(&self, segment: c_int) -> Result<Vec<u8>, WhisperError> {
+        Ok(self.full_get_segment_raw(segment)?.to_bytes().to_vec())
     }
 
     /// Get the text of the specified segment.
-    /// This function differs from [WhisperState::full_get_segment_text]
-    /// in that it ignores invalid UTF-8 in whisper strings,
-    /// instead opting to replace it with the replacement character.
     ///
     /// # Arguments
     /// * segment: Segment index.
     ///
     /// # Returns
-    /// Ok(String) on success, Err(WhisperError) on failure.
+    /// `Ok(String)` on success, with the UTF-8 validated string, or
+    /// `Err(WhisperError)` on failure (either `NullPointer` or `InvalidUtf8`)
     ///
     /// # C++ equivalent
     /// `const char * whisper_full_get_segment_text(struct whisper_context * ctx, int i_segment)`
-    pub fn full_get_segment_text_lossy(&self, segment: c_int) -> Result<String, WhisperError> {
-        let ret =
-            unsafe { whisper_rs_sys::whisper_full_get_segment_text_from_state(self.ptr, segment) };
-        if ret.is_null() {
-            return Err(WhisperError::NullPointer);
-        }
-        let c_str = unsafe { CStr::from_ptr(ret) };
-        Ok(c_str.to_string_lossy().to_string())
+    pub fn full_get_segment_text(&self, segment: c_int) -> Result<String, WhisperError> {
+        Ok(self.full_get_segment_raw(segment)?.to_str()?.to_string())
     }
 
-    /// Get the bytes of the specified segment.
+    /// Get the text of the specified segment.
+    /// This function differs from [WhisperState::full_get_segment_text]
+    /// in that it ignores invalid UTF-8 in whisper strings,
+    /// instead opting to replace it with the replacement character.
     ///
     /// # Arguments
     /// * segment: Segment index.
     ///
     /// # Returns
-    /// `Ok(Vec<u8>)` on success, `Err(WhisperError)` on failure.
+    /// `Ok(String)` on success, or
+    /// `Err(WhisperError::NullPointer)` on failure (this is the only possible error)
     ///
     /// # C++ equivalent
     /// `const char * whisper_full_get_segment_text(struct whisper_context * ctx, int i_segment)`
-    pub fn full_get_segment_bytes(&self, segment: c_int) -> Result<Vec<u8>, WhisperError> {
-        let ret =
-            unsafe { whisper_rs_sys::whisper_full_get_segment_text_from_state(self.ptr, segment) };
-        if ret.is_null() {
-            return Err(WhisperError::NullPointer);
-        }
-        let c_str = unsafe { CStr::from_ptr(ret) };
-        Ok(c_str.to_bytes().to_vec())
+    pub fn full_get_segment_text_lossy(&self, segment: c_int) -> Result<String, WhisperError> {
+        Ok(self
+            .full_get_segment_raw(segment)?
+            .to_string_lossy()
+            .to_string())
     }
 
     /// Get number of tokens in the specified segment.