search_path: Vec<PyObjectRef>,

search_cache: HashMap<String, PyCodec>,

errors: HashMap<String, PyObjectRef>,

#[inline]

pub fn from_tuple(tuple: PyTupleRef) -> Result<Self, PyTupleRef> {

if tuple.len() == 4 {

Ok(Self(tuple))

} else {

Err(tuple)

}

}

#[inline]

pub fn into_tuple(self) -> PyTupleRef {

self.0

}

#[inline]

pub fn as_tuple(&self) -> &Py<PyTuple> {

&self.0

}

#[inline]

pub fn get_encode_func(&self) -> &PyObject {

&self.0[0]

}

#[inline]

pub fn get_decode_func(&self) -> &PyObject {

&self.0[1]

}

pub fn is_text_codec(&self, vm: &VirtualMachine) -> PyResult<bool> {

let is_text = vm.get_attribute_opt(self.0.clone().into(), "_is_text_encoding")?;

is_text.map_or(Ok(true), |is_text| is_text.try_to_bool(vm))

}

pub fn encode(

&self,

obj: PyObjectRef,

errors: Option<PyUtf8StrRef>,

vm: &VirtualMachine,

) -> PyResult {

let args = match errors {

Some(errors) => vec![obj, errors.into_wtf8().into()],

None => vec![obj],

};

let res = self.get_encode_func().call(args, vm)?;

let res = res

.downcast::<PyTuple>()

.ok()

.filter(|tuple| tuple.len() == 2)

.ok_or_else(|| vm.new_type_error("encoder must return a tuple (object, integer)"))?;

// we don't actually care about the integer

Ok(res[0].clone())

}

pub fn decode(

&self,

obj: PyObjectRef,

errors: Option<PyUtf8StrRef>,

vm: &VirtualMachine,

) -> PyResult {

let args = match errors {

Some(errors) => vec![obj, errors.into_wtf8().into()],

None => vec![obj],

};

let res = self.get_decode_func().call(args, vm)?;

let res = res

.downcast::<PyTuple>()

.ok()

.filter(|tuple| tuple.len() == 2)

.ok_or_else(|| vm.new_type_error("decoder must return a tuple (object,integer)"))?;

// we don't actually care about the integer

Ok(res[0].clone())

}

pub fn get_incremental_encoder(

&self,

errors: Option<PyStrRef>,

vm: &VirtualMachine,

) -> PyResult {

let args = errors.map_or_else(Vec::new, |e| vec![e.into()]);

vm.call_method(self.0.as_object(), "incrementalencoder", args)

}

pub fn get_incremental_decoder(

&self,

errors: Option<PyStrRef>,

vm: &VirtualMachine,

) -> PyResult {

let args = errors.map_or_else(Vec::new, |e| vec![e.into()]);

vm.call_method(self.0.as_object(), "incrementaldecoder", args)

}

fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {

obj.downcast::<PyTuple>()

.ok()

.and_then(|tuple| Self::from_tuple(tuple).ok())

.ok_or_else(|| vm.new_type_error("codec search functions must return 4-tuples"))

}

#[inline]

fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {

self.0.into()

}

/// Reset the inner RwLock to unlocked state after fork().

///

/// # Safety

/// Must only be called after fork() in the child process when no other

/// threads exist.

#[cfg(all(unix, feature = "threading", feature = "host_env"))]

pub(crate) unsafe fn reinit_after_fork(&self) {

unsafe { crate::common::lock::reinit_rwlock_after_fork(&self.inner) };

}

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

::rustpython_vm::common::static_cell! {

static METHODS: Box<[PyMethodDef]>;

}

let methods = METHODS.get_or_init(|| {

crate::define_methods![

"strict_errors" => strict_errors as EMPTY,

"ignore_errors" => ignore_errors as EMPTY,

"replace_errors" => replace_errors as EMPTY,

"xmlcharrefreplace_errors" => xmlcharrefreplace_errors as EMPTY,

"backslashreplace_errors" => backslashreplace_errors as EMPTY,

"namereplace_errors" => namereplace_errors as EMPTY,

"surrogatepass_errors" => surrogatepass_errors as EMPTY,

"surrogateescape_errors" => surrogateescape_errors as EMPTY

]

.into_boxed_slice()

});

let errors = [

("strict", methods[0].build_function(ctx)),

("ignore", methods[1].build_function(ctx)),

("replace", methods[2].build_function(ctx)),

("xmlcharrefreplace", methods[3].build_function(ctx)),

("backslashreplace", methods[4].build_function(ctx)),

("namereplace", methods[5].build_function(ctx)),

("surrogatepass", methods[6].build_function(ctx)),

("surrogateescape", methods[7].build_function(ctx)),

]

.into_iter()

.map(|(name, f)| (name.to_owned(), f.into()))

.collect();

let inner = RegistryInner {

search_path: Vec::new(),

search_cache: HashMap::new(),

errors,

};

Self {

inner: PyRwLock::new(inner),

}

}

pub fn register(&self, search_function: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {

if !search_function.is_callable() {

return Err(vm.new_type_error("argument must be callable"));

}

self.inner.write().search_path.push(search_function);

Ok(())

}

pub fn unregister(&self, search_function: PyObjectRef) {

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;

}

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;

}

}

}

pub(crate) fn register_manual(&self, name: &str, codec: PyCodec) {

let name = normalize_encoding_name(name);

self.inner

.write()

.search_cache

.insert(name.into_owned(), codec);

}

pub fn lookup(&self, encoding: &str, vm: &VirtualMachine) -> PyResult<PyCodec> {

let encoding = normalize_encoding_name(encoding);

let search_path = {

let inner = self.inner.read();

if let Some(codec) = inner.search_cache.get(encoding.as_ref()) {

// hit cache

return Ok(codec.clone());

}

inner.search_path.clone()

};

let encoding: PyUtf8StrRef = vm.ctx.new_utf8_str(encoding.as_ref());

for func in search_path {

let res = func.call((encoding.clone(),), vm)?;

let res: Option<PyCodec> = res.try_into_value(vm)?;

if let Some(codec) = res {

let mut inner = self.inner.write();

// someone might have raced us to this, so use theirs

let codec = inner

.search_cache

.entry(encoding.as_str().to_owned())

.or_insert(codec);

return Ok(codec.clone());

}

}

Err(vm.new_lookup_error(format!("unknown encoding: {encoding}")))

}

fn _lookup_text_encoding(

&self,

encoding: &str,

generic_func: &str,

vm: &VirtualMachine,

) -> PyResult<PyCodec> {

let codec = self.lookup(encoding, vm)?;

if codec.is_text_codec(vm)? {

Ok(codec)

} else {

Err(vm.new_lookup_error(format!(

"'{encoding}' is not a text encoding; use {generic_func} to handle arbitrary codecs"

)))

}

}

pub fn forget(&self, encoding: &str) -> Option<PyCodec> {

let encoding = normalize_encoding_name(encoding);

self.inner.write().search_cache.remove(encoding.as_ref())

}

pub fn encode(

&self,

obj: PyObjectRef,

encoding: &str,

errors: Option<PyUtf8StrRef>,

vm: &VirtualMachine,

) -> PyResult {

let codec = self.lookup(encoding, vm)?;

codec.encode(obj, errors, vm).inspect_err(|exc| {

Self::add_codec_note(exc, "encoding", encoding, vm);

})

}

pub fn decode(

&self,

obj: PyObjectRef,

encoding: &str,

errors: Option<PyUtf8StrRef>,

vm: &VirtualMachine,

) -> PyResult {

let codec = self.lookup(encoding, vm)?;

codec.decode(obj, errors, vm).inspect_err(|exc| {

Self::add_codec_note(exc, "decoding", encoding, vm);

})

}

pub fn encode_text(

&self,

obj: PyStrRef,

encoding: &str,

errors: Option<PyUtf8StrRef>,

vm: &VirtualMachine,

) -> PyResult<PyBytesRef> {

let codec = self._lookup_text_encoding(encoding, "codecs.encode()", vm)?;

codec

.encode(obj.into(), errors, vm)

.inspect_err(|exc| {

Self::add_codec_note(exc, "encoding", encoding, vm);

})?

.downcast()

.map_err(|obj| {

vm.new_type_error(format!(

"'{}' encoder returned '{}' instead of 'bytes'; use codecs.encode() to \

encoding,

obj.class().name(),

))

})

}

/// The text an encoded object holds. An object holding nothing decodes to

/// the empty string without the encoding ever being looked up.

/// = PyUnicode_FromEncodedObject

pub fn decode_text_object(

&self,

obj: PyObjectRef,

encoding: &str,

errors: Option<PyUtf8StrRef>,

vm: &VirtualMachine,

) -> PyResult<PyStrRef> {

let empty = if let Some(bytes) = obj.downcast_ref::<PyBytes>() {

bytes.as_bytes().is_empty()

} else if let Some(bytes) = obj.downcast_ref::<PyByteArray>() {

bytes.borrow_buf().is_empty()

} else {

false

};

if empty {

return Ok(vm.ctx.empty_str.to_owned());

}

self.decode_text(obj, encoding, errors, vm)

}

pub fn decode_text(

&self,

obj: PyObjectRef,

encoding: &str,

errors: Option<PyUtf8StrRef>,

vm: &VirtualMachine,

) -> PyResult<PyStrRef> {

let codec = self._lookup_text_encoding(encoding, "codecs.decode()", vm)?;

codec

.decode(obj, errors, vm)

.inspect_err(|exc| {

Self::add_codec_note(exc, "decoding", encoding, vm);

})?

.downcast()

.map_err(|obj| {

vm.new_type_error(format!(

"'{}' decoder returned '{}' instead of 'str'; use codecs.decode() to \

encoding,

obj.class().name(),

))

})

}

fn add_codec_note(

exc: &crate::builtins::PyBaseExceptionRef,

operation: &str,

encoding: &str,

vm: &VirtualMachine,

) {

let note = format!("{operation} with '{encoding}' codec failed");

let _ = vm.call_method(exc.as_object(), "add_note", (vm.ctx.new_str(note),));

}

pub fn register_error(&self, name: String, handler: PyObjectRef) -> Option<PyObjectRef> {

self.inner.write().errors.insert(name, handler)

}

pub fn unregister_error(&self, name: &str, vm: &VirtualMachine) -> PyResult<bool> {

const BUILTIN_ERROR_HANDLERS: &[&str] = &[

"strict",

"ignore",

"replace",

"xmlcharrefreplace",

"backslashreplace",

"namereplace",

"surrogatepass",

"surrogateescape",

];

if BUILTIN_ERROR_HANDLERS.contains(&name) {

return Err(vm.new_value_error(format!(

"cannot un-register built-in error handler '{name}'"

)));

}

Ok(self.inner.write().errors.remove(name).is_some())

}

pub fn lookup_error_opt(&self, name: &str) -> Option<PyObjectRef> {

self.inner.read().errors.get(name).cloned()

}

pub fn lookup_error(&self, name: &str, vm: &VirtualMachine) -> PyResult<PyObjectRef> {

self.lookup_error_opt(name)

.ok_or_else(|| vm.new_lookup_error(format!("unknown error handler name '{name}'")))

}

// _Py_normalize_encoding: collapse non-alphanumeric/non-dot chars into

// single underscore, strip non-ASCII, lowercase ASCII letters.

let needs_transform = encoding

.bytes()

.any(|b| b.is_ascii_uppercase() || !b.is_ascii_alphanumeric() && b != b'.');

if !needs_transform {

return encoding.into();

}

let mut out = String::with_capacity(encoding.len());

let mut punct = false;

for c in encoding.chars() {

if c.is_ascii_alphanumeric() || c == '.' {

if punct && !out.is_empty() {

out.push('_');

}

out.push(c.to_ascii_lowercase());

punct = false;

} else {

punct = true;

}

}

out.into()

Utf8,

Utf16Be,

Utf16Le,

Utf32Be,

Utf32Le,

const UTF_16_NE: Self = cfg_select! {

target_endian = "little" => Self::Utf16Le,

target_endian = "big" => Self::Utf16Be,

};

const UTF_32_NE: Self = cfg_select! {

target_endian = "little" => Self::Utf32Le,

target_endian = "big" => Self::Utf32Be,

};

fn parse(encoding: &str) -> Option<Self> {

if let Some(encoding) = encoding.to_lowercase().strip_prefix("utf") {

let encoding = encoding

.strip_prefix(|c| ['-', '_'].contains(&c))

.unwrap_or(encoding);

if encoding == "8" {

Some(Self::Utf8)

} else if let Some(encoding) = encoding.strip_prefix("16") {

if encoding.is_empty() {

return Some(Self::UTF_16_NE);

}

let encoding = encoding.strip_prefix(['-', '_']).unwrap_or(encoding);

match encoding {

"be" => Some(Self::Utf16Be),

"le" => Some(Self::Utf16Le),

_ => None,

}

} else if let Some(encoding) = encoding.strip_prefix("32") {

if encoding.is_empty() {

return Some(Self::UTF_32_NE);

}

let encoding = encoding.strip_prefix(['-', '_']).unwrap_or(encoding);

match encoding {

"be" => Some(Self::Utf32Be),

"le" => Some(Self::Utf32Le),

_ => None,

}

} else {

None

}

} else if encoding == "cp65001" {

Some(Self::Utf8)

} else {

None

}

}

fn handle_encode_error(

&self,

ctx: &mut PyEncodeContext<'a>,

range: Range<StrSize>,

reason: Option<&str>,

) -> PyResult<(EncodeReplace<PyEncodeContext<'a>>, StrSize)> {

let standard_encoding = StandardEncoding::parse(ctx.encoding)

.ok_or_else(|| ctx.error_encoding(range.clone(), reason))?;

let err_str = &ctx.full_data()[range.start.bytes..range.end.bytes];

let num_chars = range.end.chars - range.start.chars;

let mut out: Vec<u8> = Vec::with_capacity(num_chars * 4);

for ch in err_str.code_points() {

let c = ch.to_u32();

if !(0xd800..=0xdfff).contains(&c) {

// Not a surrogate, fail with original exception

return Err(ctx.error_encoding(range, reason));

}

match standard_encoding {

StandardEncoding::Utf8 => out.extend(ch.encode_wtf8(&mut [0; 4]).as_bytes()),

StandardEncoding::Utf16Le => out.extend((c as u16).to_le_bytes()),

StandardEncoding::Utf16Be => out.extend((c as u16).to_be_bytes()),

StandardEncoding::Utf32Le => out.extend(c.to_le_bytes()),

StandardEncoding::Utf32Be => out.extend(c.to_be_bytes()),

}

}

Ok((EncodeReplace::Bytes(ctx.bytes(out)), range.end))

}

fn handle_decode_error(

&self,

ctx: &mut PyDecodeContext<'a>,

byte_range: Range<usize>,

reason: Option<&str>,

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

let standard_encoding = StandardEncoding::parse(ctx.encoding)

.ok_or_else(|| ctx.error_decoding(byte_range.clone(), reason))?;

let s = ctx.full_data();

debug_assert!(byte_range.start <= s.len().saturating_sub(1));

debug_assert!(byte_range.end >= 1.min(s.len()));

debug_assert!(byte_range.end <= s.len());

// Try decoding a single surrogate character. If there are more,

// let the codec call us again.

let p = &s[byte_range.start..];

fn slice<const N: usize>(p: &[u8]) -> Option<[u8; N]> {

p.first_chunk().copied()

}

let c = match standard_encoding {

StandardEncoding::Utf8 => {

// it's a three-byte code

slice::<3>(p)

.filter(|&[a, b, c]| {

(u32::from(a) & 0xf0) == 0xe0

&& (u32::from(b) & 0xc0) == 0x80

&& (u32::from(c) & 0xc0) == 0x80

})

.map(|[a, b, c]| {

((u32::from(a) & 0x0f) << 12)

+ ((u32::from(b) & 0x3f) << 6)

+ (u32::from(c) & 0x3f)

})

}

StandardEncoding::Utf16Le => slice(p).map(u16::from_le_bytes).map(u32::from),

StandardEncoding::Utf16Be => slice(p).map(u16::from_be_bytes).map(u32::from),

StandardEncoding::Utf32Le => slice(p).map(u32::from_le_bytes),

StandardEncoding::Utf32Be => slice(p).map(u32::from_be_bytes),

};

let byte_length = match standard_encoding {

StandardEncoding::Utf8 => 3,

StandardEncoding::Utf16Be | StandardEncoding::Utf16Le => 2,

StandardEncoding::Utf32Be | StandardEncoding::Utf32Le => 4,

};

// !Py_UNICODE_IS_SURROGATE

let c = c

.and_then(CodePoint::from_u32)

.filter(|c| matches!(c.to_u32(), 0xd800..=0xdfff))

.ok_or_else(|| ctx.error_decoding(byte_range.clone(), reason))?;

Ok((ctx.string(c.into()), byte_range.start + byte_length))

}

vm: &'a VirtualMachine,

encoding: &'a str,

data: &'a Py<PyStr>,

pos: StrSize,

exception: OnceCell<PyBaseExceptionRef>,

pub(crate) fn new(encoding: &'a str, data: &'a Py<PyStr>, vm: &'a VirtualMachine) -> Self {

Self {

vm,

encoding,

data,

pos: StrSize::default(),

exception: OnceCell::new(),

}

}

type Error = PyBaseExceptionRef;

type StrBuf = PyStrRef;

type BytesBuf = PyBytesRef;

fn string(&self, s: Wtf8Buf) -> Self::StrBuf {

self.vm.ctx.new_str(s)

}

fn bytes(&self, b: Vec<u8>) -> Self::BytesBuf {

self.vm.ctx.new_bytes(b)

}

fn full_data(&self) -> &Wtf8 {

self.data.as_wtf8()

}

fn data_len(&self) -> StrSize {

StrSize {

bytes: self.data.byte_len(),

chars: self.data.char_len(),

}

}

fn remaining_data(&self) -> &Wtf8 {

&self.full_data()[self.pos.bytes..]

}

fn position(&self) -> StrSize {

self.pos

}

fn restart_from(&mut self, pos: StrSize) -> Result<(), Self::Error> {

if pos.chars > self.data.char_len() {

return Err(self.vm.new_index_error(format!(

"position {} from error handler out of bounds",

pos.chars

)));

}

assert!(

self.data.as_wtf8().is_code_point_boundary(pos.bytes),

"invalid pos {pos:?} for {:?}",

self.data.as_wtf8()

);

self.pos = pos;

Ok(())

}

fn error_encoding(&self, range: Range<StrSize>, reason: Option<&str>) -> Self::Error {

let vm = self.vm;

match self.exception.get() {

Some(exc) => {

match update_unicode_error_attrs(

exc.as_object(),

range.start.chars,

range.end.chars,

reason,

vm,

) {

Ok(()) => exc.clone(),

Err(e) => e,

}

}

None => self

.exception

.get_or_init(|| {

let reason = reason.expect(

"should only ever pass reason: None if an exception is already set",

);

vm.new_unicode_encode_error_real(

vm.ctx.new_str(self.encoding),

self.data.to_owned(),

range.start.chars,

range.end.chars,

vm.ctx.new_str(reason),

)

})

.clone(),

}

}

vm: &'a VirtualMachine,

encoding: &'a str,

data: PyDecodeData<'a>,

orig_bytes: Option<&'a Py<PyBytes>>,

pos: usize,

exception: OnceCell<PyBaseExceptionRef>,

type Target = [u8];

fn deref(&self) -> &Self::Target {

match self {

PyDecodeData::Original(data) => data,

PyDecodeData::Modified(data) => data,

}

}

pub(crate) fn new(encoding: &'a str, data: &'a ArgBytesLike, vm: &'a VirtualMachine) -> Self {

Self {

vm,

encoding,

data: PyDecodeData::Original(data.borrow_buf()),

orig_bytes: data.as_object().downcast_ref(),

pos: 0,

exception: OnceCell::new(),

}

}

type Error = PyBaseExceptionRef;

type StrBuf = PyStrRef;

type BytesBuf = PyBytesRef;

fn string(&self, s: Wtf8Buf) -> Self::StrBuf {

self.vm.ctx.new_str(s)

}

fn bytes(&self, b: Vec<u8>) -> Self::BytesBuf {

self.vm.ctx.new_bytes(b)

}

fn full_data(&self) -> &[u8] {

&self.data

}

fn remaining_data(&self) -> &[u8] {

&self.data[self.pos..]

}

fn position(&self) -> usize {

self.pos

}