From 23b1e17cf2a8ddb8ce6a62161a5fa8bacf9a53c5 Mon Sep 17 00:00:00 2001
From: Laurent <laurent.mazare@gmail.com>
Date: Tue, 20 Aug 2024 21:46:56 +0200
Subject: [PATCH] Potentially resample when writing opus files.

---
 .gitignore    |  2 ++
 src/lib.rs    | 10 +++++++---
 src/opus.rs   | 17 ++++++++++++++++-
 test/basic.py |  1 +
 4 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index d75ea22..e0325a4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,3 +16,5 @@ Cargo.lock
 __pycache__
 
 bria.mp3
+bria.wav
+bria.opus
diff --git a/src/lib.rs b/src/lib.rs
index 1fcff2e..5cafb14 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -191,11 +191,12 @@ fn write_wav(
 }
 
 #[pyfunction]
-#[pyo3(signature = (filename, data, sample_rate))]
+#[pyo3(signature = (filename, data, sample_rate, resample_to=None))]
 fn write_opus(
     filename: std::path::PathBuf,
     data: numpy::PyReadonlyArray1<f32>,
     sample_rate: u32,
+    resample_to: Option<u32>,
 ) -> PyResult<()> {
     let w = std::fs::File::create(&filename).w_f(filename.as_path())?;
     let mut w = std::io::BufWriter::new(w);
@@ -203,9 +204,12 @@ fn write_opus(
     match data.as_slice() {
         None => {
             let data = data.to_vec();
-            opus::write_ogg(&mut w, data.as_ref(), sample_rate).w_f(filename.as_path())?
+            opus::write_ogg(&mut w, data.as_ref(), sample_rate, resample_to)
+                .w_f(filename.as_path())?
+        }
+        Some(data) => {
+            opus::write_ogg(&mut w, data, sample_rate, resample_to).w_f(filename.as_path())?
         }
-        Some(data) => opus::write_ogg(&mut w, data, sample_rate).w_f(filename.as_path())?,
     }
     Ok(())
 }
diff --git a/src/opus.rs b/src/opus.rs
index 8ca28ac..4799068 100644
--- a/src/opus.rs
+++ b/src/opus.rs
@@ -120,7 +120,7 @@ fn write_opus_tags<W: std::io::Write>(w: &mut W) -> std::io::Result<()> {
     Ok(())
 }
 
-pub fn write_ogg<W: std::io::Write>(w: &mut W, pcm: &[f32], sample_rate: u32) -> Result<()> {
+fn write_ogg_<W: std::io::Write>(w: &mut W, pcm: &[f32], sample_rate: u32) -> Result<()> {
     let mut pw = ogg::PacketWriter::new(w);
 
     // Write the opus headers and tags
@@ -152,3 +152,18 @@ pub fn write_ogg<W: std::io::Write>(w: &mut W, pcm: &[f32], sample_rate: u32) ->
 
     Ok(())
 }
+
+pub fn write_ogg<W: std::io::Write>(
+    w: &mut W,
+    pcm: &[f32],
+    sample_rate: u32,
+    resample_to: Option<u32>,
+) -> Result<()> {
+    match resample_to {
+        None => write_ogg_(w, pcm, sample_rate),
+        Some(resample_to) => {
+            let pcm = crate::audio::resample(pcm, sample_rate as usize, resample_to as usize)?;
+            write_ogg_(w, &pcm, resample_to)
+        }
+    }
+}
diff --git a/test/basic.py b/test/basic.py
index 8c99eeb..7307bd2 100644
--- a/test/basic.py
+++ b/test/basic.py
@@ -14,3 +14,4 @@
 print(data.shape, sr)
 
 sphn.write_wav("bria.wav", data[0], sr)
+sphn.write_opus("bria.opus", data[0], sr, resample_to=48000)