refactor(yt_dlp/lib): De-duplicate the info json sanitize code

1 files changed, 41 insertions, 14 deletions

diff --git a/crates/yt_dlp/src/lib.rs b/crates/yt_dlp/src/lib.rs
index 536b0d1..0f40f0a 100644
--- a/crates/yt_dlp/src/lib.rs
+++ b/crates/yt_dlp/src/lib.rs
@@ -332,14 +332,7 @@ impl YoutubeDL {
                 }
             }
 
-            let result = {
-                let sanitize = self.youtube_dl_class.get_attr("sanitize_info", vm)?;
-                let value = sanitize.call((result,), vm)?;
-
-                value.downcast::<PyDict>().expect("This should stay a dict")
-            };
-
-            let result_json = json_dumps(result, vm);
+            let result = self.prepare_info_json(result, vm)?;
 
             Ok(result)
         })
@@ -384,19 +377,37 @@ impl YoutubeDL {
                 .downcast::<PyDict>()
                 .expect("This is a dict");
 
-            let result = {
-                let sanitize = self.youtube_dl_class.get_attr("sanitize_info", vm)?;
-                let value = sanitize.call((result,), vm)?;
+            let result = self.prepare_info_json(result, vm)?;
 
-                value.downcast::<PyDict>().expect("This should stay a dict")
-            };
             Ok(result)
         })
     }
 
+    fn prepare_info_json(
+        &self,
+        info: PyRef<PyDict>,
+        vm: &VirtualMachine,
+    ) -> Result<InfoJson, prepare::Error> {
+        let sanitize = self
+            .youtube_dl_class
+            .get_attr("sanitize_info", vm)
+            .map_err(|exc| PythonError::from_exception(vm, &exc))?;
+
+        let value = sanitize
+            .call((info,), vm)
+            .map_err(|exc| PythonError::from_exception(vm, &exc))?;
 
+        let result = value.downcast::<PyDict>().expect("This should stay a dict");
+
+        let json = json_dumps(result, vm);
+
+        {
         }
 
+        Ok(json)
+    }
+}
+
 #[derive(thiserror::Error, Debug)]
 pub struct PythonError(pub String);
 
@@ -421,11 +432,27 @@ pub mod process_ie_result {
     pub enum Error {
         #[error(transparent)]
         Python(#[from] PythonError),
+
+        #[error("Failed to prepare the info json")]
+        InfoJsonPrepare(#[from] prepare::Error),
     }
 }
 #[allow(missing_docs)]
 pub mod extract_info {
-    use crate::{PythonError};
+    use crate::{PythonError, prepare};
+
+    #[derive(Debug, thiserror::Error)]
+    pub enum Error {
+        #[error(transparent)]
+        Python(#[from] PythonError),
+
+        #[error("Failed to prepare the info json")]
+        InfoJsonPrepare(#[from] prepare::Error),
+    }
+}
+#[allow(missing_docs)]
+pub mod prepare {
+    use crate::{PythonError, post_processors};
 
     #[derive(Debug, thiserror::Error)]
     pub enum Error {