Update warning warn_explicit

Author: mumallaeng

vm/src/vm/mod.rs

@@ -31,7 +31,9 @@ use crate::{
     import,
     protocol::PyIterIter,
     scope::Scope,
-    signal, stdlib, AsObject, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
+    signal, stdlib,
+    warn::WarningsState,
+    AsObject, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
 };
 use crossbeam_utils::atomic::AtomicCell;
 use std::sync::atomic::AtomicBool;
@@ -88,6 +90,7 @@ pub struct PyGlobalState {
     pub atexit_funcs: PyMutex<Vec<(PyObjectRef, FuncArgs)>>,
     pub codec_registry: CodecsRegistry,
     pub finalizing: AtomicBool,
+    pub warnings: WarningsState,
 }
 
 pub fn process_hash_secret_seed() -> u32 {
@@ -136,6 +139,8 @@ impl VirtualMachine {
 
         let codec_registry = CodecsRegistry::new(&ctx);
 
+        let warnings = WarningsState::init_state(&ctx);
+
         let mut vm = VirtualMachine {
             builtins,
             sys_module,
@@ -161,6 +166,7 @@ impl VirtualMachine {
                 atexit_funcs: PyMutex::default(),
                 codec_registry,
                 finalizing: AtomicBool::new(false),
+                warnings,
             }),
             initialized: false,
             recursion_depth: Cell::new(0),

vm/src/warn.rs

@@ -1,8 +1,39 @@
+use itertools::Itertools;
+
 use crate::{
-    builtins::{PyDict, PyStrRef, PyType, PyTypeRef},
-    AsObject, Py, PyObjectRef, PyResult, VirtualMachine,
+    builtins::{PyDict, PyDictRef, PyListRef, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef},
+    convert::TryFromObject,
+    AsObject, Context, Py, PyObjectRef, PyResult, VirtualMachine,
 };
 
+pub struct WarningsState {
+    filters: PyListRef,
+    _once_registry: PyDictRef,
+    _default_action: PyStrRef,
+    filters_version: usize,
+}
+
+impl WarningsState {
+    fn create_filter(ctx: &Context) -> PyListRef {
+        ctx.new_list(vec![
+            ctx.new_str("__main__").into(),
+            ctx.types.none_type.as_object().to_owned(),
+            ctx.exceptions.warning.as_object().to_owned(),
+            ctx.new_str("ACTION").into(),
+            ctx.new_int(0).into(),
+        ])
+    }
+
+    pub fn init_state(ctx: &Context) -> WarningsState {
+        WarningsState {
+            filters: Self::create_filter(ctx),
+            _once_registry: PyDict::new_ref(ctx),
+            _default_action: ctx.new_str("default"),
+            filters_version: 0,
+        }
+    }
+}
+
 pub fn py_warn(
     category: &Py<PyType>,
     message: String,
@@ -31,19 +62,130 @@ pub fn warn(
     )
 }
 
+fn get_filter(
+    _category: PyObjectRef,
+    _text: PyObjectRef,
+    _lineno: usize,
+    _module: PyObjectRef,
+    mut _item: PyObjectRef,
+    vm: &VirtualMachine,
+) -> PyResult<()> {
+    let filters = vm.state.warnings.filters.as_object().to_owned();
+
+    let filters: PyListRef = filters
+        .try_into_value(vm)
+        .map_err(|_| vm.new_value_error("_warnings.filters must be a list".to_string()))?;
+
+    /* WarningsState.filters could change while we are iterating over it. */
+    for i in 0..filters.borrow_vec().len() {
+        let tmp_item = filters.borrow_vec().get(i).cloned();
+        let tmp_item = if let Some(tmp_item) = tmp_item {
+            let tmp_item = PyTupleRef::try_from_object(vm, tmp_item)?;
+            if tmp_item.len() != 5 {
+                Err(vm.new_value_error(format!("_warnings.filters item {} isn't a 5-tuple", i)))
+            } else {
+                Ok(tmp_item)
+            }
+        } else {
+            Err(vm.new_value_error(format!("_warnings.filters item {} isn't a 5-tuple", i)))
+        }?;
+
+        /* Python code: action, msg, cat, mod, ln = item */
+        let _action = tmp_item.get(0);
+        let _msg = tmp_item.get(1);
+        let _cat = tmp_item.get(2);
+        let _item_mod = tmp_item.get(3);
+        let _ln_obj = tmp_item.get(4);
+    }
+
+    Ok(())
+}
+
+fn already_warned(
+    registry: PyObjectRef,
+    key: PyObjectRef,
+    should_set: usize,
+    vm: &VirtualMachine,
+) -> bool {
+    let version_obj = registry.get_item("version", vm).ok();
+    let filters_version = vm
+        .ctx
+        .new_int(vm.state.warnings.filters_version)
+        .as_object()
+        .to_owned();
+
+    if version_obj.is_none()
+        || version_obj.as_ref().unwrap().try_int(vm).is_err()
+        || !version_obj.unwrap().is(&filters_version)
+    {
+        let registry = registry.dict();
+        registry.as_ref().map(|registry| {
+            registry.into_iter().collect_vec().clear();
+            registry
+        });
+
+        if let Some(registry) = registry {
+            if registry.set_item("version", filters_version, vm).is_err() {
+                return false;
+            }
+        }
+    } else {
+        let already_warned = registry.get_item(key.as_ref(), vm);
+        return already_warned.is_ok();
+    }
+
+    /* This warning wasn't found in the registry, set it. */
+    if should_set != 0 {
+        registry.set_item(key.as_ref(), vm.new_pyobj(""), vm).ok();
+    }
+
+    false
+}
+
+fn normalize_module(filename: PyStrRef, vm: &VirtualMachine) -> Option<PyObjectRef> {
+    let len = filename.char_len();
+
+    if len == 0 {
+        Some(vm.new_pyobj("<unknown>"))
+    } else if len >= 3 && filename.as_str().contains(".py") {
+        Some(vm.new_pyobj(filename.as_str().replace(".py", "")))
+    } else {
+        Some(filename.as_object().to_owned())
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 fn warn_explicit(
     category: Option<PyTypeRef>,
     message: PyStrRef,
-    _filename: PyStrRef,
-    _lineno: usize,
-    _module: PyObjectRef,
-    _registry: PyObjectRef,
+    filename: PyStrRef,
+    lineno: usize,
+    module: Option<PyObjectRef>,
+    registry: PyObjectRef,
     _source_line: Option<PyObjectRef>,
     _source: Option<PyObjectRef>,
     vm: &VirtualMachine,
 ) -> PyResult<()> {
-    // TODO: Implement correctly
+    if registry.clone().is_true(vm)?
+        && !registry.class().is(vm.ctx.types.dict_type)
+        && !registry.class().is(vm.ctx.types.none_type)
+    {
+        return Err(vm.new_type_error("'registry' must be a dict or None".to_string()));
+    }
+
+    // Normalize module.
+    let module = module.and(normalize_module(filename, vm));
+    let module = if let Some(module) = module {
+        module
+    } else {
+        return Ok(());
+    };
+
+    // Normalize message.
+    let text = message.as_str();
+    if text.is_empty() {
+        return Ok(());
+    }
     let category = if let Some(category) = category {
         if !category.fast_issubclass(vm.ctx.exceptions.warning) {
             return Err(vm.new_type_error(format!(
@@ -55,8 +197,35 @@ fn warn_explicit(
     } else {
         vm.ctx.exceptions.user_warning.to_owned()
     };
+
+    // Create key.
+    let key = PyTuple::new_ref(
+        vec![
+            vm.ctx.new_int(3).into(),
+            vm.ctx.new_str(text).into(),
+            category.as_object().to_owned(),
+            vm.ctx.new_int(lineno).into(),
+        ],
+        &vm.ctx,
+    );
+
+    if already_warned(registry, key.as_object().to_owned(), 0, vm) {
+        return Ok(());
+    }
+    // Else this warning hasn't been generated before.
+
+    let item = vm.ctx.new_tuple(vec![]).into();
+    let _action = get_filter(
+        category.as_object().to_owned(),
+        vm.ctx.new_str(text).into(),
+        lineno,
+        module,
+        item,
+        vm,
+    );
+
     let stderr = crate::stdlib::sys::PyStderr(vm);
-    writeln!(stderr, "{}: {}", category.name(), message.as_str(),);
+    writeln!(stderr, "{}: {}", category.name(), text,);
     Ok(())
 }
 
@@ -67,7 +236,7 @@ fn setup_context(
     vm: &VirtualMachine,
 ) -> PyResult<
     // filename, lineno, module, registry
-    (PyStrRef, usize, PyObjectRef, PyObjectRef),
+    (PyStrRef, usize, Option<PyObjectRef>, PyObjectRef),
 > {
     let __warningregistry__ = "__warningregistry__";
     let __name__ = "__name__";
@@ -120,5 +289,5 @@ fn setup_context(
     let module = globals
         .get_item(__name__, vm)
         .unwrap_or_else(|_| vm.new_pyobj("<string>"));
-    Ok((filename.to_owned(), lineno, module, registry))
+    Ok((filename.to_owned(), lineno, Some(module), registry))
 }