path: root/crates/yt_dlp/src/lib.rs

// yt - A fully featured command line YouTube client
//
// Copyright (C) 2025 Benedikt Peetz <benedikt.peetz@b-peetz.de>
// SPDX-License-Identifier: GPL-3.0-or-later
//
// This file is part of Yt.
//
// You should have received a copy of the License along with this program.
// If not, see <https://www.gnu.org/licenses/gpl-3.0.txt>.

//! The `yt_dlp` interface is completely contained in the [`YoutubeDL`] structure.

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 url::Url;

use crate::{
    info_json::{InfoJson, json_dumps, json_loads},
    python_error::PythonError,
};

pub mod info_json;
pub mod options;
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),
            None => panic!(
                concat!(
                    "Expected '",
                    $name,
                    "' to be a key for the'",
                    stringify!($value),
                    "' object: {:#?}"
                ),
                $value
            ),
        }
    }};
}

#[macro_export]
macro_rules! json_cast {
    ($value:expr, $into:ident) => {{
        match $value.$into() {
            Some(result) => result,
            None => panic!(
                concat!(
                    "Expected to be able to cast value ({:#?}) ",
                    stringify!($into)
                ),
                $value
            ),
        }
    }};
}

/// The core of the `yt_dlp` interface.
pub struct YoutubeDL {
    interpreter: Interpreter,
    youtube_dl_class: PyObjectRef,
    yt_dlp_module: PyObjectRef,
    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
    ///
    /// 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()
    }

    /// Download a given list of URLs.
    /// Returns the paths they were downloaded to.
    ///
    /// # Errors
    /// If one of the downloads error.
    pub fn download(&self, urls: &[Url]) -> Result<Vec<PathBuf>, extract_info::Error> {
        let mut out_paths = Vec::with_capacity(urls.len());

        for url in urls {
            info!("Started downloading url: '{url}'");
            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))
            } else {
                PathBuf::from(json_get!(
                    json_cast!(
                        json_get!(info_json, "requested_downloads", as_array)[0],
                        as_object
                    ),
                    "filename",
                    as_str
                ))
            };

            out_paths.push(result_string);
            info!("Finished downloading url");
        }

        Ok(out_paths)
    }

    /// `extract_info(self, url, download=True, ie_key=None, extra_info=None, process=True, force_generic_extractor=False)`
    ///
    /// Extract and return the information dictionary of the URL
    ///
    /// Arguments:
    /// - `url`          URL to extract
    ///
    /// Keyword arguments:
    /// :`download`     Whether to download videos
    /// :`process`      Whether to resolve all unresolved references (URLs, playlist items).
    ///                 Must be True for download to work
    ///
    /// # Panics
    /// If expectations about python fail to hold.
    ///
    /// # Errors
    /// If python operations fail.
    pub fn extract_info(
        &self,
        url: &Url,
        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);

            let inner = self
                .youtube_dl_class
                .get_attr("extract_info", vm)
                .map_err(|exc| PythonError::from_exception(vm, &exc))?;
            let result = inner
                .call_with_args(fun_args, vm)
                .map_err(|exc| PythonError::from_exception(vm, &exc))?
                .downcast::<PyDict>()
                .expect("This is a dict");

            // Resolve the generator object
            if let Ok(generator) = result.get_item("entries", vm) {
                if generator.payload_is::<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")
                    });

                    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);

                        if out.len() == max_backlog {
                            break;
                        }
                    }
                    result
                        .set_item("entries", vm.new_pyobj(out), vm)
                        .map_err(|exc| PythonError::from_exception(vm, &exc))?;
                }
            }

            let result = self.prepare_info_json(result, vm)?;

            Ok(result)
        })
    }

    /// Take the (potentially modified) result of the information extractor (i.e.,
    /// [`Self::extract_info`] with `process` and `download` set to false)
    /// and resolve all unresolved references (URLs,
    /// playlist items).
    ///
    /// It will also download the videos if 'download' is true.
    /// Returns the resolved `ie_result`.
    ///
    /// # Panics
    /// If expectations about python fail to hold.
    ///
    /// # Errors
    /// If python operations fail.
    pub fn process_ie_result(
        &self,
        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);

            let inner = self
                .youtube_dl_class
                .get_attr("process_ie_result", vm)
                .map_err(|exc| PythonError::from_exception(vm, &exc))?;
            let result = inner
                .call_with_args(fun_args, vm)
                .map_err(|exc| PythonError::from_exception(vm, &exc))?
                .downcast::<PyDict>()
                .expect("This is a dict");

            let result = self.prepare_info_json(result, vm)?;

            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");

        Ok(json_dumps(result, vm))
    }
}

#[allow(missing_docs)]
pub mod process_ie_result {
    use crate::{prepare, python_error::PythonError};

    #[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 extract_info {
    use crate::{prepare, python_error::PythonError};

    #[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::python_error::PythonError;

    #[derive(Debug, thiserror::Error)]
    pub enum Error {
        #[error(transparent)]
        Python(#[from] PythonError),
    }
}