diff options
Diffstat (limited to '')
| -rw-r--r-- | crates/yt_dlp/src/lib.rs | 260 |
1 files changed, 162 insertions, 98 deletions
diff --git a/crates/yt_dlp/src/lib.rs b/crates/yt_dlp/src/lib.rs index a1db606..4b252de 100644 --- a/crates/yt_dlp/src/lib.rs +++ b/crates/yt_dlp/src/lib.rs @@ -12,18 +12,16 @@ use std::path::PathBuf; -use indexmap::IndexMap; -use log::info; -use rustpython::vm::{ - Interpreter, PyObjectRef, PyRef, VirtualMachine, - builtins::{PyDict, PyList, PyStr}, - function::{FuncArgs, KwArgs, PosArgs}, +use log::{debug, info}; +use pyo3::{ + Bound, Py, PyAny, Python, intern, + types::{PyAnyMethods, PyDict, PyIterator, PyList}, }; use url::Url; use crate::{ info_json::{InfoJson, json_dumps, json_loads}, - python_error::PythonError, + python_error::{IntoPythonError, PythonError}, }; pub mod info_json; @@ -32,18 +30,16 @@ pub mod post_processors; pub mod progress_hook; pub mod python_error; -mod logging; - #[macro_export] macro_rules! json_get { ($value:expr, $name:literal, $into:ident) => {{ match $value.get($name) { - Some(val) => $crate::json_cast!(val, $into), + Some(val) => $crate::json_cast!(@log_key $name, val, $into), None => panic!( concat!( "Expected '", $name, - "' to be a key for the'", + "' to be a key for the '", stringify!($value), "' object: {:#?}" ), @@ -54,51 +50,86 @@ macro_rules! json_get { } #[macro_export] +macro_rules! json_try_get { + ($value:expr, $name:literal, $into:ident) => {{ + if let Some(val) = $value.get($name) { + if val.is_null() { + None + } else { + Some(json_cast!(@log_key $name, val, $into)) + } + } else { + None + } + }}; +} + +#[macro_export] macro_rules! json_cast { ($value:expr, $into:ident) => {{ + let value_name = stringify!($value); + json_cast!(@log_key value_name, $value, $into) + }}; + + (@log_key $name:expr, $value:expr, $into:ident) => {{ match $value.$into() { Some(result) => result, None => panic!( concat!( - "Expected to be able to cast value ({:#?}) ", + "Expected to be able to cast '{}' value (which is '{:?}') ", stringify!($into) ), + $name, $value ), } }}; } +macro_rules! py_kw_args { + ($py:expr => $($kw_arg_name:ident = $kw_arg_val:expr),*) => {{ + use $crate::python_error::IntoPythonError; + + let dict = PyDict::new($py); + + $( + dict.set_item(stringify!($kw_arg_name), $kw_arg_val).wrap_exc($py)?; + )* + + Some(dict) + } + .as_ref()}; +} +pub(crate) use py_kw_args; + /// The core of the `yt_dlp` interface. +#[derive(Debug)] pub struct YoutubeDL { - interpreter: Interpreter, - youtube_dl_class: PyObjectRef, - yt_dlp_module: PyObjectRef, + inner: Py<PyAny>, options: serde_json::Map<String, serde_json::Value>, } -impl std::fmt::Debug for YoutubeDL { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - // TODO(@bpeetz): Use something useful here. <2025-06-13> - f.write_str("YoutubeDL") - } -} - impl YoutubeDL { - /// # Panics + /// Fetch the underlying `yt_dlp` and `python` version. /// - /// If `yt_dlp` changed their location or type of `__version__`. - pub fn version(&self) -> String { - let str_ref: PyRef<PyStr> = self.interpreter.enter_and_expect( - |vm| { - let version_module = self.yt_dlp_module.get_attr("version", vm)?; - let version = version_module.get_attr("__version__", vm)?; - let version = version.downcast().expect("This should always be a string"); - Ok(version) - }, - "yt_dlp version location has changed", - ); - str_ref.to_string() + /// # Errors + /// If python attribute access fails. + pub fn version(&self) -> Result<(String, String), PythonError> { + Python::attach(|py| { + let yt_dlp = py + .import(intern!(py, "yt_dlp")) + .wrap_exc(py)? + .getattr(intern!(py, "version")) + .wrap_exc(py)? + .getattr(intern!(py, "__version__")) + .wrap_exc(py)? + .extract() + .wrap_exc(py)?; + + let python = py.version(); + + Ok((yt_dlp, python.to_owned())) + }) } /// Download a given list of URLs. @@ -114,8 +145,9 @@ impl YoutubeDL { let info_json = self.extract_info(url, true, true)?; // Try to work around yt-dlp type weirdness - let result_string = if let Some(filename) = info_json.get("filename") { - PathBuf::from(json_cast!(filename, as_str)) + let result_string = if let Some(filename) = json_try_get!(info_json, "filename", as_str) + { + PathBuf::from(filename) } else { PathBuf::from(json_get!( json_cast!( @@ -157,55 +189,63 @@ impl YoutubeDL { download: bool, process: bool, ) -> Result<InfoJson, extract_info::Error> { - self.interpreter.enter(|vm| { - let pos_args = PosArgs::new(vec![vm.new_pyobj(url.to_string())]); - - let kw_args = KwArgs::new({ - let mut map = IndexMap::new(); - map.insert("download".to_owned(), vm.new_pyobj(download)); - map.insert("process".to_owned(), vm.new_pyobj(process)); - map - }); - - let fun_args = FuncArgs::new(pos_args, kw_args); - + Python::attach(|py| { let inner = self - .youtube_dl_class - .get_attr("extract_info", vm) - .map_err(|exc| PythonError::from_exception(vm, &exc))?; + .inner + .bind(py) + .getattr(intern!(py, "extract_info")) + .wrap_exc(py)?; + let result = inner - .call_with_args(fun_args, vm) - .map_err(|exc| PythonError::from_exception(vm, &exc))? - .downcast::<PyDict>() + .call( + (url.to_string(),), + py_kw_args!(py => download = download, process = process), + ) + .wrap_exc(py)? + .cast_into::<PyDict>() .expect("This is a dict"); // Resolve the generator object - if let Ok(generator) = result.get_item("entries", vm) { - if generator.payload_is::<PyList>() { + if let Ok(generator) = result.get_item(intern!(py, "entries")) { + if generator.is_instance_of::<PyList>() { // already resolved. Do nothing - } else { - let max_backlog = self.options.get("playlistend").map_or(10000, |value| { - usize::try_from(value.as_u64().expect("Works")).expect("Should work") - }); + } else if let Ok(generator) = generator.cast::<PyIterator>() { + // A python generator object. + let max_backlog = json_try_get!(self.options, "playlistend", as_u64) + .map_or(10000, |playlistend| { + usize::try_from(playlistend).expect("Should work") + }); let mut out = vec![]; - let next = generator - .get_attr("__next__", vm) - .map_err(|exc| PythonError::from_exception(vm, &exc))?; - while let Ok(output) = next.call((), vm) { - out.push(output); + for output in generator { + out.push(output.wrap_exc(py)?); if out.len() == max_backlog { break; } } + + result.set_item(intern!(py, "entries"), out).wrap_exc(py)?; + } else { + // Probably some sort of paged list (`OnDemand` or otherwise) + let max_backlog = json_try_get!(self.options, "playlistend", as_u64) + .map_or(10000, |playlistend| { + usize::try_from(playlistend).expect("Should work") + }); + + let next = generator.getattr(intern!(py, "getslice")).wrap_exc(py)?; + + let output = next + .call((), py_kw_args!(py => start = 0, end = max_backlog)) + .wrap_exc(py)?; + result - .set_item("entries", vm.new_pyobj(out), vm) - .map_err(|exc| PythonError::from_exception(vm, &exc))?; + .set_item(intern!(py, "entries"), output) + .wrap_exc(py)?; } } - let result = self.prepare_info_json(result, vm)?; + let result = self.prepare_info_json(&result, py)?; Ok(result) }) @@ -229,54 +269,78 @@ impl YoutubeDL { ie_result: InfoJson, download: bool, ) -> Result<InfoJson, process_ie_result::Error> { - self.interpreter.enter(|vm| { - let pos_args = PosArgs::new(vec![vm.new_pyobj(json_loads(ie_result, vm))]); - - let kw_args = KwArgs::new({ - let mut map = IndexMap::new(); - map.insert("download".to_owned(), vm.new_pyobj(download)); - map - }); - - let fun_args = FuncArgs::new(pos_args, kw_args); - + Python::attach(|py| { let inner = self - .youtube_dl_class - .get_attr("process_ie_result", vm) - .map_err(|exc| PythonError::from_exception(vm, &exc))?; + .inner + .bind(py) + .getattr(intern!(py, "process_ie_result")) + .wrap_exc(py)?; + let result = inner - .call_with_args(fun_args, vm) - .map_err(|exc| PythonError::from_exception(vm, &exc))? - .downcast::<PyDict>() + .call( + (json_loads(ie_result, py),), + py_kw_args!(py => download = download), + ) + .wrap_exc(py)? + .cast_into::<PyDict>() .expect("This is a dict"); - let result = self.prepare_info_json(result, vm)?; + let result = self.prepare_info_json(&result, py)?; Ok(result) }) } - fn prepare_info_json( + /// Close this [`YoutubeDL`] instance, and stop all currently running downloads. + /// + /// # Errors + /// If python operations fail. + pub fn close(&self) -> Result<(), close::Error> { + Python::attach(|py| { + debug!("Closing YoutubeDL."); + + let inner = self + .inner + .bind(py) + .getattr(intern!(py, "close")) + .wrap_exc(py)?; + + inner.call0().wrap_exc(py)?; + + Ok(()) + }) + } + + fn prepare_info_json<'py>( &self, - info: PyRef<PyDict>, - vm: &VirtualMachine, + info: &Bound<'py, PyDict>, + py: Python<'py>, ) -> Result<InfoJson, prepare::Error> { let sanitize = self - .youtube_dl_class - .get_attr("sanitize_info", vm) - .map_err(|exc| PythonError::from_exception(vm, &exc))?; + .inner + .bind(py) + .getattr(intern!(py, "sanitize_info")) + .wrap_exc(py)?; - let value = sanitize - .call((info,), vm) - .map_err(|exc| PythonError::from_exception(vm, &exc))?; + let value = sanitize.call((info,), None).wrap_exc(py)?; - let result = value.downcast::<PyDict>().expect("This should stay a dict"); + let result = value.cast::<PyDict>().expect("This should stay a dict"); - Ok(json_dumps(result, vm)) + Ok(json_dumps(result)) } } #[allow(missing_docs)] +pub mod close { + use crate::python_error::PythonError; + + #[derive(Debug, thiserror::Error)] + pub enum Error { + #[error(transparent)] + Python(#[from] PythonError), + } +} +#[allow(missing_docs)] pub mod process_ie_result { use crate::{prepare, python_error::PythonError}; |