use crate::obj::objint::PyIntRef;
use crate::obj::objslice::PySliceRef;
use crate::obj::objstr::PyStringRef;
use crate::obj::objtuple::PyTupleRef;
use crate::pyhash;
use crate::pyobject::Either;
use crate::vm::VirtualMachine;
use core::cell::Cell;
use std::ops::Deref;
use crate::function::OptionalArg;
use crate::pyobject::{
IntoPyObject, PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue,
TryFromObject,
};
use super::objbyteinner::{
ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
ByteInnerPosition, ByteInnerSplitOptions, ByteInnerSplitlinesOptions,
ByteInnerTranslateOptions, PyByteInner,
};
use super::objiter;
use super::objtype::PyClassRef;
use wtf8;
#[pyclass(name = "bytes")]
#[derive(Clone, Debug)]
pub struct PyBytes {
inner: PyByteInner,
}
pub type PyBytesRef = PyRef<PyBytes>;
impl PyBytes {
pub fn new(elements: Vec<u8>) -> Self {
PyBytes {
inner: PyByteInner { elements },
}
}
pub fn from_string(value: &str, encoding: &str, vm: &VirtualMachine) -> PyResult<Self> {
Ok(PyBytes {
inner: PyByteInner::from_string(value, encoding, vm)?,
})
}
pub fn get_value(&self) -> &[u8] {
&self.inner.elements
}
}
impl IntoPyObject for Vec<u8> {
fn into_pyobject(self, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self))
}
}
impl Deref for PyBytes {
type Target = [u8];
fn deref(&self) -> &[u8] {
&self.inner.elements
}
}
impl PyValue for PyBytes {
fn class(vm: &VirtualMachine) -> PyClassRef {
vm.ctx.bytes_type()
}
}
pub fn get_value<'a>(obj: &'a PyObjectRef) -> impl Deref<Target = Vec<u8>> + 'a {
&obj.payload::<PyBytes>().unwrap().inner.elements
}
pub fn init(context: &PyContext) {
PyBytesRef::extend_class(context, &context.types.bytes_type);
let bytes_type = &context.types.bytes_type;
extend_class!(context, bytes_type, {
"fromhex" => context.new_rustfunc(PyBytesRef::fromhex),
"maketrans" => context.new_rustfunc(PyByteInner::maketrans),
});
PyBytesIterator::extend_class(context, &context.types.bytesiterator_type);
}
#[pyimpl]
impl PyBytesRef {
#[pymethod(name = "__new__")]
fn bytes_new(
cls: PyClassRef,
options: ByteInnerNewOptions,
vm: &VirtualMachine,
) -> PyResult<PyBytesRef> {
PyBytes {
inner: options.get_value(vm)?,
}
.into_ref_with_type(vm, cls)
}
#[pymethod(name = "__repr__")]
fn repr(self, vm: &VirtualMachine) -> PyResult {
Ok(vm.new_str(format!("b'{}'", self.inner.repr()?)))
}
#[pymethod(name = "__len__")]
fn len(self, _vm: &VirtualMachine) -> usize {
self.inner.len()
}
#[pymethod(name = "__eq__")]
fn eq(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
self.inner.eq(other, vm)
}
#[pymethod(name = "__ge__")]
fn ge(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
self.inner.ge(other, vm)
}
#[pymethod(name = "__le__")]
fn le(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
self.inner.le(other, vm)
}
#[pymethod(name = "__gt__")]
fn gt(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
self.inner.gt(other, vm)
}
#[pymethod(name = "__lt__")]
fn lt(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
self.inner.lt(other, vm)
}
#[pymethod(name = "__hash__")]
fn hash(self, _vm: &VirtualMachine) -> pyhash::PyHash {
self.inner.hash()
}
#[pymethod(name = "__iter__")]
fn iter(self, _vm: &VirtualMachine) -> PyBytesIterator {
PyBytesIterator {
position: Cell::new(0),
bytes: self,
}
}
#[pymethod(name = "__add__")]
fn add(self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
if let Ok(other) = PyByteInner::try_from_object(vm, other) {
Ok(vm.ctx.new_bytearray(self.inner.add(other)))
} else {
Ok(vm.ctx.not_implemented())
}
}
#[pymethod(name = "__contains__")]
fn contains(self, needle: Either<PyByteInner, PyIntRef>, vm: &VirtualMachine) -> PyResult {
self.inner.contains(needle, vm)
}
#[pymethod(name = "__getitem__")]
fn getitem(self, needle: Either<PyIntRef, PySliceRef>, vm: &VirtualMachine) -> PyResult {
self.inner.getitem(needle, vm)
}
#[pymethod(name = "isalnum")]
fn isalnum(self, vm: &VirtualMachine) -> PyResult {
self.inner.isalnum(vm)
}
#[pymethod(name = "isalpha")]
fn isalpha(self, vm: &VirtualMachine) -> PyResult {
self.inner.isalpha(vm)
}
#[pymethod(name = "isascii")]
fn isascii(self, vm: &VirtualMachine) -> PyResult {
self.inner.isascii(vm)
}
#[pymethod(name = "isdigit")]
fn isdigit(self, vm: &VirtualMachine) -> PyResult {
self.inner.isdigit(vm)
}
#[pymethod(name = "islower")]
fn islower(self, vm: &VirtualMachine) -> PyResult {
self.inner.islower(vm)
}
#[pymethod(name = "isspace")]
fn isspace(self, vm: &VirtualMachine) -> PyResult {
self.inner.isspace(vm)
}
#[pymethod(name = "isupper")]
fn isupper(self, vm: &VirtualMachine) -> PyResult {
self.inner.isupper(vm)
}
#[pymethod(name = "istitle")]
fn istitle(self, vm: &VirtualMachine) -> PyResult {
self.inner.istitle(vm)
}
#[pymethod(name = "lower")]
fn lower(self, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.lower(vm)))
}
#[pymethod(name = "upper")]
fn upper(self, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.upper(vm)))
}
#[pymethod(name = "capitalize")]
fn capitalize(self, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.capitalize(vm)))
}
#[pymethod(name = "swapcase")]
fn swapcase(self, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.swapcase(vm)))
}
#[pymethod(name = "hex")]
fn hex(self, vm: &VirtualMachine) -> PyResult {
self.inner.hex(vm)
}
fn fromhex(string: PyStringRef, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(PyByteInner::fromhex(string.as_str(), vm)?))
}
#[pymethod(name = "center")]
fn center(self, options: ByteInnerPaddingOptions, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.center(options, vm)?))
}
#[pymethod(name = "ljust")]
fn ljust(self, options: ByteInnerPaddingOptions, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.ljust(options, vm)?))
}
#[pymethod(name = "rjust")]
fn rjust(self, options: ByteInnerPaddingOptions, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.rjust(options, vm)?))
}
#[pymethod(name = "count")]
fn count(self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<usize> {
self.inner.count(options, vm)
}
#[pymethod(name = "join")]
fn join(self, iter: PyIterable, vm: &VirtualMachine) -> PyResult {
self.inner.join(iter, vm)
}
#[pymethod(name = "endswith")]
fn endswith(
self,
suffix: Either<PyByteInner, PyTupleRef>,
start: OptionalArg<PyObjectRef>,
end: OptionalArg<PyObjectRef>,
vm: &VirtualMachine,
) -> PyResult {
self.inner.startsendswith(suffix, start, end, true, vm)
}
#[pymethod(name = "startswith")]
fn startswith(
self,
prefix: Either<PyByteInner, PyTupleRef>,
start: OptionalArg<PyObjectRef>,
end: OptionalArg<PyObjectRef>,
vm: &VirtualMachine,
) -> PyResult {
self.inner.startsendswith(prefix, start, end, false, vm)
}
#[pymethod(name = "find")]
fn find(self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
self.inner.find(options, false, vm)
}
#[pymethod(name = "index")]
fn index(self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
let res = self.inner.find(options, false, vm)?;
if res == -1 {
return Err(vm.new_value_error("substring not found".to_string()));
}
Ok(res)
}
#[pymethod(name = "rfind")]
fn rfind(self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
self.inner.find(options, true, vm)
}
#[pymethod(name = "rindex")]
fn rindex(self, options: ByteInnerFindOptions, vm: &VirtualMachine) -> PyResult<isize> {
let res = self.inner.find(options, true, vm)?;
if res == -1 {
return Err(vm.new_value_error("substring not found".to_string()));
}
Ok(res)
}
#[pymethod(name = "translate")]
fn translate(self, options: ByteInnerTranslateOptions, vm: &VirtualMachine) -> PyResult {
self.inner.translate(options, vm)
}
#[pymethod(name = "strip")]
fn strip(self, chars: OptionalArg<PyByteInner>, vm: &VirtualMachine) -> PyResult {
Ok(vm
.ctx
.new_bytes(self.inner.strip(chars, ByteInnerPosition::All, vm)?))
}
#[pymethod(name = "lstrip")]
fn lstrip(self, chars: OptionalArg<PyByteInner>, vm: &VirtualMachine) -> PyResult {
Ok(vm
.ctx
.new_bytes(self.inner.strip(chars, ByteInnerPosition::Left, vm)?))
}
#[pymethod(name = "rstrip")]
fn rstrip(self, chars: OptionalArg<PyByteInner>, vm: &VirtualMachine) -> PyResult {
Ok(vm
.ctx
.new_bytes(self.inner.strip(chars, ByteInnerPosition::Right, vm)?))
}
#[pymethod(name = "split")]
fn split(self, options: ByteInnerSplitOptions, vm: &VirtualMachine) -> PyResult {
let as_bytes = self
.inner
.split(options, false)?
.iter()
.map(|x| vm.ctx.new_bytes(x.to_vec()))
.collect::<Vec<PyObjectRef>>();
Ok(vm.ctx.new_list(as_bytes))
}
#[pymethod(name = "rsplit")]
fn rsplit(self, options: ByteInnerSplitOptions, vm: &VirtualMachine) -> PyResult {
let as_bytes = self
.inner
.split(options, true)?
.iter()
.map(|x| vm.ctx.new_bytes(x.to_vec()))
.collect::<Vec<PyObjectRef>>();
Ok(vm.ctx.new_list(as_bytes))
}
#[pymethod(name = "partition")]
fn partition(self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let sepa = PyByteInner::try_from_object(vm, sep.clone())?;
let (left, right) = self.inner.partition(&sepa, false)?;
Ok(vm
.ctx
.new_tuple(vec![vm.ctx.new_bytes(left), sep, vm.ctx.new_bytes(right)]))
}
#[pymethod(name = "rpartition")]
fn rpartition(self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let sepa = PyByteInner::try_from_object(vm, sep.clone())?;
let (left, right) = self.inner.partition(&sepa, true)?;
Ok(vm
.ctx
.new_tuple(vec![vm.ctx.new_bytes(left), sep, vm.ctx.new_bytes(right)]))
}
#[pymethod(name = "expandtabs")]
fn expandtabs(self, options: ByteInnerExpandtabsOptions, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.expandtabs(options)))
}
#[pymethod(name = "splitlines")]
fn splitlines(self, options: ByteInnerSplitlinesOptions, vm: &VirtualMachine) -> PyResult {
let as_bytes = self
.inner
.splitlines(options)
.iter()
.map(|x| vm.ctx.new_bytes(x.to_vec()))
.collect::<Vec<PyObjectRef>>();
Ok(vm.ctx.new_list(as_bytes))
}
#[pymethod(name = "zfill")]
fn zfill(self, width: PyIntRef, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.zfill(width)))
}
#[pymethod(name = "replace")]
fn replace(
self,
old: PyByteInner,
new: PyByteInner,
count: OptionalArg<PyIntRef>,
vm: &VirtualMachine,
) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.replace(old, new, count)?))
}
#[pymethod(name = "title")]
fn title(self, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.title()))
}
#[pymethod(name = "__mul__")]
fn repeat(self, n: PyIntRef, vm: &VirtualMachine) -> PyResult {
Ok(vm.ctx.new_bytes(self.inner.repeat(n, vm)?))
}
#[pymethod(name = "__rmul__")]
fn rmul(self, n: PyIntRef, vm: &VirtualMachine) -> PyResult {
self.repeat(n, vm)
}
#[pymethod(name = "decode")]
fn decode(
self,
encoding: OptionalArg<PyStringRef>,
errors: OptionalArg<PyStringRef>,
vm: &VirtualMachine,
) -> PyResult<String> {
let mut strict_mod = true;
let replacing_char = match errors {
OptionalArg::Present(ref input) => match input.as_str() {
"replace" => {
strict_mod = false;
Some('\u{FFFD}')
}
"ignore" => {
strict_mod = false;
None
}
_ => None,
},
OptionalArg::Missing => None,
};
let encoding_type = match encoding {
OptionalArg::Present(ref input) => input.as_str(),
OptionalArg::Missing => "utf-8",
};
let decode_error = Err(vm.new_value_error("DecodeError".to_string()));
let mut decode_content = String::new();
match encoding_type {
"ascii" => {
for &b in self.get_value() {
if b.is_ascii() {
decode_content.push(b as char)
} else if !strict_mod && replacing_char.is_some() {
decode_content.push(replacing_char.unwrap())
}
}
}
"utf-8" | "utf8" | "" => {
let mut p: u32 = 0u32;
let mut remaining_bytes = 0;
for &b in self.get_value() {
if (b as u8) & 128 == 0 {
if b.is_ascii() {
decode_content.push(b as char)
} else if !strict_mod && replacing_char.is_some() {
decode_content.push(replacing_char.unwrap())
}
} else if (b as u8) & 192 == 128 {
remaining_bytes -= 1;
p += u32::from(b as u8 & 63) << (6 * remaining_bytes);
if remaining_bytes == 0 {
match wtf8::CodePoint::from_u32(p) {
Some(cp) => {
if !strict_mod && replacing_char.is_some() {
decode_content.push(cp.to_char_lossy());
} else {
match cp.to_char() {
Some(c) => decode_content.push(c),
None => {
if strict_mod {
return decode_error;
}
}
}
}
}
None => {
if replacing_char.is_some() {
decode_content.push(replacing_char.unwrap())
}
}
}
p = 0u32;
}
} else if (b as u8) & 224 == 192 {
remaining_bytes = 1;
p = u32::from(b as u8 & 31) << 6;
} else if (b as u8) & 240 == 224 {
remaining_bytes = 2;
p = u32::from(b as u8 & 15) << 12;
} else if (b as u8) & 248 == 240 {
remaining_bytes = 3;
p = u32::from(b as u8 & 7) << 18;
} else if !strict_mod && replacing_char.is_some() {
decode_content.push(replacing_char.unwrap())
}
}
}
_ => {
return Err(vm.new_lookup_error(format!("unknown encoding: {}", encoding_type)));
}
}
Ok(decode_content)
}
}
#[pyclass]
#[derive(Debug)]
pub struct PyBytesIterator {
position: Cell<usize>,
bytes: PyBytesRef,
}
impl PyValue for PyBytesIterator {
fn class(vm: &VirtualMachine) -> PyClassRef {
vm.ctx.bytesiterator_type()
}
}
#[pyimpl]
impl PyBytesIterator {
#[pymethod(name = "__next__")]
fn next(&self, vm: &VirtualMachine) -> PyResult<u8> {
if self.position.get() < self.bytes.inner.len() {
let ret = self.bytes[self.position.get()];
self.position.set(self.position.get() + 1);
Ok(ret)
} else {
Err(objiter::new_stop_iteration(vm))
}
}
#[pymethod(name = "__iter__")]
fn iter(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyRef<Self> {
zelf
}
}