use crate::{

AsObject, Context, Py, PyObject, PyObjectRef, PyResult, VirtualMachine,

builtins::{

PyBaseExceptionRef, PyDictRef, PyListRef, PyStr, PyStrInterned, PyStrRef, PyTuple,

PyTupleRef, PyTypeRef,

},

convert::TryFromObject,

};

use core::sync::atomic::{AtomicUsize, Ordering};

use rustpython_common::lock::OnceCell;

pub struct WarningsState {

pub filters: PyListRef,

pub once_registry: PyDictRef,

pub default_action: PyStrRef,

pub filters_version: AtomicUsize,

pub context_var: OnceCell<PyObjectRef>,

lock_count: AtomicUsize,

}

impl WarningsState {

fn create_default_filters(ctx: &Context) -> PyListRef {

// init_filters(): non-debug default filter set.

ctx.new_list(vec![

ctx.new_tuple(vec![

ctx.new_str("default").into(),

ctx.none(),

ctx.exceptions.deprecation_warning.as_object().to_owned(),

ctx.new_str("__main__").into(),

ctx.new_int(0).into(),

])

.into(),

ctx.new_tuple(vec![

ctx.new_str("ignore").into(),

ctx.none(),

ctx.exceptions.deprecation_warning.as_object().to_owned(),

ctx.none(),

ctx.new_int(0).into(),

])

.into(),

ctx.new_tuple(vec![

ctx.new_str("ignore").into(),

ctx.none(),

ctx.exceptions

.pending_deprecation_warning

.as_object()

.to_owned(),

ctx.none(),

ctx.new_int(0).into(),

])

.into(),

ctx.new_tuple(vec![

ctx.new_str("ignore").into(),

ctx.none(),

ctx.exceptions.import_warning.as_object().to_owned(),

ctx.none(),

ctx.new_int(0).into(),

])

.into(),

ctx.new_tuple(vec![

ctx.new_str("ignore").into(),

ctx.none(),

ctx.exceptions.resource_warning.as_object().to_owned(),

ctx.none(),

ctx.new_int(0).into(),

])

.into(),

])

}

pub fn init_state(ctx: &Context) -> Self {

Self {

filters: Self::create_default_filters(ctx),

once_registry: ctx.new_dict(),

default_action: ctx.new_str("default"),

filters_version: AtomicUsize::new(0),

context_var: OnceCell::new(),

lock_count: AtomicUsize::new(0),

}

}

pub fn acquire_lock(&self) {

self.lock_count.fetch_add(1, Ordering::SeqCst);

}

pub fn release_lock(&self) -> bool {

let prev = self.lock_count.load(Ordering::SeqCst);

if prev == 0 {

return false;

}

self.lock_count.fetch_sub(1, Ordering::SeqCst);

true

}

pub fn filters_mutated(&self) {

self.filters_version.fetch_add(1, Ordering::SeqCst);

}

}

/// None matches everything; plain strings do exact comparison;

/// regex objects use .match().

fn check_matched(obj: &PyObject, arg: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {

if vm.is_none(obj) {

return Ok(true);

}

if obj.class().is(vm.ctx.types.str_type) {

return obj.rich_compare_bool(arg, crate::types::PyComparisonOp::Eq, vm);

}

let result = vm.call_method(obj, "match", (arg.to_owned(),))?;

result.is_true(vm)

}

fn get_warnings_attr(

vm: &VirtualMachine,

attr_name: &'static PyStrInterned,

try_import: bool,

) -> PyResult<Option<PyObjectRef>> {

let module = if try_import

&& !vm

.state

.finalizing

.load(core::sync::atomic::Ordering::SeqCst)

{

match vm.import("warnings", 0) {

Ok(module) => module,

Err(_) => return Ok(None),

}

} else {

match vm.sys_module.get_attr(identifier!(vm, modules), vm) {

Ok(modules) => match modules.get_item(vm.ctx.intern_str("warnings"), vm) {

Ok(module) => module,

Err(_) => return Ok(None),

},

Err(_) => return Ok(None),

}

};

match module.get_attr(attr_name, vm) {

Ok(attr) => Ok(Some(attr)),

Err(_) => Ok(None),

}

}

/// Get the warnings filters list from sys.modules['warnings'].filters,

/// falling back to vm.state.warnings.filters.

fn get_warnings_filters(vm: &VirtualMachine) -> PyResult<PyListRef> {

if let Some(filters_obj) = get_warnings_attr(vm, identifier!(&vm.ctx, filters), false)?

&& let Ok(filters) = filters_obj.try_into_value::<PyListRef>(vm)

{

return Ok(filters);

}

Ok(vm.state.warnings.filters.clone())

}

/// Get the default action from sys.modules['warnings']._defaultaction,

/// falling back to vm.state.warnings.default_action.

fn get_default_action(vm: &VirtualMachine) -> PyResult<PyObjectRef> {

if let Some(action) = get_warnings_attr(vm, identifier!(&vm.ctx, defaultaction), false)? {

if !action.class().is(vm.ctx.types.str_type) {

return Err(vm.new_type_error(format!(

"_warnings.defaultaction must be a string, not '{}'",

action.class().name()

)));

}

return Ok(action);

}

Ok(vm.state.warnings.default_action.clone().into())

}

/// Get the once registry from sys.modules['warnings']._onceregistry,

/// falling back to vm.state.warnings.once_registry.

fn get_once_registry(vm: &VirtualMachine) -> PyResult<PyObjectRef> {

if let Some(registry) = get_warnings_attr(vm, identifier!(&vm.ctx, onceregistry), false)? {

if !registry.class().is(vm.ctx.types.dict_type) {

return Err(vm.new_type_error(format!(

"_warnings.onceregistry must be a dict, not '{}'",

registry.class().name()

)));

}

return Ok(registry);

}

Ok(vm.state.warnings.once_registry.clone().into())

}

fn already_warned(

registry: &PyObject,

key: PyObjectRef,

should_set: bool,

vm: &VirtualMachine,

) -> PyResult<bool> {

if vm.is_none(registry) {

return Ok(false);

}

let current_version = vm.state.warnings.filters_version.load(Ordering::SeqCst);

let version_obj = registry.get_item(identifier!(&vm.ctx, version), vm).ok();

let version_matches = version_obj.as_ref().is_some_and(|v| {

v.try_int(vm)

.map(|i| i.as_u32_mask() as usize == current_version)

.unwrap_or(false)

});

if version_matches {

if let Ok(val) = registry.get_item(key.as_ref(), vm)

&& val.is_true(vm)?

{

return Ok(true);

}

} else if let Ok(dict) = PyDictRef::try_from_object(vm, registry.to_owned()) {

dict.clear();

dict.set_item(

identifier!(&vm.ctx, version),

vm.ctx.new_int(current_version).into(),

vm,

)?;

}

if should_set {

registry.set_item(key.as_ref(), vm.ctx.true_value.clone().into(), vm)?;

}

Ok(false)

}

/// Create a `(text, category)` or `(text, category, 0)` key and record

/// it in the registry via `already_warned`.

fn update_registry(

registry: &PyObject,

text: &PyObject,

category: &PyObject,

add_zero: bool,

vm: &VirtualMachine,

) -> PyResult<bool> {

let altkey: PyObjectRef = if add_zero {

PyTuple::new_ref(

vec![

text.to_owned(),

category.to_owned(),

vm.ctx.new_int(0).into(),

],

&vm.ctx,

)

.into()

} else {

PyTuple::new_ref(vec![text.to_owned(), category.to_owned()], &vm.ctx).into()

};

already_warned(registry, altkey, true, vm)

}

fn normalize_module(filename: &Py<PyStr>, vm: &VirtualMachine) -> PyObjectRef {

match filename.byte_len() {

0 => vm.new_pyobj("<unknown>"),

len if len >= 3 && filename.as_bytes().ends_with(b".py") => {

vm.new_pyobj(&filename.as_wtf8()[..len - 3])

}

_ => filename.as_object().to_owned(),

}

}

/// Search the global filters list for a matching action.

// TODO: split into filter_search() + get_filter() and support

// context-aware filters (get_warnings_context_filters).

fn get_filter(

category: PyObjectRef,

text: PyObjectRef,

lineno: usize,

module: PyObjectRef,

vm: &VirtualMachine,

) -> PyResult {

let filters = get_warnings_filters(vm)?;

// filters could change while we are iterating over it.

// Re-check list length each iteration (matches C behavior).

let mut i = 0;

while i < filters.borrow_vec().len() {

let Some(tmp_item) = filters.borrow_vec().get(i).cloned() else {

break;

};

let tmp_item = PyTupleRef::try_from_object(vm, tmp_item)

.ok()

.filter(|t| t.len() == 5)

.ok_or_else(|| {

vm.new_value_error(format!("_warnings.filters item {i} isn't a 5-tuple"))

})?;

/* action, msg, cat, mod, ln = item */

let action = &tmp_item[0];

let good_msg = check_matched(&tmp_item[1], &text, vm)?;

let is_subclass = category.is_subclass(&tmp_item[2], vm)?;

let good_mod = check_matched(&tmp_item[3], &module, vm)?;

let ln: usize = tmp_item[4].try_int(vm).map_or(0, |v| v.as_u32_mask() as _);

if good_msg && is_subclass && good_mod && (ln == 0 || lineno == ln) {

return Ok(action.to_owned());

}

i += 1;

}

get_default_action(vm)

}

pub fn warn(

message: PyObjectRef,

category: Option<PyTypeRef>,

stack_level: isize,

source: Option<PyObjectRef>,

vm: &VirtualMachine,

) -> PyResult<()> {

warn_with_skip(message, category, stack_level, source, None, vm)

}

/// do_warn: resolve context via setup_context, then call warn_explicit.

pub fn warn_with_skip(

message: PyObjectRef,

category: Option<PyTypeRef>,

mut stack_level: isize,

source: Option<PyObjectRef>,

skip_file_prefixes: Option<PyTupleRef>,

vm: &VirtualMachine,

) -> PyResult<()> {

if let Some(ref prefixes) = skip_file_prefixes

&& !prefixes.is_empty()

&& stack_level < 2

{

stack_level = 2;

}

let (filename, lineno, module, registry) =

setup_context(stack_level, skip_file_prefixes.as_ref(), vm)?;

warn_explicit(

category, message, filename, lineno, module, registry, None, source, vm,

)

}

/// Core warning logic matching `warn_explicit()` in `_warnings.c`.

#[allow(clippy::too_many_arguments)]

pub(crate) fn warn_explicit(

category: Option<PyTypeRef>,

message: PyObjectRef,

filename: PyStrRef,

lineno: usize,

module: Option<PyObjectRef>,

registry: PyObjectRef,

source_line: Option<PyObjectRef>,

source: Option<PyObjectRef>,

vm: &VirtualMachine,

) -> PyResult<()> {

// Normalize module. None → silent return (late-shutdown safety).

let module = module.unwrap_or_else(|| normalize_module(&filename, vm));

if vm.is_none(&module) {

return Ok(());

}

// Normalize message.

let is_warning = message.fast_isinstance(vm.ctx.exceptions.warning);

let (text, category, message) = if is_warning {

let text = message.str(vm)?;

let cat = message.class().to_owned();

(text, cat, message)

} else {

// For non-Warning messages, convert to string via str()

let text = message.str(vm)?;

let cat = category.unwrap_or_else(|| vm.ctx.exceptions.user_warning.to_owned());

let instance = cat.as_object().call((text.clone(),), vm)?;

(text, cat, instance)

};

let lineno_obj: PyObjectRef = vm.ctx.new_int(lineno).into();

// key = (text, category, lineno)

let key: PyObjectRef = PyTuple::new_ref(

vec![

text.clone().into(),

category.as_object().to_owned(),

lineno_obj.clone(),

],

&vm.ctx,

)

.into();

// Check if already warned

if !vm.is_none(&registry) && already_warned(&registry, key.clone(), false, vm)? {

return Ok(());

}

// Get filter action

let action = get_filter(

category.as_object().to_owned(),

text.clone().into(),

lineno,

module,

vm,

)?;

let action_str = PyStrRef::try_from_object(vm, action)

.map_err(|_| vm.new_type_error("action must be a string"))?;

if action_str.as_bytes() == b"error" {

let exc = PyBaseExceptionRef::try_from_object(vm, message)?;

return Err(exc);

}

if action_str.as_bytes() == b"ignore" {

return Ok(());

}

// For everything except "always"/"all", record in registry then

// check per-action registries.

let already = if action_str.as_wtf8() != "always" && action_str.as_wtf8() != "all" {

if !vm.is_none(&registry) {

registry.set_item(&*key, vm.ctx.true_value.clone().into(), vm)?;

}

let action_s = action_str.to_str();

match action_s {

Some("once") => {

let reg = if vm.is_none(&registry) {

get_once_registry(vm)?

} else {

registry.clone()

};

update_registry(&reg, text.as_ref(), category.as_object(), false, vm)?

}

Some("module") => {

if !vm.is_none(&registry) {

update_registry(&registry, text.as_ref(), category.as_object(), false, vm)?

} else {

false

}

}

Some("default") => false,

_ => {

return Err(vm.new_runtime_error(format!(

"Unrecognized action ({action_str}) in warnings.filters:\n {action_str}"

)));

}

}

} else {

false

};

if already {

return Ok(());

}

call_show_warning(

category,

text,

message,

filename,

lineno,

lineno_obj,

source_line,

source,

vm,

)

}

#[allow(clippy::too_many_arguments)]

fn call_show_warning(

category: PyTypeRef,

text: PyStrRef,

message: PyObjectRef,

filename: PyStrRef,

lineno: usize,

lineno_obj: PyObjectRef,

source_line: Option<PyObjectRef>,

source: Option<PyObjectRef>,

vm: &VirtualMachine,

) -> PyResult<()> {

let Some(show_fn) =

get_warnings_attr(vm, identifier!(&vm.ctx, _showwarnmsg), source.is_some())?

else {

return show_warning(filename, lineno, text, category, source_line, vm);

};

if !show_fn.is_callable() {

return Err(vm.new_type_error("warnings._showwarnmsg() must be set to a callable"));

}

let Some(warnmsg_cls) = get_warnings_attr(vm, identifier!(&vm.ctx, WarningMessage), false)?

else {

return Err(vm.new_runtime_error("unable to get warnings.WarningMessage"));

};

let msg = warnmsg_cls.call(

vec![

message,

category.into(),

filename.into(),

lineno_obj,

vm.ctx.none(),

vm.ctx.none(),

vm.unwrap_or_none(source),

],

vm,

)?;

show_fn.call((msg,), vm)?;

Ok(())

}

fn show_warning(

filename: PyStrRef,

lineno: usize,

text: PyStrRef,

category: PyTypeRef,

_source_line: Option<PyObjectRef>,

vm: &VirtualMachine,

) -> PyResult<()> {

let stderr = crate::stdlib::sys::PyStderr(vm);

writeln!(

stderr,

"{}:{}: {}: {}",

filename,

lineno,

category.name(),

text

);

Ok(())

}

/// Check if a frame's filename starts with any of the given prefixes.

fn is_filename_to_skip(frame: &crate::frame::Frame, prefixes: &PyTupleRef) -> bool {

let filename = frame.f_code().co_filename();

let filename_bytes = filename.as_bytes();

prefixes.iter().any(|prefix| {

prefix

.downcast_ref::<PyStr>()

.is_some_and(|s| filename_bytes.starts_with(s.as_bytes()))

})

}

/// Like Frame::next_external_frame but also skips frames matching prefixes.

fn next_external_frame_with_skip(

frame: &crate::frame::FrameRef,

skip_file_prefixes: Option<&PyTupleRef>,

vm: &VirtualMachine,

) -> Option<crate::frame::FrameRef> {

let mut f = frame.f_back(vm);

loop {

let current: crate::frame::FrameRef = f.take()?;

if current.is_internal_frame()

|| skip_file_prefixes.is_some_and(|p| is_filename_to_skip(&current, p))

{

f = current.f_back(vm);

} else {

return Some(current);

}

}

}

/// filename, module, and registry are new refs, globals is borrowed

/// Returns `Ok` on success, or `Err` on error (no new refs)

fn setup_context(

mut stack_level: isize,

skip_file_prefixes: Option<&PyTupleRef>,

vm: &VirtualMachine,

) -> PyResult<(PyStrRef, usize, Option<PyObjectRef>, PyObjectRef)> {

let mut f = vm.current_frame();

// Stack level comparisons to Python code is off by one as there is no

// warnings-related stack level to avoid.

if stack_level <= 0 || f.as_ref().is_some_and(|frame| frame.is_internal_frame()) {

while {

stack_level -= 1;

stack_level > 0

} {

match f {

Some(tmp) => f = tmp.f_back(vm),

None => break,

}

}

} else {

while {

stack_level -= 1;

stack_level > 0

} {

match f {

Some(tmp) => f = next_external_frame_with_skip(&tmp, skip_file_prefixes, vm),

None => break,

}

}

}

let (globals, filename, lineno) = if let Some(f) = f {

(f.globals.clone(), f.code.source_path(), f.f_lineno())

} else if let Some(frame) = vm.current_frame() {

// We have a frame but it wasn't found during stack walking

(frame.globals.clone(), vm.ctx.intern_str("<sys>"), 1)

} else {

// No frames on the stack - use sys.__dict__ (interp->sysdict)

let globals = vm

.sys_module

.as_object()

.get_attr(identifier!(vm, __dict__), vm)

.and_then(|d| {

d.downcast::<crate::builtins::PyDict>()

.map_err(|_| vm.new_type_error("sys.__dict__ is not a dictionary"))

})?;

(globals, vm.ctx.intern_str("<sys>"), 0)

};

let registry = match globals.get_item("__warningregistry__", vm) {

Ok(r) => r,

Err(_) => {

let r = vm.ctx.new_dict();

globals.set_item("__warningregistry__", r.clone().into(), vm)?;

r.into()

}

};

// Setup module.

let module = globals

.get_item("__name__", vm)

.unwrap_or_else(|_| vm.new_pyobj("<string>"));

Ok((filename.to_owned(), lineno, Some(module), registry))

}