diff options
| -rw-r--r-- | crates/yt_dlp/src/lib.rs | 29 |
1 files changed, 24 insertions, 5 deletions
diff --git a/crates/yt_dlp/src/lib.rs b/crates/yt_dlp/src/lib.rs index b0e157c..cd1d0d2 100644 --- a/crates/yt_dlp/src/lib.rs +++ b/crates/yt_dlp/src/lib.rs @@ -360,14 +360,33 @@ pub async fn extract_info( kwargs.set_item("download", download)?; kwargs.set_item("process", process)?; - let result = instance.call_method("extract_info", args, Some(&kwargs))?; + let result = instance + .call_method("extract_info", args, Some(&kwargs))? + .downcast_into::<PyDict>() + .expect("This is a dict"); + + // Resolve the generator object + if let Some(generator) = result.get_item("entries")? { + if generator.is_instance_of::<PyList>() { + // already resolved. Do nothing + } else { + let max_backlog = yt_dlp_opts.get("playlistend").map_or(10000, |value| { + usize::try_from(value.as_u64().expect("Works")).expect("Should work") + }); + + let mut out = vec![]; + while let Ok(output) = generator.call_method0("__next__") { + out.push(output); - // Remove the `<generator at 0xsome_hex>`, by setting it to null - if !process { - result.set_item("entries", ())?; + if out.len() == max_backlog { + break; + } + } + result.set_item("entries", out)?; + } } - let result_str = json_dumps(py, result)?; + let result_str = json_dumps(py, result.into_any())?; if let Ok(confirm) = env::var("YT_STORE_INFO_JSON") { if confirm == "yes" { |