path: root/crates/yt_dlp/src/options.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>.

use std::env;

use indexmap::IndexMap;
use log::{Level, debug, error, log_enabled};
use rustpython::{
    InterpreterConfig,
    vm::{
        self, PyObjectRef, PyRef, PyResult, VirtualMachine,
        builtins::{PyBaseException, PyStr},
        function::{FuncArgs, KwArgs, PosArgs},
    },
};

use crate::{
    YoutubeDL, json_loads, logging::setup_logging, post_processors, python_error::process_exception,
};

/// Wrap your function with [`mk_python_function`].
pub type ProgressHookFunction = fn(input: FuncArgs, vm: &VirtualMachine);

pub type PostProcessorFunction = fn(vm: &VirtualMachine) -> PyResult<PyObjectRef>;

/// Options, that are used to customize the download behaviour.
///
/// In the future, this might get a Builder api.
///
/// See `help(yt_dlp.YoutubeDL())` from python for a full list of available options.
#[derive(Default, Debug)]
pub struct YoutubeDLOptions {
    options: serde_json::Map<String, serde_json::Value>,
    progress_hook: Option<ProgressHookFunction>,
    post_processors: Vec<PostProcessorFunction>,
}

impl YoutubeDLOptions {
    #[must_use]
    pub fn new() -> Self {
        let me = Self {
            options: serde_json::Map::new(),
            progress_hook: None,
            post_processors: vec![],
        };

        me.with_post_processor(post_processors::dearrow::process)
    }

    #[must_use]
    pub fn set(self, key: impl Into<String>, value: impl Into<serde_json::Value>) -> Self {
        let mut options = self.options;
        options.insert(key.into(), value.into());

        Self { options, ..self }
    }

    #[must_use]
    pub fn with_progress_hook(self, progress_hook: ProgressHookFunction) -> Self {
        if let Some(_previous_hook) = self.progress_hook {
            todo!()
        } else {
            Self {
                progress_hook: Some(progress_hook),
                ..self
            }
        }
    }

    #[must_use]
    pub fn with_post_processor(mut self, pp: PostProcessorFunction) -> Self {
        self.post_processors.push(pp);
        self
    }

    /// # Errors
    /// If the underlying [`YoutubeDL::from_options`] errors.
    pub fn build(self) -> Result<YoutubeDL, build::Error> {
        YoutubeDL::from_options(self)
    }

    #[must_use]
    pub fn from_json_options(options: serde_json::Map<String, serde_json::Value>) -> Self {
        Self {
            options,
            ..Self::new()
        }
    }

    #[must_use]
    pub fn get(&self, key: &str) -> Option<&serde_json::Value> {
        self.options.get(key)
    }
}

impl YoutubeDL {
    /// Construct this instance from options.
    ///
    /// # Panics
    /// If `yt_dlp` changed their interface.
    ///
    /// # Errors
    /// If a python call fails.
    #[allow(clippy::too_many_lines)]
    pub fn from_options(options: YoutubeDLOptions) -> Result<Self, build::Error> {
        let mut settings = vm::Settings::default();
        if let Ok(python_path) = env::var("PYTHONPATH") {
            for path in python_path.split(':') {
                settings.path_list.push(path.to_owned());
            }
        } else {
            error!(
                "No PYTHONPATH found or invalid utf8. \
                This means, that you probably did not \
                supply a yt_dlp python package!"
            );
        }

        settings.install_signal_handlers = false;

        // NOTE(@bpeetz): Another value leads to an internal codegen error. <2025-06-13>
        settings.optimize = 0;

        settings.isolated = true;

        let interpreter = InterpreterConfig::new()
            .init_stdlib()
            .settings(settings)
            .interpreter();

        let output_options = options.options.clone();

        let (yt_dlp_module, youtube_dl_class) = match interpreter.enter(|vm| {
            let yt_dlp_module = vm.import("yt_dlp", 0)?;
            let class = yt_dlp_module.get_attr("YoutubeDL", vm)?;

            let opts = json_loads(options.options, vm);

            {
                // Setup the progress hook
                if let Some(function) = options.progress_hook {
                    opts.get_or_insert(vm, vm.new_pyobj("progress_hooks"), || {
                        let hook: PyObjectRef = vm.new_function("progress_hook", function).into();
                        vm.new_pyobj(vec![hook])
                    })
                    .expect("Should work?");
                }
            }

            {
                // Unconditionally set a logger.
                // Otherwise, yt_dlp will log to stderr.

                /// Is the specified record to be logged? Returns false for no,
                /// true for yes. Filters can either modify log records in-place or
                /// return a completely different record instance which will replace
                /// the original log record in any future processing of the event.
                fn filter_error_log(mut input: FuncArgs, vm: &VirtualMachine) -> bool {
                    let record = input.args.remove(0);

                    // Filter out all error logs (they are propagated as rust errors)
                    let levelname: PyRef<PyStr> = record
                        .get_attr("levelname", vm)
                        .expect("This should exist")
                        .downcast()
                        .expect("This should be a String");

                    let return_value = levelname.as_str() != "ERROR";

                    if log_enabled!(Level::Debug) && !return_value {
                        let message: String = {
                            let get_message = record.get_attr("getMessage", vm).expect("Is set");
                            let message: PyRef<PyStr> = get_message
                                .call((), vm)
                                .expect("Can be called")
                                .downcast()
                                .expect("Downcasting works");

                            message.as_str().to_owned()
                        };

                        debug!("Swollowed error message: '{message}'");
                    }
                    return_value
                }

                let logging = setup_logging(vm, "yt_dlp")?;
                let ytdl_logger = {
                    let get_logger = logging.get_item("getLogger", vm)?;
                    get_logger.call(("yt_dlp",), vm)?
                };

                {
                    let args = FuncArgs::new(
                        PosArgs::new(vec![]),
                        KwArgs::new({
                            let mut map = IndexMap::new();
                            // Ensure that all events are logged by setting
                            // the log level to NOTSET (we filter on rust's side)
                            map.insert("level".to_owned(), vm.new_pyobj(0));
                            map
                        }),
                    );

                    let basic_config = logging.get_item("basicConfig", vm)?;
                    basic_config.call(args, vm)?;
                }

                {
                    let add_filter = ytdl_logger.get_attr("addFilter", vm)?;
                    add_filter.call(
                        (vm.new_function("yt_dlp_error_filter", filter_error_log),),
                        vm,
                    )?;
                }

                opts.set_item("logger", ytdl_logger, vm)?;
            }

            let youtube_dl_class = class.call((opts,), vm)?;

            {
                // Setup the post processors

                let add_post_processor_fun = youtube_dl_class.get_attr("add_post_processor", vm)?;

                for pp in options.post_processors {
                    let args = {
                        FuncArgs::new(
                            PosArgs::new(vec![pp(vm)?]),
                            KwArgs::new({
                                let mut map = IndexMap::new();
                                //  "when" can take any value in yt_dlp.utils.POSTPROCESS_WHEN
                                map.insert("when".to_owned(), vm.new_pyobj("pre_process"));
                                map
                            }),
                        )
                    };

                    add_post_processor_fun.call(args, vm)?;
                }
            }

            Ok::<_, PyRef<PyBaseException>>((yt_dlp_module, youtube_dl_class))
        }) {
            Ok(ok) => Ok(ok),
            Err(err) => {
                // TODO(@bpeetz): Do we want to run `interpreter.finalize` here? <2025-06-14>
                // interpreter.finalize(Some(err));
                interpreter.enter(|vm| {
                    let buffer = process_exception(vm, &err);
                    Err(build::Error::Python(buffer))
                })
            }
        }?;

        Ok(Self {
            interpreter,
            youtube_dl_class,
            yt_dlp_module,
            options: output_options,
        })
    }
}

#[allow(missing_docs)]
pub mod build {
    #[derive(Debug, thiserror::Error)]
    pub enum Error {
        #[error("Python threw an exception: {0}")]
        Python(String),
    }
}