diff options
Diffstat (limited to 'crates/yt_dlp/crates/pyo3-pylogger/src/lib.rs')
| -rw-r--r-- | crates/yt_dlp/crates/pyo3-pylogger/src/lib.rs | 201 |
1 files changed, 201 insertions, 0 deletions
diff --git a/crates/yt_dlp/crates/pyo3-pylogger/src/lib.rs b/crates/yt_dlp/crates/pyo3-pylogger/src/lib.rs new file mode 100644 index 0000000..b16e448 --- /dev/null +++ b/crates/yt_dlp/crates/pyo3-pylogger/src/lib.rs @@ -0,0 +1,201 @@ +use std::{ + ffi::CString, + sync::{self, OnceLock}, +}; + +use log::{debug, log_enabled}; +use pyo3::{ + Bound, Py, PyAny, PyResult, Python, pyfunction, + sync::OnceLockExt, + types::{PyAnyMethods, PyDict, PyListMethods, PyModuleMethods}, + wrap_pyfunction, +}; + +mod kv; +mod level; + +static LOGGER: sync::OnceLock<Py<PyAny>> = OnceLock::new(); + +/// 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. +#[pyfunction] +fn filter_error_log<'py>(record: Bound<'py, PyAny>) -> bool { + // Filter out all error logs (they are propagated as rust errors) + let levelname: String = record + .getattr("levelname") + .expect("This should exist") + .extract() + .expect("This should be a String"); + + let return_value = levelname.as_str() != "ERROR"; + + if log_enabled!(log::Level::Debug) && !return_value { + let message: String = { + let get_message = record.getattr("getMessage").expect("Is set"); + let message: String = get_message + .call((), None) + .expect("Can be called") + .extract() + .expect("Downcasting works"); + + message.as_str().to_owned() + }; + + debug!("Swollowed error message: '{message}'"); + } + return_value +} + +/// Consume a Python `logging.LogRecord` and emit a Rust `Log` instead. +#[pyfunction] +fn host_log(record: Bound<'_, PyAny>, rust_target: &str) -> PyResult<()> { + let level = record.getattr("levelno")?.extract()?; + let message = record.getattr("getMessage")?.call0()?.to_string(); + let pathname = record.getattr("pathname")?.extract::<String>()?; + let lineno = record.getattr("lineno")?.extract::<u32>()?; + + let logger_name = record.getattr("name")?.extract::<String>()?; + + let full_target: Option<String> = if logger_name.trim().is_empty() || logger_name == "root" { + None + } else { + // Libraries (ex: tracing_subscriber::filter::Directive) expect rust-style targets like foo::bar, + // and may not deal well with "." as a module separator: + let logger_name = logger_name.replace('.', "::"); + Some(format!("{rust_target}::{logger_name}")) + }; + let target = full_target.as_deref().unwrap_or(rust_target); + + handle_record(record, target, &message, lineno, &pathname, level)?; + + Ok(()) +} + +fn handle_record( + #[allow(unused_variables)] record: Bound<'_, PyAny>, + target: &str, + message: &str, + lineno: u32, + pathname: &str, + level: u8, +) -> PyResult<()> { + // If log feature is enabled, use log::logger + let level = crate::level::get_level(level).0; + + { + let mut metadata_builder = log::MetadataBuilder::new(); + metadata_builder.target(target); + metadata_builder.level(level); + + let mut record_builder = log::Record::builder(); + + { + let kv_args = kv::find_kv_args(&record)?; + + let kv_source = kv_args.map(kv::KVSource); + if let Some(kv_source) = kv_source { + log::logger().log( + &record_builder + .metadata(metadata_builder.build()) + .args(format_args!("{}", &message)) + .line(Some(lineno)) + .file(Some(pathname)) + .module_path(Some(pathname)) + .key_values(&kv_source) + .build(), + ); + return Ok(()); + } + } + + log::logger().log( + &record_builder + .metadata(metadata_builder.build()) + .args(format_args!("{}", &message)) + .line(Some(lineno)) + .file(Some(pathname)) + .module_path(Some(pathname)) + .build(), + ); + } + + Ok(()) +} + +/// Registers the host_log function in rust as the event handler for Python's logging logger +/// This function needs to be called from within a pyo3 context as early as possible to ensure logging messages +/// arrive to the rust consumer. +pub fn setup_logging<'py>(py: Python<'py>, target: &str) -> PyResult<Bound<'py, PyAny>> { + let logger = LOGGER + .get_or_init_py_attached(py, || match setup_logging_inner(py, target) { + Ok(ok) => ok.unbind(), + Err(err) => { + panic!("Failed to initialize logger: {}", err); + } + }) + .clone_ref(py); + + Ok(logger.into_bound(py)) +} + +fn setup_logging_inner<'py>(py: Python<'py>, target: &str) -> PyResult<Bound<'py, PyAny>> { + let logging = py.import("logging")?; + + logging.setattr("host_log", wrap_pyfunction!(host_log, &logging)?)?; + + #[allow(clippy::uninlined_format_args)] + let code = CString::new(format!( + r#" +class HostHandler(Handler): + def __init__(self, level=0): + super().__init__(level=level) + + def emit(self, record: LogRecord): + host_log(record, "{}") + +oldBasicConfig = basicConfig +def basicConfig(*pargs, **kwargs): + if "handlers" not in kwargs: + kwargs["handlers"] = [HostHandler()] + return oldBasicConfig(*pargs, **kwargs) +"#, + target + ))?; + + let logging_scope = logging.dict(); + py.run(&code, Some(&logging_scope), None)?; + + let all = logging.index()?; + all.append("HostHandler")?; + + let logger = { + let get_logger = logging_scope.get_item("getLogger")?; + get_logger.call((target,), None)? + }; + + { + let basic_config = logging_scope.get_item("basicConfig")?; + basic_config.call( + (), + { + let dict = PyDict::new(py); + + // Ensure that all events are logged by setting + // the log level to NOTSET (we filter on rust's side) + dict.set_item("level", 0)?; + + Some(dict) + } + .as_ref(), + )?; + } + + { + let add_filter = logger.getattr("addFilter")?; + add_filter.call((wrap_pyfunction!(filter_error_log, &logging)?,), None)?; + } + + Ok(logger) +} |