diff --git a/vm/src/codecs.rs b/vm/src/codecs.rs index 11fb587a2..9609e305e 100644 --- a/vm/src/codecs.rs +++ b/vm/src/codecs.rs @@ -2,7 +2,7 @@ use crate::{ builtins::{PyBaseExceptionRef, PyBytesRef, PyStr, PyStrRef, PyTuple, PyTupleRef}, common::{ascii, lock::PyRwLock}, function::IntoPyObject, - PyContext, PyObject, PyObjectRef, PyResult, PyValue, TryFromObject, TypeProtocol, + IdProtocol, PyContext, PyObject, PyObjectRef, PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine, }; use std::borrow::Cow; @@ -195,6 +195,24 @@ impl CodecsRegistry { Ok(()) } + pub fn unregister(&self, search_function: PyObjectRef) -> PyResult<()> { + let mut inner = self.inner.write(); + // Do nothing if search_path is not created yet or was cleared. + if inner.search_path.is_empty() { + return Ok(()); + } + for (i, item) in inner.search_path.iter().enumerate() { + if item.get_id() == search_function.get_id() { + if !inner.search_cache.is_empty() { + inner.search_cache.clear(); + } + inner.search_path.remove(i); + return Ok(()); + } + } + Ok(()) + } + pub fn lookup(&self, encoding: &str, vm: &VirtualMachine) -> PyResult { let encoding = normalize_encoding_name(encoding); let inner = self.inner.read(); diff --git a/vm/src/stdlib/codecs.rs b/vm/src/stdlib/codecs.rs index bd422242a..15ab8e60e 100644 --- a/vm/src/stdlib/codecs.rs +++ b/vm/src/stdlib/codecs.rs @@ -16,6 +16,11 @@ mod _codecs { vm.state.codec_registry.register(search_function, vm) } + #[pyfunction] + fn unregister(search_function: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + vm.state.codec_registry.unregister(search_function) + } + #[pyfunction] fn lookup(encoding: PyStrRef, vm: &VirtualMachine) -> PyResult { vm.state