diff --git a/.github/workflows/cron-ci.yaml b/.github/workflows/cron-ci.yaml index 492051de8..edb7eb662 100644 --- a/.github/workflows/cron-ci.yaml +++ b/.github/workflows/cron-ci.yaml @@ -56,3 +56,33 @@ jobs: uses: codecov/codecov-action@v1 with: file: ${{ steps.coverage.outputs.report }} + + testdata: + name: Collect regression test data + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@master + - name: build rustpython + uses: actions-rs/cargo@v1 + with: + command: build + args: --release --verbose --all + - name: collect tests data + run: cargo run --release tests/jsontests.py + env: + RUSTPYTHONPATH: ${{ github.workspace }}/Lib + - name: upload tests data to the website + env: + SSHKEY: ${{ secrets.ACTIONS_TESTS_DATA_DEPLOY_KEY }} + GITHUB_ACTOR: ${{ github.actor }} + run: | + echo "$SSHKEY" >~/github_key + chmod 600 ~/github_key + export GIT_SSH_COMMAND="ssh -i ~/github_key" + + git clone git@github.com:RustPython/rustpython.github.io.git website + cd website + cp ../tests/cpython_tests_results.json ./_data/regrtests_results.json + git add ./_data/regrtests_results.json + git -c user.name="Github Actions" -c user.email="actions@github.com" commit -m "Update regression test results" --author="$GITHUB_ACTOR" + git push diff --git a/Lib/test/string_tests.py b/Lib/test/string_tests.py index 7e544d6c0..495044433 100644 --- a/Lib/test/string_tests.py +++ b/Lib/test/string_tests.py @@ -500,7 +500,6 @@ class BaseTest: self.checkraises(ValueError, 'hello', 'rsplit', '') self.checkraises(ValueError, 'hello', 'rsplit', '', 0) - @unittest.skip("TODO: RUSTPYTHON test_bytes") def test_replace(self): EQ = self.checkequal diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py index d49760939..76628739a 100644 --- a/Lib/test/test_bytes.py +++ b/Lib/test/test_bytes.py @@ -1291,7 +1291,8 @@ class ByteArrayTest(BaseBytesTest, unittest.TestCase): del b[4] self.assertEqual(b, bytearray([1, 2, 3, 4, 6, 7, 8])) - @unittest.skip("TODO: RUSTPYTHON") + # TODO: RUSTPYTHON + @unittest.expectedFailure def test_setslice(self): b = bytearray(range(10)) self.assertEqual(list(b), list(range(10))) diff --git a/Lib/types.py b/Lib/types.py index fdb71f9c1..4567549a9 100644 --- a/Lib/types.py +++ b/Lib/types.py @@ -24,11 +24,10 @@ _c = _c() CoroutineType = type(_c) _c.close() # Prevent ResourceWarning -# XXX RUSTPYTHON TODO: async generators -# async def _ag(): -# yield -# _ag = _ag() -# AsyncGeneratorType = type(_ag) +async def _ag(): + yield +_ag = _ag() +AsyncGeneratorType = type(_ag) class _C: def _m(self): pass diff --git a/compiler/src/peephole.rs b/compiler/src/peephole.rs index 67a4b9efe..aaec0470c 100644 --- a/compiler/src/peephole.rs +++ b/compiler/src/peephole.rs @@ -92,7 +92,7 @@ impl PeepholeOptimizer { } fn optimize(&mut self) { - apply_optimizations!(self, operator, unpack); + apply_optimizations!(self, operator /* , unpack */); } } diff --git a/compiler/src/peephole/optimizations.rs b/compiler/src/peephole/optimizations.rs index 28bac862c..461baa7cc 100644 --- a/compiler/src/peephole/optimizations.rs +++ b/compiler/src/peephole/optimizations.rs @@ -73,26 +73,27 @@ pub fn operator(buf: &mut impl OptimizationBuffer) { } } -pub fn unpack(buf: &mut impl OptimizationBuffer) { - let (instruction, meta) = buf.pop(); - if let Instruction::UnpackSequence { size } = instruction { - let (arg, arg_meta) = buf.pop(); - match arg { - Instruction::BuildTuple { - size: tup_size, - unpack, - } if !unpack && tup_size == size => { - buf.emit( - Instruction::Reverse { amount: size }, - vec![arg_meta, meta].into(), - ); - } - arg => { - buf.emit(arg, arg_meta); - buf.emit(instruction, meta); - } - } - } else { - buf.emit(instruction, meta) - } -} +// TODO: make a version of this that doesn't miscompile `a, b = (1, 2) if True else (3, 4)` +// pub fn unpack(buf: &mut impl OptimizationBuffer) { +// let (instruction, meta) = buf.pop(); +// if let Instruction::UnpackSequence { size } = instruction { +// let (arg, arg_meta) = buf.pop(); +// match arg { +// Instruction::BuildTuple { +// size: tup_size, +// unpack, +// } if !unpack && tup_size == size => { +// buf.emit( +// Instruction::Reverse { amount: size }, +// vec![arg_meta, meta].into(), +// ); +// } +// arg => { +// buf.emit(arg, arg_meta); +// buf.emit(instruction, meta); +// } +// } +// } else { +// buf.emit(instruction, meta) +// } +// } diff --git a/compiler/src/symboltable.rs b/compiler/src/symboltable.rs index d08aa4f2e..22acbe1ac 100644 --- a/compiler/src/symboltable.rs +++ b/compiler/src/symboltable.rs @@ -95,7 +95,6 @@ pub struct Symbol { pub name: String, // pub table: SymbolTableRef, pub scope: SymbolScope, - pub is_param: bool, pub is_referenced: bool, pub is_assigned: bool, pub is_parameter: bool, @@ -108,7 +107,6 @@ impl Symbol { name: name.to_owned(), // table, scope: SymbolScope::Unknown, - is_param: false, is_referenced: false, is_assigned: false, is_parameter: false, @@ -741,7 +739,6 @@ impl SymbolTableBuilder { Ok(()) } - #[allow(clippy::single_match)] fn register_name(&mut self, name: &str, role: SymbolUsage) -> SymbolTableResult { let scope_depth = self.tables.len(); let table = self.tables.last_mut().unwrap(); @@ -777,13 +774,11 @@ impl SymbolTableBuilder { // Some more checks: match role { - SymbolUsage::Nonlocal => { - if scope_depth < 2 { - return Err(SymbolTableError { - error: format!("cannot define nonlocal '{}' at top level.", name), - location, - }); - } + SymbolUsage::Nonlocal if scope_depth < 2 => { + return Err(SymbolTableError { + error: format!("cannot define nonlocal '{}' at top level.", name), + location, + }) } _ => { // Ok! diff --git a/derive/src/from_args.rs b/derive/src/from_args.rs index 99d701188..42559fba2 100644 --- a/derive/src/from_args.rs +++ b/derive/src/from_args.rs @@ -16,14 +16,11 @@ enum ParameterKind { impl ParameterKind { fn from_ident(ident: &Ident) -> Option { - if ident == "positional_only" { - Some(ParameterKind::PositionalOnly) - } else if ident == "positional_or_keyword" { - Some(ParameterKind::PositionalOrKeyword) - } else if ident == "keyword_only" { - Some(ParameterKind::KeywordOnly) - } else { - None + match ident.to_string().as_str() { + "positional_only" => Some(ParameterKind::PositionalOnly), + "positional_or_keyword" => Some(ParameterKind::PositionalOrKeyword), + "keyword_only" => Some(ParameterKind::KeywordOnly), + _ => None, } } } @@ -150,6 +147,13 @@ fn generate_field(field: &Field) -> Result { }; let name = &field.ident; + if let Some(name) = name { + if name.to_string().starts_with("_phantom") { + return Ok(quote! { + #name: std::marker::PhantomData, + }); + } + } let middle = quote! { .map(|x| ::rustpython_vm::pyobject::TryFromObject::try_from_object(vm, x)).transpose()? }; @@ -210,8 +214,9 @@ pub fn impl_from_args(input: DeriveInput) -> Result { }; let name = input.ident; + let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl(); let output = quote! { - impl ::rustpython_vm::function::FromArgs for #name { + impl #impl_generics ::rustpython_vm::function::FromArgs for #name #ty_generics #where_clause { fn from_args( vm: &::rustpython_vm::VirtualMachine, args: &mut ::rustpython_vm::function::PyFuncArgs diff --git a/derive/src/lib.rs b/derive/src/lib.rs index cfda0f773..46547071e 100644 --- a/derive/src/lib.rs +++ b/derive/src/lib.rs @@ -12,6 +12,7 @@ mod error; mod compile_bytecode; mod from_args; mod pyclass; +mod pymodule; mod util; use error::{extract_spans, Diagnostic}; @@ -44,6 +45,13 @@ pub fn pyimpl(attr: TokenStream, item: TokenStream) -> TokenStream { result_to_tokens(pyclass::impl_pyimpl(attr, item)) } +#[proc_macro_attribute] +pub fn pymodule(attr: TokenStream, item: TokenStream) -> TokenStream { + let attr = parse_macro_input!(attr as AttributeArgs); + let item = parse_macro_input!(item as Item); + result_to_tokens(pymodule::impl_pymodule(attr, item)) +} + #[proc_macro_attribute] pub fn pystruct_sequence(attr: TokenStream, item: TokenStream) -> TokenStream { let attr = parse_macro_input!(attr as AttributeArgs); diff --git a/derive/src/pyclass.rs b/derive/src/pyclass.rs index 42194395a..1d6474048 100644 --- a/derive/src/pyclass.rs +++ b/derive/src/pyclass.rs @@ -1,11 +1,11 @@ use super::Diagnostic; -use crate::util::path_eq; +use crate::util::{def_to_name, path_eq, strip_prefix, ItemIdent, ItemMeta}; use proc_macro2::{Span, TokenStream as TokenStream2}; use quote::{quote, quote_spanned, ToTokens}; use std::collections::{HashMap, HashSet}; use syn::{ parse_quote, spanned::Spanned, Attribute, AttributeArgs, Ident, Index, Item, Lit, Meta, - NestedMeta, Signature, + NestedMeta, }; fn meta_to_vec(meta: Meta) -> Result, Meta> { @@ -42,181 +42,6 @@ enum ClassItem { }, } -struct ClassItemMeta<'a> { - sig: &'a Signature, - parent_type: &'static str, - meta: HashMap>, -} - -impl<'a> ClassItemMeta<'a> { - const METHOD_NAMES: &'static [&'static str] = &["name", "magic"]; - const PROPERTY_NAMES: &'static [&'static str] = &["name", "magic", "setter"]; - - fn from_nested_meta( - parent_type: &'static str, - sig: &'a Signature, - nested_meta: &[NestedMeta], - names: &[&'static str], - ) -> Result { - let mut extracted = Self { - sig, - parent_type, - meta: HashMap::new(), - }; - - let validate_name = |name: &str, extracted: &Self| -> Result<(), Diagnostic> { - if names.contains(&name) { - if extracted.meta.contains_key(name) { - bail_span!( - &sig.ident, - "#[{}] must have only one '{}'", - parent_type, - name - ); - } else { - Ok(()) - } - } else { - bail_span!( - &sig.ident, - "#[{}({})] is not one of allowed attributes {}", - parent_type, - name, - names.join(", ") - ); - } - }; - - for meta in nested_meta { - let meta = match meta { - NestedMeta::Meta(meta) => meta, - NestedMeta::Lit(_) => continue, - }; - - match meta { - Meta::NameValue(name_value) => { - if let Some(ident) = name_value.path.get_ident() { - let name = ident.to_string(); - validate_name(&name, &extracted)?; - extracted.meta.insert(name, Some(name_value.lit.clone())); - } - } - Meta::Path(path) => { - if let Some(ident) = path.get_ident() { - let name = ident.to_string(); - validate_name(&name, &extracted)?; - extracted.meta.insert(name, None); - } else { - continue; - } - } - _ => (), - } - } - - Ok(extracted) - } - - fn _str(&self, key: &str) -> Result, Diagnostic> { - Ok(match self.meta.get(key) { - Some(Some(lit)) => { - if let Lit::Str(s) = lit { - Some(s.value()) - } else { - bail_span!( - &self.sig.ident, - "#[{}({} = ...)] must be a string", - self.parent_type, - key - ); - } - } - Some(None) => { - bail_span!( - &self.sig.ident, - "#[{}({} = ...)] is expected", - self.parent_type, - key, - ); - } - None => None, - }) - } - - fn _bool(&self, key: &str) -> Result { - Ok(match self.meta.get(key) { - Some(Some(_)) => { - bail_span!( - &self.sig.ident, - "#[{}({})] is expected", - self.parent_type, - key, - ); - } - Some(None) => true, - None => false, - }) - } - - fn method_name(&self) -> Result { - let name = self._str("name")?; - let magic = self._bool("magic")?; - Ok(if let Some(name) = name { - name - } else { - let name = self.sig.ident.to_string(); - if magic { - format!("__{}__", name) - } else { - name - } - }) - } - - fn setter(&self) -> Result { - self._bool("setter") - } - - fn property_name(&self) -> Result { - let magic = self._bool("magic")?; - let setter = self._bool("setter")?; - let name = self._str("name")?; - - Ok(if let Some(name) = name { - name - } else { - let sig_name = self.sig.ident.to_string(); - let name = if setter { - if let Some(name) = strip_prefix(&sig_name, "set_") { - if name.is_empty() { - bail_span!( - &self.sig.ident, - "A #[{}(setter)] fn with a set_* name must \ - have something after \"set_\"", - self.parent_type - ) - } - name.to_string() - } else { - bail_span!( - &self.sig.ident, - "A #[{}(setter)] fn must either have a `name` \ - parameter or a fn name along the lines of \"set_*\"", - self.parent_type - ) - } - } else { - sig_name - }; - if magic { - format!("__{}__", name) - } else { - name - } - }) - } -} - impl Class { fn add_item(&mut self, item: ClassItem, span: Span) -> Result<(), Diagnostic> { if self.items.insert(item) { @@ -229,7 +54,7 @@ impl Class { } } - fn extract_method(sig: &Signature, meta: Meta) -> Result { + fn extract_method(ident: &Ident, meta: Meta) -> Result { let nesteds = meta_to_vec(meta).map_err(|meta| { err_span!( meta, @@ -238,19 +63,15 @@ impl Class { ) })?; - let item_meta = ClassItemMeta::from_nested_meta( - "pymethod", - sig, - &nesteds, - ClassItemMeta::METHOD_NAMES, - )?; + let item_meta = + ItemMeta::from_nested_meta("pymethod", &ident, &nesteds, ItemMeta::ATTRIBUTE_NAMES)?; Ok(ClassItem::Method { - item_ident: sig.ident.clone(), + item_ident: ident.clone(), py_name: item_meta.method_name()?, }) } - fn extract_classmethod(sig: &Signature, meta: Meta) -> Result { + fn extract_classmethod(ident: &Ident, meta: Meta) -> Result { let nesteds = meta_to_vec(meta).map_err(|meta| { err_span!( meta, @@ -258,19 +79,19 @@ impl Class { #[pyclassmethod(name = \"...\")]", ) })?; - let item_meta = ClassItemMeta::from_nested_meta( + let item_meta = ItemMeta::from_nested_meta( "pyclassmethod", - sig, + &ident, &nesteds, - ClassItemMeta::METHOD_NAMES, + ItemMeta::ATTRIBUTE_NAMES, )?; Ok(ClassItem::ClassMethod { - item_ident: sig.ident.clone(), + item_ident: ident.clone(), py_name: item_meta.method_name()?, }) } - fn extract_property(sig: &Signature, meta: Meta) -> Result { + fn extract_property(ident: &Ident, meta: Meta) -> Result { let nesteds = meta_to_vec(meta).map_err(|meta| { err_span!( meta, @@ -278,31 +99,27 @@ impl Class { #[pyproperty(name = \"...\")]" ) })?; - let item_meta = ClassItemMeta::from_nested_meta( - "pyproperty", - sig, - &nesteds, - ClassItemMeta::PROPERTY_NAMES, - )?; + let item_meta = + ItemMeta::from_nested_meta("pyproperty", &ident, &nesteds, ItemMeta::PROPERTY_NAMES)?; Ok(ClassItem::Property { py_name: item_meta.property_name()?, - item_ident: sig.ident.clone(), + item_ident: ident.clone(), setter: item_meta.setter()?, }) } - fn extract_slot(sig: &Signature, meta: Meta) -> Result { + fn extract_slot(ident: &Ident, meta: Meta) -> Result { let pyslot_err = "#[pyslot] must be of the form #[pyslot] or #[pyslot(slotname)]"; let nesteds = meta_to_vec(meta).map_err(|meta| err_span!(meta, "{}", pyslot_err))?; if nesteds.len() > 1 { return Err(Diagnostic::spanned_error("e!(#(#nesteds)*), pyslot_err)); } let slot_ident = if nesteds.is_empty() { - let ident_str = sig.ident.to_string(); + let ident_str = ident.to_string(); if let Some(stripped) = strip_prefix(&ident_str, "tp_") { - proc_macro2::Ident::new(stripped, sig.ident.span()) + proc_macro2::Ident::new(stripped, ident.span()) } else { - sig.ident.clone() + ident.clone() } } else { match nesteds.into_iter().next().unwrap() { @@ -315,14 +132,14 @@ impl Class { }; Ok(ClassItem::Slot { slot_ident, - item_ident: sig.ident.clone(), + item_ident: ident.clone(), }) } fn extract_item_from_syn( &mut self, attrs: &mut Vec, - sig: &Signature, + ident: &Ident, ) -> Result<(), Diagnostic> { let mut attr_idxs = Vec::new(); for (i, meta) in attrs @@ -335,17 +152,16 @@ impl Class { Some(name) => name, None => continue, }; - if name == "pymethod" { - self.add_item(Self::extract_method(sig, meta)?, meta_span)?; - } else if name == "pyclassmethod" { - self.add_item(Self::extract_classmethod(sig, meta)?, meta_span)?; - } else if name == "pyproperty" { - self.add_item(Self::extract_property(sig, meta)?, meta_span)?; - } else if name == "pyslot" { - self.add_item(Self::extract_slot(sig, meta)?, meta_span)?; - } else { - continue; - } + let item = match name.to_string().as_str() { + "pymethod" => Self::extract_method(ident, meta)?, + "pyclassmethod" => Self::extract_classmethod(ident, meta)?, + "pyproperty" => Self::extract_property(ident, meta)?, + "pyslot" => Self::extract_slot(ident, meta)?, + _ => { + continue; + } + }; + self.add_item(item, meta_span)?; attr_idxs.push(i); } let mut i = 0; @@ -365,12 +181,7 @@ impl Class { } } -struct ItemSig<'a> { - attrs: &'a mut Vec, - sig: &'a Signature, -} - -fn extract_impl_items(mut items: Vec) -> Result { +fn extract_impl_items(mut items: Vec) -> Result { let mut diagnostics: Vec = Vec::new(); let mut class = Class::default(); @@ -378,7 +189,7 @@ fn extract_impl_items(mut items: Vec) -> Result Result { - Some(ItemSig { attrs, sig }) + Some(ItemIdent { + attrs, + ident: &sig.ident, + }) } _ => None, }) @@ -574,7 +388,10 @@ pub fn impl_pyimpl(attr: AttributeArgs, item: Item) -> Result { - Some(ItemSig { attrs, sig }) + Some(ItemIdent { + attrs, + ident: &sig.ident, + }) } _ => None, }) @@ -597,30 +414,9 @@ pub fn impl_pyimpl(attr: AttributeArgs, item: Item) -> Result Result { - let mut class_name = None; - for attr in attr { - if let NestedMeta::Meta(meta) = attr { - if let Meta::NameValue(name_value) = meta { - if path_eq(&name_value.path, "name") { - if let Lit::Str(s) = name_value.lit { - class_name = Some(s.value()); - } else { - bail_span!( - name_value.lit, - "#[{}(name = ...)] must be a string", - attr_name - ); - } - } - } - } - } - let class_name = class_name.unwrap_or_else(|| ident.to_string()); - let mut doc: Option> = None; for attr in attrs.iter() { if attr.path.is_ident("doc") { @@ -646,7 +442,7 @@ fn generate_class_def( let ret = quote! { impl ::rustpython_vm::pyobject::PyClassDef for #ident { - const NAME: &'static str = #class_name; + const NAME: &'static str = #name; const DOC: Option<&'static str> = #doc; } }; @@ -663,7 +459,8 @@ pub fn impl_pyclass(attr: AttributeArgs, item: Item) -> Result Result Result(s: &'a str, prefix: &str) -> Option<&'a str> { - if s.starts_with(prefix) { - Some(&s[prefix.len()..]) - } else { - None - } -} diff --git a/derive/src/pymodule.rs b/derive/src/pymodule.rs new file mode 100644 index 000000000..112715213 --- /dev/null +++ b/derive/src/pymodule.rs @@ -0,0 +1,217 @@ +use super::Diagnostic; +use crate::util::{def_to_name, ItemIdent, ItemMeta}; +use proc_macro2::{Span, TokenStream as TokenStream2}; +use quote::{quote, quote_spanned}; +use std::collections::HashSet; +use syn::{parse_quote, spanned::Spanned, Attribute, AttributeArgs, Ident, Item, Meta, NestedMeta}; + +fn meta_to_vec(meta: Meta) -> Result, Meta> { + match meta { + Meta::Path(_) => Ok(Vec::new()), + Meta::List(list) => Ok(list.nested.into_iter().collect()), + Meta::NameValue(_) => Err(meta), + } +} + +#[derive(Default)] +struct Module { + items: HashSet, +} + +#[derive(PartialEq, Eq, Hash)] +enum ModuleItem { + Function { item_ident: Ident, py_name: String }, + Class { item_ident: Ident, py_name: String }, +} + +impl Module { + fn add_item(&mut self, item: ModuleItem, span: Span) -> Result<(), Diagnostic> { + if self.items.insert(item) { + Ok(()) + } else { + Err(Diagnostic::span_error( + span, + "Duplicate #[py*] attribute on pyimpl".to_owned(), + )) + } + } + + fn extract_function(ident: &Ident, meta: Meta) -> Result { + let nesteds = meta_to_vec(meta).map_err(|meta| { + err_span!( + meta, + "#[pyfunction = \"...\"] cannot be a name/value, you probably meant \ + #[pyfunction(name = \"...\")]", + ) + })?; + + let item_meta = + ItemMeta::from_nested_meta("pyfunction", &ident, &nesteds, ItemMeta::SIMPLE_NAMES)?; + Ok(ModuleItem::Function { + item_ident: ident.clone(), + py_name: item_meta.simple_name()?, + }) + } + + fn extract_class(ident: &Ident, meta: Meta) -> Result { + let nesteds = meta_to_vec(meta).map_err(|meta| { + err_span!( + meta, + "#[pyclass = \"...\"] cannot be a name/value, you probably meant \ + #[pyclass(name = \"...\")]", + ) + })?; + + let item_meta = + ItemMeta::from_nested_meta("pyclass", &ident, &nesteds, ItemMeta::SIMPLE_NAMES)?; + Ok(ModuleItem::Class { + item_ident: ident.clone(), + py_name: item_meta.simple_name()?, + }) + } + + fn extract_item_from_syn( + &mut self, + attrs: &mut Vec, + ident: &Ident, + ) -> Result<(), Diagnostic> { + let mut attr_idxs = Vec::new(); + for (i, meta) in attrs + .iter() + .filter_map(|attr| attr.parse_meta().ok()) + .enumerate() + { + let meta_span = meta.span(); + let name = match meta.path().get_ident() { + Some(name) => name, + None => continue, + }; + let item = match name.to_string().as_str() { + "pyfunction" => { + attr_idxs.push(i); + Self::extract_function(ident, meta)? + } + "pyclass" => Self::extract_class(ident, meta)?, + _ => { + continue; + } + }; + self.add_item(item, meta_span)?; + } + let mut i = 0; + let mut attr_idxs = &*attr_idxs; + attrs.retain(|_| { + let drop = attr_idxs.first().copied() == Some(i); + if drop { + attr_idxs = &attr_idxs[1..]; + } + i += 1; + !drop + }); + for (i, idx) in attr_idxs.iter().enumerate() { + attrs.remove(idx - i); + } + Ok(()) + } +} + +fn extract_module_items(mut items: Vec) -> Result { + let mut diagnostics: Vec = Vec::new(); + + let mut module = Module::default(); + + for item in items.iter_mut() { + push_diag_result!( + diagnostics, + module.extract_item_from_syn(&mut item.attrs, item.ident), + ); + } + + let functions = module.items.into_iter().map(|item| match item { + ModuleItem::Function { + item_ident, + py_name, + } => { + let new_func = quote_spanned!(item_ident.span() => .new_function(#item_ident)); + quote! { + vm.__module_set_attr(&module, #py_name, vm.ctx#new_func).unwrap(); + } + } + ModuleItem::Class { + item_ident, + py_name, + } => { + let new_class = quote_spanned!(item_ident.span() => #item_ident::make_class(&vm.ctx)); + quote! { + vm.__module_set_attr(&module, #py_name, #new_class).unwrap(); + } + } + }); + + Diagnostic::from_vec(diagnostics)?; + + Ok(quote! { + #(#functions)* + }) +} + +pub fn impl_pymodule(attr: AttributeArgs, item: Item) -> Result { + match item { + Item::Mod(mut module) => { + let module_name = def_to_name(&module.ident, "pymodule", attr)?; + + if let Some(content) = module.content.as_mut() { + let items = content + .1 + .iter_mut() + .filter_map(|item| match item { + Item::Fn(syn::ItemFn { attrs, sig, .. }) => Some(ItemIdent { + attrs, + ident: &sig.ident, + }), + Item::Struct(syn::ItemStruct { attrs, ident, .. }) => { + Some(ItemIdent { attrs, ident }) + } + Item::Enum(syn::ItemEnum { attrs, ident, .. }) => { + Some(ItemIdent { attrs, ident }) + } + _ => None, + }) + .collect(); + + let extend_mod = extract_module_items(items)?; + content.1.push(parse_quote! { + const MODULE_NAME: &str = #module_name; + }); + content.1.push(parse_quote! { + pub(crate) fn extend_module( + vm: &::rustpython_vm::vm::VirtualMachine, + module: &::rustpython_vm::pyobject::PyObjectRef, + ) { + #extend_mod + } + }); + content.1.push(parse_quote! { + #[allow(dead_code)] + pub(crate) fn make_module( + vm: &::rustpython_vm::vm::VirtualMachine + ) -> ::rustpython_vm::pyobject::PyObjectRef { + let module = vm.new_module(MODULE_NAME, vm.ctx.new_dict()); + extend_module(vm, &module); + module + } + }); + + Ok(quote! { + #module + }) + } else { + bail_span!( + module, + "#[pymodule] can only be on a module declaration with body" + ) + } + } + other => bail_span!(other, "#[pymodule] can only be on a module declaration"), + } +} diff --git a/derive/src/util.rs b/derive/src/util.rs index 7284d8014..deb2ef046 100644 --- a/derive/src/util.rs +++ b/derive/src/util.rs @@ -1,3 +1,216 @@ -pub fn path_eq(path: &syn::Path, s: &str) -> bool { +use super::Diagnostic; +use std::collections::HashMap; +use syn::{Attribute, AttributeArgs, Ident, Lit, Meta, NestedMeta, Path}; + +pub fn path_eq(path: &Path, s: &str) -> bool { path.get_ident().map_or(false, |id| id == s) } + +pub fn def_to_name( + ident: &Ident, + attr_name: &'static str, + attr: AttributeArgs, +) -> Result { + let mut name = None; + for attr in attr { + if let NestedMeta::Meta(meta) = attr { + if let Meta::NameValue(name_value) = meta { + if path_eq(&name_value.path, "name") { + if let Lit::Str(s) = name_value.lit { + name = Some(s.value()); + } else { + bail_span!( + name_value.lit, + "#[{}(name = ...)] must be a string", + attr_name + ); + } + } + } + } + } + Ok(name.unwrap_or_else(|| ident.to_string())) +} + +pub fn strip_prefix<'a>(s: &'a str, prefix: &str) -> Option<&'a str> { + if s.starts_with(prefix) { + Some(&s[prefix.len()..]) + } else { + None + } +} + +pub struct ItemIdent<'a> { + pub attrs: &'a mut Vec, + pub ident: &'a Ident, +} + +pub struct ItemMeta<'a> { + ident: &'a Ident, + parent_type: &'static str, + meta: HashMap>, +} + +impl<'a> ItemMeta<'a> { + pub const SIMPLE_NAMES: &'static [&'static str] = &["name"]; + pub const ATTRIBUTE_NAMES: &'static [&'static str] = &["name", "magic"]; + pub const PROPERTY_NAMES: &'static [&'static str] = &["name", "magic", "setter"]; + + pub fn from_nested_meta( + parent_type: &'static str, + ident: &'a Ident, + nested_meta: &[NestedMeta], + names: &[&'static str], + ) -> Result { + let mut extracted = Self { + ident, + parent_type, + meta: HashMap::new(), + }; + + let validate_name = |name: &str, extracted: &Self| -> Result<(), Diagnostic> { + if names.contains(&name) { + if extracted.meta.contains_key(name) { + bail_span!(ident, "#[{}] must have only one '{}'", parent_type, name); + } else { + Ok(()) + } + } else { + bail_span!( + ident, + "#[{}({})] is not one of allowed attributes {}", + parent_type, + name, + names.join(", ") + ); + } + }; + + for meta in nested_meta { + let meta = match meta { + NestedMeta::Meta(meta) => meta, + NestedMeta::Lit(_) => continue, + }; + + match meta { + Meta::NameValue(name_value) => { + if let Some(ident) = name_value.path.get_ident() { + let name = ident.to_string(); + validate_name(&name, &extracted)?; + extracted.meta.insert(name, Some(name_value.lit.clone())); + } + } + Meta::Path(path) => { + if let Some(ident) = path.get_ident() { + let name = ident.to_string(); + validate_name(&name, &extracted)?; + extracted.meta.insert(name, None); + } else { + continue; + } + } + _ => (), + } + } + + Ok(extracted) + } + + fn _str(&self, key: &str) -> Result, Diagnostic> { + Ok(match self.meta.get(key) { + Some(Some(lit)) => { + if let Lit::Str(s) = lit { + Some(s.value()) + } else { + bail_span!( + &self.ident, + "#[{}({} = ...)] must be a string", + self.parent_type, + key + ); + } + } + Some(None) => { + bail_span!( + &self.ident, + "#[{}({} = ...)] is expected", + self.parent_type, + key, + ); + } + None => None, + }) + } + + fn _bool(&self, key: &str) -> Result { + Ok(match self.meta.get(key) { + Some(Some(_)) => { + bail_span!(&self.ident, "#[{}({})] is expected", self.parent_type, key,); + } + Some(None) => true, + None => false, + }) + } + + pub fn simple_name(&self) -> Result { + Ok(self._str("name")?.unwrap_or_else(|| self.ident.to_string())) + } + + pub fn method_name(&self) -> Result { + let name = self._str("name")?; + let magic = self._bool("magic")?; + Ok(if let Some(name) = name { + name + } else { + let name = self.ident.to_string(); + if magic { + format!("__{}__", name) + } else { + name + } + }) + } + + pub fn property_name(&self) -> Result { + let magic = self._bool("magic")?; + let setter = self._bool("setter")?; + let name = self._str("name")?; + + Ok(if let Some(name) = name { + name + } else { + let sig_name = self.ident.to_string(); + let name = if setter { + if let Some(name) = strip_prefix(&sig_name, "set_") { + if name.is_empty() { + bail_span!( + &self.ident, + "A #[{}(setter)] fn with a set_* name must \ + have something after \"set_\"", + self.parent_type + ) + } + name.to_string() + } else { + bail_span!( + &self.ident, + "A #[{}(setter)] fn must either have a `name` \ + parameter or a fn name along the lines of \"set_*\"", + self.parent_type + ) + } + } else { + sig_name + }; + if magic { + format!("__{}__", name) + } else { + name + } + }) + } + + pub fn setter(&self) -> Result { + self._bool("setter") + } +} diff --git a/parser/build.rs b/parser/build.rs index 19f102721..0bfc79276 100644 --- a/parser/build.rs +++ b/parser/build.rs @@ -1,5 +1,3 @@ -use lalrpop; - fn main() { lalrpop::process_root().unwrap() } diff --git a/parser/src/function.rs b/parser/src/function.rs index 96366d3c5..d62ea28a2 100644 --- a/parser/src/function.rs +++ b/parser/src/function.rs @@ -6,7 +6,6 @@ use crate::error::{LexicalError, LexicalErrorType}; type ParameterDefs = (Vec, Vec); type ParameterDef = (ast::Parameter, Option); -#[allow(clippy::collapsible_if)] pub fn parse_params( params: (Vec, Vec), ) -> Result { @@ -16,15 +15,13 @@ pub fn parse_params( let mut try_default = |name: &ast::Parameter, default| { if let Some(default) = default { defaults.push(default); - } else { - if !defaults.is_empty() { - // Once we have started with defaults, all remaining arguments must - // have defaults - return Err(LexicalError { - error: LexicalErrorType::DefaultArgumentError, - location: name.location, - }); - } + } else if !defaults.is_empty() { + // Once we have started with defaults, all remaining arguments must + // have defaults + return Err(LexicalError { + error: LexicalErrorType::DefaultArgumentError, + location: name.location, + }); } Ok(()) }; diff --git a/tests/snippets/if_expressions.py b/tests/snippets/if_expressions.py index 143cb6458..ac5f49273 100644 --- a/tests/snippets/if_expressions.py +++ b/tests/snippets/if_expressions.py @@ -24,3 +24,7 @@ def func2(): return 20 assert ret(func1() or func2()) == 20 + +a, b = (1, 2) if True else (3, 4) +assert a == 1 +assert b == 2 diff --git a/vm/src/dictdatatype.rs b/vm/src/dictdatatype.rs index 5541f1758..43a9fceb0 100644 --- a/vm/src/dictdatatype.rs +++ b/vm/src/dictdatatype.rs @@ -1,6 +1,6 @@ use crate::obj::objstr::PyString; use crate::pyhash; -use crate::pyobject::{IdProtocol, IntoPyObject, PyObjectRef, PyResult}; +use crate::pyobject::{IdProtocol, IntoPyObject, PyObjectRef, PyResult, ThreadSafe}; use crate::vm::VirtualMachine; use num_bigint::ToBigInt; /// Ordered dictionary implementation. @@ -10,6 +10,7 @@ use num_bigint::ToBigInt; use std::collections::{hash_map::DefaultHasher, HashMap}; use std::hash::{Hash, Hasher}; use std::mem::size_of; +use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; /// hash value of an object returned by __hash__ type HashValue = pyhash::PyHash; @@ -18,30 +19,64 @@ type HashIndex = pyhash::PyHash; /// entry index mapped in indices type EntryIndex = usize; -#[derive(Clone)] pub struct Dict { + inner: RwLock>, +} + +impl ThreadSafe for Dict {} + +struct InnerDict { size: usize, indices: HashMap, entries: Vec>>, } +impl Clone for InnerDict { + fn clone(&self) -> Self { + InnerDict { + size: self.size, + indices: self.indices.clone(), + entries: self.entries.clone(), + } + } +} + +impl Clone for Dict { + fn clone(&self) -> Self { + Dict { + inner: RwLock::new(self.inner.read().unwrap().clone()), + } + } +} + impl Default for Dict { fn default() -> Self { Dict { - size: 0, - indices: HashMap::new(), - entries: Vec::new(), + inner: RwLock::new(InnerDict { + size: 0, + indices: HashMap::new(), + entries: Vec::new(), + }), } } } -#[derive(Clone)] struct DictEntry { hash: HashValue, key: PyObjectRef, value: T, } +impl Clone for DictEntry { + fn clone(&self) -> Self { + DictEntry { + hash: self.hash, + key: self.key.clone(), + value: self.value.clone(), + } + } +} + #[derive(Debug)] pub struct DictSize { size: usize, @@ -49,10 +84,19 @@ pub struct DictSize { } impl Dict { - fn resize(&mut self) { - let mut new_indices = HashMap::with_capacity(self.size); - let mut new_entries = Vec::with_capacity(self.size); - for maybe_entry in self.entries.drain(0..) { + fn borrow_value(&self) -> RwLockReadGuard<'_, InnerDict> { + self.inner.read().unwrap() + } + + fn borrow_value_mut(&self) -> RwLockWriteGuard<'_, InnerDict> { + self.inner.write().unwrap() + } + + fn resize(&self) { + let mut inner = self.borrow_value_mut(); + let mut new_indices = HashMap::with_capacity(inner.size); + let mut new_entries = Vec::with_capacity(inner.size); + for maybe_entry in inner.entries.drain(0..) { if let Some(entry) = maybe_entry { let mut hash_index = entry.hash; // Faster version of lookup. No equality checks here. @@ -64,12 +108,12 @@ impl Dict { new_entries.push(Some(entry)); } } - self.indices = new_indices; - self.entries = new_entries; + inner.indices = new_indices; + inner.entries = new_entries; } fn unchecked_push( - &mut self, + &self, hash_index: HashIndex, hash_value: HashValue, key: PyObjectRef, @@ -80,44 +124,46 @@ impl Dict { key, value, }; - let entry_index = self.entries.len(); - self.entries.push(Some(entry)); - self.indices.insert(hash_index, entry_index); - self.size += 1; - } - - fn unchecked_delete(&mut self, entry_index: EntryIndex) { - self.entries[entry_index] = None; - self.size -= 1; + let mut inner = self.borrow_value_mut(); + let entry_index = inner.entries.len(); + inner.entries.push(Some(entry)); + inner.indices.insert(hash_index, entry_index); + inner.size += 1; } /// Store a key pub fn insert( - &mut self, + &self, vm: &VirtualMachine, key: K, value: T, ) -> PyResult<()> { - if self.indices.len() > 2 * self.size { + // This does not need to be accurate so we can take the lock mutiple times. + if self.borrow_value().indices.len() > 2 * self.borrow_value().size { self.resize(); } - match self.lookup(vm, key)? { - LookupResult::Existing(index) => { - // Update existing key - if let Some(ref mut entry) = self.entries[index] { - entry.value = value; - Ok(()) - } else { - panic!("Lookup returned invalid index into entries!"); + loop { + match self.lookup(vm, key)? { + LookupResult::Existing(index) => { + let mut inner = self.borrow_value_mut(); + // Update existing key + if let Some(ref mut entry) = inner.entries[index] { + // They entry might have changed since we did lookup. Should we update the key? + entry.value = value; + break Ok(()); + } else { + // The dict was changed since we did lookup. Let's try again. + continue; + } + } + LookupResult::NewIndex { + hash_index, + hash_value, + } => { + // New key: + self.unchecked_push(hash_index, hash_value, key.into_pyobject(vm)?, value); + break Ok(()); } - } - LookupResult::NewIndex { - hash_index, - hash_value, - } => { - // New key: - self.unchecked_push(hash_index, hash_value, key.into_pyobject(vm)?, value); - Ok(()) } } } @@ -130,32 +176,32 @@ impl Dict { } } - fn unchecked_get(&self, index: EntryIndex) -> T { - if let Some(entry) = &self.entries[index] { - entry.value.clone() - } else { - panic!("Lookup returned invalid index into entries!"); - } - } - /// Retrieve a key #[cfg_attr(feature = "flame-it", flame("Dict"))] pub fn get(&self, vm: &VirtualMachine, key: K) -> PyResult> { - if let LookupResult::Existing(index) = self.lookup(vm, key)? { - Ok(Some(self.unchecked_get(index))) - } else { - Ok(None) + loop { + if let LookupResult::Existing(index) = self.lookup(vm, key)? { + if let Some(entry) = &self.borrow_value().entries[index] { + break Ok(Some(entry.value.clone())); + } else { + // The dict was changed since we did lookup. Let's try again. + continue; + } + } else { + break Ok(None); + } } } - pub fn clear(&mut self) { - self.entries.clear(); - self.indices.clear(); - self.size = 0 + pub fn clear(&self) { + let mut inner = self.borrow_value_mut(); + inner.entries.clear(); + inner.indices.clear(); + inner.size = 0 } /// Delete a key - pub fn delete(&mut self, vm: &VirtualMachine, key: &PyObjectRef) -> PyResult<()> { + pub fn delete(&self, vm: &VirtualMachine, key: &PyObjectRef) -> PyResult<()> { if self.delete_if_exists(vm, key)? { Ok(()) } else { @@ -163,33 +209,56 @@ impl Dict { } } - pub fn delete_if_exists(&mut self, vm: &VirtualMachine, key: &PyObjectRef) -> PyResult { - if let LookupResult::Existing(entry_index) = self.lookup(vm, key)? { - self.unchecked_delete(entry_index); - Ok(true) - } else { - Ok(false) + pub fn delete_if_exists(&self, vm: &VirtualMachine, key: &PyObjectRef) -> PyResult { + loop { + if let LookupResult::Existing(entry_index) = self.lookup(vm, key)? { + let mut inner = self.borrow_value_mut(); + if inner.entries[entry_index].is_some() { + inner.entries[entry_index] = None; + inner.size -= 1; + break Ok(true); + } else { + // The dict was changed since we did lookup. Let's try again. + continue; + } + } else { + break Ok(false); + } } } pub fn delete_or_insert( - &mut self, + &self, vm: &VirtualMachine, key: &PyObjectRef, value: T, ) -> PyResult<()> { - match self.lookup(vm, key)? { - LookupResult::Existing(entry_index) => self.unchecked_delete(entry_index), - LookupResult::NewIndex { - hash_value, - hash_index, - } => self.unchecked_push(hash_index, hash_value, key.clone(), value), - }; - Ok(()) + loop { + match self.lookup(vm, key)? { + LookupResult::Existing(entry_index) => { + let mut inner = self.borrow_value_mut(); + if inner.entries[entry_index].is_some() { + inner.entries[entry_index] = None; + inner.size -= 1; + break Ok(()); + } else { + // The dict was changed since we did lookup. Let's try again. + continue; + } + } + LookupResult::NewIndex { + hash_value, + hash_index, + } => { + self.unchecked_push(hash_index, hash_value, key.clone(), value); + break Ok(()); + } + }; + } } pub fn len(&self) -> usize { - self.size + self.borrow_value().size } pub fn is_empty(&self) -> bool { @@ -197,33 +266,42 @@ impl Dict { } pub fn size(&self) -> DictSize { + let inner = self.borrow_value(); DictSize { - size: self.size, - entries_size: self.entries.len(), + size: inner.size, + entries_size: inner.entries.len(), } } - pub fn next_entry(&self, position: &mut EntryIndex) -> Option<(&PyObjectRef, &T)> { - self.entries[*position..].iter().find_map(|entry| { - *position += 1; - entry - .as_ref() - .map(|DictEntry { key, value, .. }| (key, value)) - }) + pub fn next_entry(&self, position: &mut EntryIndex) -> Option<(PyObjectRef, T)> { + self.borrow_value().entries[*position..] + .iter() + .find_map(|entry| { + *position += 1; + entry + .as_ref() + .map(|DictEntry { key, value, .. }| (key.clone(), value.clone())) + }) } pub fn len_from_entry_index(&self, position: EntryIndex) -> usize { - self.entries[position..].iter().flatten().count() + self.borrow_value().entries[position..] + .iter() + .flatten() + .count() } pub fn has_changed_size(&self, position: &DictSize) -> bool { - position.size != self.size || self.entries.len() != position.entries_size + let inner = self.borrow_value(); + position.size != inner.size || inner.entries.len() != position.entries_size } - pub fn keys<'a>(&'a self) -> impl Iterator + 'a { - self.entries + pub fn keys(&self) -> Vec { + self.borrow_value() + .entries .iter() .filter_map(|v| v.as_ref().map(|v| v.key.clone())) + .collect() } /// Lookup the index for the given key. @@ -232,33 +310,41 @@ impl Dict { let hash_value = key.do_hash(vm)?; let perturb = hash_value; let mut hash_index: HashIndex = hash_value; - loop { - if self.indices.contains_key(&hash_index) { - // Now we have an index, lets check the key. - let index = self.indices[&hash_index]; - if let Some(entry) = &self.entries[index] { - // Okay, we have an entry at this place - if key.do_is(&entry.key) { - // Literally the same object - break Ok(LookupResult::Existing(index)); - } else if entry.hash == hash_value { - if key.do_eq(vm, &entry.key)? { - break Ok(LookupResult::Existing(index)); + 'outer: loop { + let (entry, index) = loop { + let inner = self.borrow_value(); + if inner.indices.contains_key(&hash_index) { + // Now we have an index, lets check the key. + let index = inner.indices[&hash_index]; + if let Some(entry) = &inner.entries[index] { + // Okay, we have an entry at this place + if key.do_is(&entry.key) { + // Literally the same object + break 'outer Ok(LookupResult::Existing(index)); + } else if entry.hash == hash_value { + break (entry.clone(), index); } else { // entry mismatch. } } else { - // entry mismatch. + // Removed entry, continue search... } } else { - // Removed entry, continue search... + // Hash not in table, we are at free slot now. + break 'outer Ok(LookupResult::NewIndex { + hash_value, + hash_index, + }); } + // Update i to next probe location: + hash_index = Self::next_index(perturb, hash_index) + // warn!("Perturb value: {}", i); + }; + // This comparison needs to be done outside the lock. + if key.do_eq(vm, &entry.key)? { + break Ok(LookupResult::Existing(index)); } else { - // Hash not in table, we are at free slot now. - break Ok(LookupResult::NewIndex { - hash_value, - hash_index, - }); + // entry mismatch. } // Update i to next probe location: @@ -275,30 +361,45 @@ impl Dict { } /// Retrieve and delete a key - pub fn pop(&mut self, vm: &VirtualMachine, key: K) -> PyResult> { - if let LookupResult::Existing(index) = self.lookup(vm, key)? { - let value = self.unchecked_get(index); - self.unchecked_delete(index); - Ok(Some(value)) - } else { - Ok(None) + pub fn pop(&self, vm: &VirtualMachine, key: K) -> PyResult> { + loop { + if let LookupResult::Existing(index) = self.lookup(vm, key)? { + let mut inner = self.borrow_value_mut(); + if let Some(entry) = &inner.entries[index] { + let value = entry.value.clone(); + inner.entries[index] = None; + inner.size -= 1; + break Ok(Some(value)); + } else { + // The dict was changed since we did lookup. Let's try again. + continue; + } + } else { + break Ok(None); + } } } - pub fn pop_front(&mut self) -> Option<(PyObjectRef, T)> { - let mut entry_index = 0; - match self.next_entry(&mut entry_index) { - Some((key, value)) => { - let item = (key.clone(), value.clone()); - self.unchecked_delete(entry_index - 1); - Some(item) - } - None => None, + pub fn pop_front(&self) -> Option<(PyObjectRef, T)> { + let mut position = 0; + let mut inner = self.borrow_value_mut(); + let first_item = inner.entries.iter().find_map(|entry| { + position += 1; + entry + .as_ref() + .map(|DictEntry { key, value, .. }| (key.clone(), value.clone())) + }); + if let Some(item) = first_item { + inner.entries[position - 1] = None; + inner.size -= 1; + Some(item) + } else { + None } } pub fn sizeof(&self) -> usize { - size_of::() + self.size * size_of::>() + size_of::() + self.borrow_value().size * size_of::>() } } @@ -389,7 +490,7 @@ mod tests { #[test] fn test_insert() { let vm: VirtualMachine = Default::default(); - let mut dict = Dict::default(); + let dict = Dict::default(); assert_eq!(0, dict.len()); let key1 = vm.new_bool(true); diff --git a/vm/src/lib.rs b/vm/src/lib.rs index 1dbf723bc..b324f1b58 100644 --- a/vm/src/lib.rs +++ b/vm/src/lib.rs @@ -6,11 +6,7 @@ //! - Base objects // for methods like vm.to_str(), not the typical use of 'to' as a method prefix -#![allow( - clippy::wrong_self_convention, - clippy::let_and_return, - clippy::implicit_hasher -)] +#![allow(clippy::wrong_self_convention, clippy::implicit_hasher)] #![doc(html_logo_url = "https://raw.githubusercontent.com/RustPython/RustPython/master/logo.png")] #![doc(html_root_url = "https://docs.rs/rustpython-vm/")] diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs index c79aa38f6..1c7d52316 100644 --- a/vm/src/obj/objbytearray.rs +++ b/vm/src/obj/objbytearray.rs @@ -4,16 +4,15 @@ use std::convert::TryFrom; use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; use super::objbyteinner::{ - ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions, - ByteInnerSplitOptions, ByteInnerSplitlinesOptions, ByteInnerTranslateOptions, ByteOr, - PyByteInner, + ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions, ByteInnerSplitOptions, + ByteInnerTranslateOptions, ByteOr, PyByteInner, }; use super::objint::PyIntRef; use super::objiter; use super::objslice::PySliceRef; use super::objstr::{PyString, PyStringRef}; use super::objtype::PyClassRef; -use super::pystr::PyCommonString; +use super::pystr::{self, PyCommonString}; use crate::cformat::CFormatString; use crate::function::{OptionalArg, OptionalOption}; use crate::obj::objstr::do_cformat_string; @@ -301,17 +300,9 @@ impl PyByteArray { } #[pymethod(name = "endswith")] - fn endswith( - &self, - suffix: PyObjectRef, - start: OptionalArg>, - end: OptionalArg>, - vm: &VirtualMachine, - ) -> PyResult { + fn endswith(&self, options: pystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult { self.borrow_value().elements[..].py_startsendswith( - suffix, - start, - end, + options, "endswith", "bytes", |s, x: &PyByteInner| s.ends_with(&x.elements[..]), @@ -322,15 +313,11 @@ impl PyByteArray { #[pymethod(name = "startswith")] fn startswith( &self, - prefix: PyObjectRef, - start: OptionalArg>, - end: OptionalArg>, + options: pystr::StartsEndsWithArgs, vm: &VirtualMachine, ) -> PyResult { self.borrow_value().elements[..].py_startsendswith( - prefix, - start, - end, + options, "startswith", "bytes", |s, x: &PyByteInner| s.starts_with(&x.elements[..]), @@ -440,12 +427,12 @@ impl PyByteArray { } #[pymethod(name = "expandtabs")] - fn expandtabs(&self, options: ByteInnerExpandtabsOptions) -> PyByteArray { + fn expandtabs(&self, options: pystr::ExpandTabsArgs) -> PyByteArray { self.borrow_value().expandtabs(options).into() } #[pymethod(name = "splitlines")] - fn splitlines(&self, options: ByteInnerSplitlinesOptions, vm: &VirtualMachine) -> PyResult { + fn splitlines(&self, options: pystr::SplitLinesArgs, vm: &VirtualMachine) -> PyResult { let as_bytes = self .borrow_value() .splitlines(options) @@ -465,9 +452,10 @@ impl PyByteArray { &self, old: PyByteInner, new: PyByteInner, - count: OptionalArg, + count: OptionalArg, + vm: &VirtualMachine, ) -> PyResult { - Ok(self.borrow_value().replace(old, new, count)?.into()) + Ok(self.borrow_value().replace(old, new, count, vm)?.into()) } #[pymethod(name = "clear")] diff --git a/vm/src/obj/objbyteinner.rs b/vm/src/obj/objbyteinner.rs index 1a4c1d305..dffd6c95e 100644 --- a/vm/src/obj/objbyteinner.rs +++ b/vm/src/obj/objbyteinner.rs @@ -14,7 +14,7 @@ use super::objnone::PyNoneRef; use super::objsequence::PySliceableSequence; use super::objslice::PySliceRef; use super::objstr::{self, PyString, PyStringRef}; -use super::pystr::{self, PyCommonString, StringRange}; +use super::pystr::{self, PyCommonString, PyCommonStringWrapper, StringRange}; use crate::function::{OptionalArg, OptionalOption}; use crate::pyhash; use crate::pyobject::{ @@ -255,43 +255,7 @@ impl ByteInnerTranslateOptions { } } -#[derive(FromArgs)] -pub struct ByteInnerSplitOptions { - #[pyarg(positional_or_keyword, default = "None")] - sep: Option, - #[pyarg(positional_or_keyword, default = "-1")] - maxsplit: isize, -} - -impl ByteInnerSplitOptions { - pub fn get_value(self, vm: &VirtualMachine) -> PyResult<(Option>, isize)> { - let sep = if let Some(s) = self.sep { - let sep = s.elements; - if sep.is_empty() { - return Err(vm.new_value_error("empty separator".to_owned())); - } - Some(sep) - } else { - None - }; - Ok((sep, self.maxsplit)) - } -} - -#[derive(FromArgs)] -pub struct ByteInnerExpandtabsOptions { - #[pyarg(positional_or_keyword, optional = true)] - tabsize: OptionalArg, -} - -impl ByteInnerExpandtabsOptions { - pub fn get_value(self) -> usize { - match self.tabsize.into_option() { - Some(int) => int.as_bigint().to_usize().unwrap_or(0), - None => 8, - } - } -} +pub type ByteInnerSplitOptions = pystr::SplitArgs; #[derive(FromArgs)] pub struct ByteInnerSplitlinesOptions { @@ -460,9 +424,12 @@ impl PyByteInner { }), }; let items = sec?; - let range = self + let mut range = self .elements .get_slice_range(&slice.start_index(vm)?, &slice.stop_index(vm)?); + if range.end < range.start { + range.end = range.start; + } self.elements.splice(range, items); Ok(vm .ctx @@ -926,36 +893,33 @@ impl PyByteInner { } pub fn strip(&self, chars: OptionalOption) -> Vec { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.elements, - None => return self.elements.trim().to_owned(), - }; self.elements - .trim_with(|c| chars.contains(&(c as u8))) - .to_owned() + .py_strip( + chars, + |s, chars| s.trim_with(|c| chars.contains(&(c as u8))), + |s| s.trim(), + ) + .to_vec() } pub fn lstrip(&self, chars: OptionalOption) -> Vec { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.elements, - None => return self.elements.trim_start().to_owned(), - }; self.elements - .trim_start_with(|c| chars.contains(&(c as u8))) - .to_owned() + .py_strip( + chars, + |s, chars| s.trim_start_with(|c| chars.contains(&(c as u8))), + |s| s.trim_start(), + ) + .to_vec() } pub fn rstrip(&self, chars: OptionalOption) -> Vec { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.elements, - None => return self.elements.trim_end().to_owned(), - }; self.elements - .trim_end_with(|c| chars.contains(&(c as u8))) - .to_owned() + .py_strip( + chars, + |s, chars| s.trim_end_with(|c| chars.contains(&(c as u8))), + |s| s.trim_end(), + ) + .to_vec() } pub fn split( @@ -967,19 +931,13 @@ impl PyByteInner { where F: Fn(&[u8], &VirtualMachine) -> PyObjectRef, { - let (sep, maxsplit) = options.get_value(vm)?; - let sep_ref = match sep { - Some(ref v) => Some(&v[..]), - None => None, - }; let elements = self.elements.py_split( - sep_ref, - maxsplit, + options, vm, |v, s, vm| v.split_str(s).map(|v| convert(v, vm)).collect(), |v, s, n, vm| v.splitn_str(n, s).map(|v| convert(v, vm)).collect(), |v, n, vm| v.py_split_whitespace(n, |v| convert(v, vm)), - ); + )?; Ok(vm.ctx.new_list(elements)) } @@ -992,19 +950,13 @@ impl PyByteInner { where F: Fn(&[u8], &VirtualMachine) -> PyObjectRef, { - let (sep, maxsplit) = options.get_value(vm)?; - let sep_ref = match sep { - Some(ref v) => Some(&v[..]), - None => None, - }; let mut elements = self.elements.py_split( - sep_ref, - maxsplit, + options, vm, |v, s, vm| v.rsplit_str(s).map(|v| convert(v, vm)).collect(), |v, s, n, vm| v.rsplitn_str(n, s).map(|v| convert(v, vm)).collect(), |v, n, vm| v.py_rsplit_whitespace(n, |v| convert(v, vm)), - ); + )?; elements.reverse(); Ok(vm.ctx.new_list(elements)) } @@ -1047,8 +999,8 @@ impl PyByteInner { Ok((front, has_mid, back)) } - pub fn expandtabs(&self, options: ByteInnerExpandtabsOptions) -> Vec { - let tabsize = options.get_value(); + pub fn expandtabs(&self, options: pystr::ExpandTabsArgs) -> Vec { + let tabsize = options.tabsize(); let mut counter: usize = 0; let mut res = vec![]; @@ -1079,9 +1031,7 @@ impl PyByteInner { res } - pub fn splitlines(&self, options: ByteInnerSplitlinesOptions) -> Vec<&[u8]> { - let keepends = options.get_value(); - + pub fn splitlines(&self, options: pystr::SplitLinesArgs) -> Vec<&[u8]> { let mut res = vec![]; if self.elements.is_empty() { @@ -1090,7 +1040,7 @@ impl PyByteInner { let mut prev_index = 0; let mut index = 0; - let keep = if keepends { 1 } else { 0 }; + let keep = if options.keepends { 1 } else { 0 }; let slice = &self.elements; while index < slice.len() { @@ -1127,41 +1077,166 @@ impl PyByteInner { bytes_zfill(&self.elements, width.to_usize().unwrap_or(0)) } - pub fn replace( - &self, - old: PyByteInner, - new: PyByteInner, - count: OptionalArg, - ) -> PyResult> { - let count = match count.into_option() { - Some(int) => int - .as_bigint() - .to_u32() - .unwrap_or(self.elements.len() as u32), - None => self.elements.len() as u32, - }; + // len(self)>=1, from="", len(to)>=1, maxcount>=1 + fn replace_interleave(&self, to: PyByteInner, maxcount: Option) -> Vec { + let place_count = self.elements.len() + 1; + let count = maxcount.map_or(place_count, |v| std::cmp::min(v, place_count)) - 1; + let capacity = self.elements.len() + count * to.len(); + let mut result = Vec::with_capacity(capacity); + let to_slice = to.elements.as_slice(); + result.extend_from_slice(to_slice); + for c in &self.elements[..count] { + result.push(*c); + result.extend_from_slice(to_slice); + } + result.extend_from_slice(&self.elements[count..]); + result + } - let mut res = vec![]; - let mut index = 0; - let mut done = 0; - - let slice = &self.elements; - loop { - if done == count || index > slice.len() - old.len() { - res.extend_from_slice(&slice[index..]); - break; - } - if &slice[index..index + old.len()] == old.elements.as_slice() { - res.extend_from_slice(&new.elements); - index += old.len(); - done += 1; - } else { - res.push(slice[index]); - index += 1 - } + fn replace_delete(&self, from: PyByteInner, maxcount: Option) -> Vec { + let count = count_substring(self.elements.as_slice(), from.elements.as_slice(), maxcount); + if count == 0 { + // no matches + return self.elements.clone(); } - Ok(res) + let result_len = self.len() - (count * from.len()); + debug_assert!(self.len() >= count * from.len()); + + let mut result = Vec::with_capacity(result_len); + let mut last_end = 0; + let mut count = count; + for offset in self.elements.find_iter(&from.elements) { + result.extend_from_slice(&self.elements[last_end..offset]); + last_end = offset + from.len(); + count -= 1; + if count == 0 { + break; + } + } + result.extend_from_slice(&self.elements[last_end..]); + result + } + + pub fn replace_in_place( + &self, + from: PyByteInner, + to: PyByteInner, + maxcount: Option, + ) -> Vec { + let len = from.len(); + let mut iter = self.elements.find_iter(&from.elements); + + let mut new = if let Some(offset) = iter.next() { + let mut new = self.elements.clone(); + new[offset..offset + len].clone_from_slice(to.elements.as_slice()); + if maxcount == Some(1) { + return new; + } else { + new + } + } else { + return self.elements.clone(); + }; + + let mut count = maxcount.unwrap_or(std::usize::MAX) - 1; + for offset in iter { + new[offset..offset + len].clone_from_slice(to.elements.as_slice()); + count -= 1; + if count == 0 { + break; + } + } + new + } + + fn replace_general( + &self, + from: PyByteInner, + to: PyByteInner, + maxcount: Option, + vm: &VirtualMachine, + ) -> PyResult> { + let count = count_substring(self.elements.as_slice(), from.elements.as_slice(), maxcount); + if count == 0 { + // no matches, return unchanged + return Ok(self.elements.clone()); + } + + // Check for overflow + // result_len = self_len + count * (to_len-from_len) + debug_assert!(count > 0); + if to.len() as isize - from.len() as isize + > (std::isize::MAX - self.elements.len() as isize) / count as isize + { + return Err(vm.new_overflow_error("replace bytes is too long".to_owned())); + } + let result_len = self.elements.len() + count * (to.len() - from.len()); + + let mut result = Vec::with_capacity(result_len); + let mut last_end = 0; + let mut count = count; + for offset in self.elements.find_iter(&from.elements) { + result.extend_from_slice(&self.elements[last_end..offset]); + result.extend_from_slice(to.elements.as_slice()); + last_end = offset + from.len(); + count -= 1; + if count == 0 { + break; + } + } + result.extend_from_slice(&self.elements[last_end..]); + Ok(result) + } + + pub fn replace( + &self, + from: PyByteInner, + to: PyByteInner, + maxcount: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult> { + // stringlib_replace in CPython + let maxcount = match maxcount { + OptionalArg::Present(maxcount) if maxcount >= 0 => { + if maxcount == 0 || self.elements.is_empty() { + // nothing to do; return the original bytes + return Ok(self.elements.clone()); + } + Some(maxcount as usize) + } + _ => None, + }; + + // Handle zero-length special cases + if from.elements.is_empty() { + if to.elements.is_empty() { + // nothing to do; return the original bytes + return Ok(self.elements.clone()); + } + // insert the 'to' bytes everywhere. + // >>> b"Python".replace(b"", b".") + // b'.P.y.t.h.o.n.' + return Ok(self.replace_interleave(to, maxcount)); + } + + // Except for b"".replace(b"", b"A") == b"A" there is no way beyond this + // point for an empty self bytes to generate a non-empty bytes + // Special case so the remaining code always gets a non-empty bytes + if self.elements.is_empty() { + return Ok(self.elements.clone()); + } + + if to.elements.is_empty() { + // delete all occurrences of 'from' bytes + Ok(self.replace_delete(from, maxcount)) + } else if from.len() == to.len() { + // Handle special case where both bytes have the same length + Ok(self.replace_in_place(from, to, maxcount)) + } else { + // Otherwise use the more generic algorithms + self.replace_general(from, to, maxcount, vm) + } } pub fn title(&self) -> Vec { @@ -1233,6 +1308,16 @@ pub fn try_as_byte(obj: &PyObjectRef) -> Option> { }) } +#[inline] +fn count_substring(haystack: &[u8], needle: &[u8], maxcount: Option) -> usize { + let substrings = haystack.find_iter(needle); + if let Some(maxcount) = maxcount { + std::cmp::min(substrings.take(maxcount).count(), maxcount) + } else { + substrings.count() + } +} + pub trait ByteOr: ToPrimitive { fn byte_or(&self, vm: &VirtualMachine) -> PyResult { match self.to_u8() { @@ -1297,13 +1382,23 @@ pub fn bytes_zfill(bytes: &[u8], width: usize) -> Vec { } } +impl PyCommonStringWrapper<[u8]> for PyByteInner { + fn as_ref(&self) -> &[u8] { + &self.elements + } +} + const ASCII_WHITESPACES: [u8; 6] = [0x20, 0x09, 0x0a, 0x0c, 0x0d, 0x0b]; -impl PyCommonString<'_, u8> for [u8] { +impl PyCommonString for [u8] { fn get_slice(&self, range: std::ops::Range) -> &Self { &self[range] } + fn is_empty(&self) -> bool { + Self::is_empty(self) + } + fn len(&self) -> usize { Self::len(self) } diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs index ffd18f545..83a26d557 100644 --- a/vm/src/obj/objbytes.rs +++ b/vm/src/obj/objbytes.rs @@ -4,15 +4,15 @@ use std::ops::Deref; use std::str::FromStr; use super::objbyteinner::{ - ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions, - ByteInnerSplitOptions, ByteInnerSplitlinesOptions, ByteInnerTranslateOptions, PyByteInner, + ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions, ByteInnerSplitOptions, + ByteInnerTranslateOptions, PyByteInner, }; use super::objint::PyIntRef; use super::objiter; use super::objslice::PySliceRef; use super::objstr::{PyString, PyStringRef}; use super::objtype::PyClassRef; -use super::pystr::PyCommonString; +use super::pystr::{self, PyCommonString}; use crate::cformat::CFormatString; use crate::function::{OptionalArg, OptionalOption}; use crate::obj::objstr::do_cformat_string; @@ -273,17 +273,9 @@ impl PyBytes { } #[pymethod(name = "endswith")] - fn endswith( - &self, - suffix: PyObjectRef, - start: OptionalArg>, - end: OptionalArg>, - vm: &VirtualMachine, - ) -> PyResult { + fn endswith(&self, options: pystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult { self.inner.elements[..].py_startsendswith( - suffix, - start, - end, + options, "endswith", "bytes", |s, x: &PyByteInner| s.ends_with(&x.elements[..]), @@ -294,15 +286,11 @@ impl PyBytes { #[pymethod(name = "startswith")] fn startswith( &self, - prefix: PyObjectRef, - start: OptionalArg>, - end: OptionalArg>, + options: pystr::StartsEndsWithArgs, vm: &VirtualMachine, ) -> PyResult { self.inner.elements[..].py_startsendswith( - prefix, - start, - end, + options, "startswith", "bytes", |s, x: &PyByteInner| s.starts_with(&x.elements[..]), @@ -401,12 +389,12 @@ impl PyBytes { } #[pymethod(name = "expandtabs")] - fn expandtabs(&self, options: ByteInnerExpandtabsOptions) -> PyBytes { + fn expandtabs(&self, options: pystr::ExpandTabsArgs) -> PyBytes { self.inner.expandtabs(options).into() } #[pymethod(name = "splitlines")] - fn splitlines(&self, options: ByteInnerSplitlinesOptions, vm: &VirtualMachine) -> PyResult { + fn splitlines(&self, options: pystr::SplitLinesArgs, vm: &VirtualMachine) -> PyResult { let as_bytes = self .inner .splitlines(options) @@ -426,9 +414,10 @@ impl PyBytes { &self, old: PyByteInner, new: PyByteInner, - count: OptionalArg, + count: OptionalArg, + vm: &VirtualMachine, ) -> PyResult { - Ok(self.inner.replace(old, new, count)?.into()) + Ok(self.inner.replace(old, new, count, vm)?.into()) } #[pymethod(name = "title")] diff --git a/vm/src/obj/objclassmethod.rs b/vm/src/obj/objclassmethod.rs index edc579698..84af11004 100644 --- a/vm/src/obj/objclassmethod.rs +++ b/vm/src/obj/objclassmethod.rs @@ -1,7 +1,7 @@ use super::objtype::PyClassRef; use crate::function::OptionalArg; use crate::pyobject::{ - PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, + PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, TypeProtocol, }; use crate::slots::SlotDescriptor; use crate::vm::VirtualMachine; @@ -32,6 +32,7 @@ pub struct PyClassMethod { callable: PyObjectRef, } pub type PyClassMethodRef = PyRef; +impl ThreadSafe for PyClassMethod {} impl PyClassMethod { pub fn new(value: PyObjectRef) -> Self { diff --git a/vm/src/obj/objcode.rs b/vm/src/obj/objcode.rs index ebe620333..5f4a36bf9 100644 --- a/vm/src/obj/objcode.rs +++ b/vm/src/obj/objcode.rs @@ -7,7 +7,9 @@ use std::ops::Deref; use super::objtype::PyClassRef; use crate::bytecode; -use crate::pyobject::{IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::pyobject::{ + IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, +}; use crate::vm::VirtualMachine; pub type PyCodeRef = PyRef; @@ -17,6 +19,8 @@ pub struct PyCode { pub code: bytecode::CodeObject, } +impl ThreadSafe for PyCode {} + impl Deref for PyCode { type Target = bytecode::CodeObject; fn deref(&self) -> &Self::Target { @@ -45,8 +49,7 @@ impl PyValue for PyCode { #[pyimpl] impl PyCodeRef { #[pyslot] - #[allow(clippy::new_ret_no_self)] - fn new(_cls: PyClassRef, vm: &VirtualMachine) -> PyResult { + fn new(_cls: PyClassRef, vm: &VirtualMachine) -> PyResult> { Err(vm.new_type_error("Cannot directly create code object".to_owned())) } diff --git a/vm/src/obj/objdict.rs b/vm/src/obj/objdict.rs index 0573253b9..a5330bdbc 100644 --- a/vm/src/obj/objdict.rs +++ b/vm/src/obj/objdict.rs @@ -1,4 +1,4 @@ -use std::cell::{Cell, RefCell}; +use std::cell::Cell; use std::fmt; use super::objiter; @@ -9,7 +9,7 @@ use crate::exceptions::PyBaseExceptionRef; use crate::function::{KwArgs, OptionalArg, PyFuncArgs}; use crate::pyobject::{ IdProtocol, IntoPyObject, ItemProtocol, PyAttributes, PyClassImpl, PyContext, PyIterable, - PyObjectRef, PyRef, PyResult, PyValue, + PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, }; use crate::vm::{ReprGuard, VirtualMachine}; @@ -20,9 +20,10 @@ pub type DictContentType = dictdatatype::Dict; #[pyclass] #[derive(Default)] pub struct PyDict { - entries: RefCell, + entries: DictContentType, } pub type PyDictRef = PyRef; +impl ThreadSafe for PyDict {} impl fmt::Debug for PyDict { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { @@ -43,7 +44,7 @@ impl PyDictRef { #[pyslot] fn tp_new(class: PyClassRef, _args: PyFuncArgs, vm: &VirtualMachine) -> PyResult { PyDict { - entries: RefCell::new(DictContentType::default()), + entries: DictContentType::default(), } .into_ref_with_type(vm, class) } @@ -59,7 +60,7 @@ impl PyDictRef { } fn merge( - dict: &RefCell, + dict: &DictContentType, dict_obj: OptionalArg, kwargs: KwArgs, vm: &VirtualMachine, @@ -68,13 +69,13 @@ impl PyDictRef { let dicted: Result = dict_obj.clone().downcast(); if let Ok(dict_obj) = dicted { for (key, value) in dict_obj { - dict.borrow_mut().insert(vm, &key, value)?; + dict.insert(vm, &key, value)?; } } else if let Some(keys) = vm.get_method(dict_obj.clone(), "keys") { let keys = objiter::get_iter(vm, &vm.invoke(&keys?, vec![])?)?; while let Some(key) = objiter::get_next_object(vm, &keys)? { let val = dict_obj.get_item(&key, vm)?; - dict.borrow_mut().insert(vm, &key, val)?; + dict.insert(vm, &key, val)?; } } else { let iter = objiter::get_iter(vm, &dict_obj)?; @@ -92,14 +93,13 @@ impl PyDictRef { if objiter::get_next_object(vm, &elem_iter)?.is_some() { return Err(err(vm)); } - dict.borrow_mut().insert(vm, &key, value)?; + dict.insert(vm, &key, value)?; } } } - let mut dict_borrowed = dict.borrow_mut(); for (key, value) in kwargs.into_iter() { - dict_borrowed.insert(vm, &vm.new_str(key), value)?; + dict.insert(vm, &vm.new_str(key), value)?; } Ok(()) } @@ -111,27 +111,26 @@ impl PyDictRef { value: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - let mut dict = DictContentType::default(); + let dict = DictContentType::default(); let value = value.unwrap_or_else(|| vm.ctx.none()); for elem in iterable.iter(vm)? { let elem = elem?; dict.insert(vm, &elem, value.clone())?; } - let entries = RefCell::new(dict); - PyDict { entries }.into_ref_with_type(vm, class) + PyDict { entries: dict }.into_ref_with_type(vm, class) } #[pymethod(magic)] fn bool(self) -> bool { - !self.entries.borrow().is_empty() + !self.entries.is_empty() } fn inner_eq(self, other: &PyDict, vm: &VirtualMachine) -> PyResult { - if other.entries.borrow().len() != self.entries.borrow().len() { + if other.entries.len() != self.entries.len() { return Ok(false); } for (k, v1) in self { - match other.entries.borrow().get(vm, &k)? { + match other.entries.get(vm, &k)? { Some(v2) => { if v1.is(&v2) { continue; @@ -170,12 +169,12 @@ impl PyDictRef { #[pymethod(magic)] fn len(self) -> usize { - self.entries.borrow().len() + self.entries.len() } #[pymethod(magic)] fn sizeof(self) -> usize { - size_of::() + self.entries.borrow().sizeof() + size_of::() + self.entries.sizeof() } #[pymethod(magic)] @@ -197,17 +196,17 @@ impl PyDictRef { #[pymethod(magic)] fn contains(self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult { - self.entries.borrow().contains(vm, &key) + self.entries.contains(vm, &key) } #[pymethod(magic)] fn delitem(self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - self.entries.borrow_mut().delete(vm, &key) + self.entries.delete(vm, &key) } #[pymethod] fn clear(self) { - self.entries.borrow_mut().clear() + self.entries.clear() } #[pymethod(magic)] @@ -243,7 +242,7 @@ impl PyDictRef { value: PyObjectRef, vm: &VirtualMachine, ) -> PyResult<()> { - self.entries.borrow_mut().insert(vm, key, value) + self.entries.insert(vm, key, value) } #[pymethod(magic)] @@ -262,7 +261,7 @@ impl PyDictRef { key: K, vm: &VirtualMachine, ) -> PyResult> { - if let Some(value) = self.entries.borrow().get(vm, key)? { + if let Some(value) = self.entries.get(vm, key)? { return Ok(Some(value)); } @@ -281,7 +280,7 @@ impl PyDictRef { default: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - match self.entries.borrow().get(vm, &key)? { + match self.entries.get(vm, &key)? { Some(value) => Ok(value), None => Ok(default.unwrap_or_else(|| vm.ctx.none())), } @@ -294,12 +293,11 @@ impl PyDictRef { default: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - let mut entries = self.entries.borrow_mut(); - match entries.get(vm, &key)? { + match self.entries.get(vm, &key)? { Some(value) => Ok(value), None => { let set_value = default.unwrap_or_else(|| vm.ctx.none()); - entries.insert(vm, &key, set_value.clone())?; + self.entries.insert(vm, &key, set_value.clone())?; Ok(set_value) } } @@ -329,7 +327,7 @@ impl PyDictRef { default: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - match self.entries.borrow_mut().pop(vm, &key)? { + match self.entries.pop(vm, &key)? { Some(value) => Ok(value), None => match default { OptionalArg::Present(default) => Ok(default), @@ -340,8 +338,7 @@ impl PyDictRef { #[pymethod] fn popitem(self, vm: &VirtualMachine) -> PyResult { - let mut entries = self.entries.borrow_mut(); - if let Some((key, value)) = entries.pop_front() { + if let Some((key, value)) = self.entries.pop_front() { Ok(vm.ctx.new_tuple(vec![key, value])) } else { let err_msg = vm.new_str("popitem(): dictionary is empty".to_owned()); @@ -360,14 +357,13 @@ impl PyDictRef { } pub fn from_attributes(attrs: PyAttributes, vm: &VirtualMachine) -> PyResult { - let mut dict = DictContentType::default(); + let dict = DictContentType::default(); for (key, value) in attrs { dict.insert(vm, &vm.ctx.new_str(key), value)?; } - let entries = RefCell::new(dict); - Ok(PyDict { entries }.into_ref(vm)) + Ok(PyDict { entries: dict }.into_ref(vm)) } #[pymethod(magic)] @@ -377,11 +373,11 @@ impl PyDictRef { pub fn contains_key(&self, key: T, vm: &VirtualMachine) -> bool { let key = key.into_pyobject(vm).unwrap(); - self.entries.borrow().contains(vm, &key).unwrap() + self.entries.contains(vm, &key).unwrap() } pub fn size(&self) -> dictdatatype::DictSize { - self.entries.borrow().size() + self.entries.size() } /// This function can be used to get an item without raising the @@ -487,8 +483,8 @@ impl Iterator for DictIter { type Item = (PyObjectRef, PyObjectRef); fn next(&mut self) -> Option { - match self.dict.entries.borrow().next_entry(&mut self.position) { - Some((key, value)) => Some((key.clone(), value.clone())), + match self.dict.entries.next_entry(&mut self.position) { + Some((key, value)) => Some((key, value)), None => None, } } @@ -524,7 +520,7 @@ macro_rules! dict_iterator { let s = if let Some(_guard) = ReprGuard::enter(zelf.as_object()) { let mut str_parts = vec![]; for (key, value) in zelf.dict.clone() { - let s = vm.to_repr(&$result_fn(vm, &key, &value))?; + let s = vm.to_repr(&$result_fn(vm, key, value))?; str_parts.push(s.as_str().to_owned()); } format!("{}([{}])", $class_name, str_parts.join(", ")) @@ -563,13 +559,12 @@ macro_rules! dict_iterator { #[allow(clippy::redundant_closure_call)] fn next(&self, vm: &VirtualMachine) -> PyResult { let mut position = self.position.get(); - let dict = self.dict.entries.borrow(); - if dict.has_changed_size(&self.size) { + if self.dict.entries.has_changed_size(&self.size) { return Err( vm.new_runtime_error("dictionary changed size during iteration".to_owned()) ); } - match dict.next_entry(&mut position) { + match self.dict.entries.next_entry(&mut position) { Some((key, value)) => { self.position.set(position); Ok($result_fn(vm, key, value)) @@ -585,10 +580,7 @@ macro_rules! dict_iterator { #[pymethod(name = "__length_hint__")] fn length_hint(&self) -> usize { - self.dict - .entries - .borrow() - .len_from_entry_index(self.position.get()) + self.dict.entries.len_from_entry_index(self.position.get()) } } @@ -607,7 +599,7 @@ dict_iterator! { dictkeyiterator_type, "dict_keys", "dictkeyiterator", - |_vm: &VirtualMachine, key: &PyObjectRef, _value: &PyObjectRef| key.clone() + |_vm: &VirtualMachine, key: PyObjectRef, _value: PyObjectRef| key } dict_iterator! { @@ -617,7 +609,7 @@ dict_iterator! { dictvalueiterator_type, "dict_values", "dictvalueiterator", - |_vm: &VirtualMachine, _key: &PyObjectRef, value: &PyObjectRef| value.clone() + |_vm: &VirtualMachine, _key: PyObjectRef, value: PyObjectRef| value } dict_iterator! { @@ -627,8 +619,8 @@ dict_iterator! { dictitemiterator_type, "dict_items", "dictitemiterator", - |vm: &VirtualMachine, key: &PyObjectRef, value: &PyObjectRef| - vm.ctx.new_tuple(vec![key.clone(), value.clone()]) + |vm: &VirtualMachine, key: PyObjectRef, value: PyObjectRef| + vm.ctx.new_tuple(vec![key, value]) } pub(crate) fn init(context: &PyContext) { diff --git a/vm/src/obj/objenumerate.rs b/vm/src/obj/objenumerate.rs index 4bbe8cacc..3c6c864c2 100644 --- a/vm/src/obj/objenumerate.rs +++ b/vm/src/obj/objenumerate.rs @@ -1,5 +1,5 @@ -use std::cell::RefCell; use std::ops::AddAssign; +use std::sync::RwLock; use num_bigint::BigInt; use num_traits::Zero; @@ -8,16 +8,17 @@ use super::objint::PyIntRef; use super::objiter; use super::objtype::PyClassRef; use crate::function::OptionalArg; -use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe}; use crate::vm::VirtualMachine; #[pyclass] #[derive(Debug)] pub struct PyEnumerate { - counter: RefCell, + counter: RwLock, iterator: PyObjectRef, } type PyEnumerateRef = PyRef; +impl ThreadSafe for PyEnumerate {} impl PyValue for PyEnumerate { fn class(vm: &VirtualMachine) -> PyClassRef { @@ -41,24 +42,19 @@ impl PyEnumerate { let iterator = objiter::get_iter(vm, &iterable)?; PyEnumerate { - counter: RefCell::new(counter), + counter: RwLock::new(counter), iterator, } .into_ref_with_type(vm, cls) } #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let iterator = &self.iterator; - let counter = &self.counter; - let next_obj = objiter::call_next(vm, iterator)?; - let result = vm - .ctx - .new_tuple(vec![vm.ctx.new_bigint(&counter.borrow()), next_obj]); - - AddAssign::add_assign(&mut counter.borrow_mut() as &mut BigInt, 1); - - Ok(result) + fn next(&self, vm: &VirtualMachine) -> PyResult<(BigInt, PyObjectRef)> { + let next_obj = objiter::call_next(vm, &self.iterator)?; + let mut counter = self.counter.write().unwrap(); + let position = counter.clone(); + AddAssign::add_assign(&mut counter as &mut BigInt, 1); + Ok((position, next_obj)) } #[pymethod(name = "__iter__")] diff --git a/vm/src/obj/objfilter.rs b/vm/src/obj/objfilter.rs index 1b3f2ae65..f5f0e4172 100644 --- a/vm/src/obj/objfilter.rs +++ b/vm/src/obj/objfilter.rs @@ -1,7 +1,9 @@ use super::objbool; use super::objiter; use super::objtype::PyClassRef; -use crate::pyobject::{IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::pyobject::{ + IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, +}; use crate::vm::VirtualMachine; pub type PyFilterRef = PyRef; @@ -16,6 +18,7 @@ pub struct PyFilter { predicate: PyObjectRef, iterator: PyObjectRef, } +impl ThreadSafe for PyFilter {} impl PyValue for PyFilter { fn class(vm: &VirtualMachine) -> PyClassRef { diff --git a/vm/src/obj/objfloat.rs b/vm/src/obj/objfloat.rs index b29e2cbc2..d8e6f25d9 100644 --- a/vm/src/obj/objfloat.rs +++ b/vm/src/obj/objfloat.rs @@ -1,4 +1,3 @@ -use hexf_parse; use num_bigint::{BigInt, ToBigInt}; use num_rational::Ratio; use num_traits::{float::Float, pow, sign::Signed, ToPrimitive, Zero}; diff --git a/vm/src/obj/objint.rs b/vm/src/obj/objint.rs index c589fafa0..a8aca1bbd 100644 --- a/vm/src/obj/objint.rs +++ b/vm/src/obj/objint.rs @@ -120,7 +120,6 @@ impl_try_from_object_int!( (u64, to_u64), ); -#[allow(clippy::collapsible_if)] fn inner_pow(int1: &BigInt, int2: &BigInt, vm: &VirtualMachine) -> PyResult { if int2.is_negative() { let v1 = try_float(int1, vm)?; diff --git a/vm/src/obj/objiter.rs b/vm/src/obj/objiter.rs index 4d48de150..c162aca21 100644 --- a/vm/src/obj/objiter.rs +++ b/vm/src/obj/objiter.rs @@ -10,8 +10,8 @@ use super::objsequence; use super::objtype::{self, PyClassRef}; use crate::exceptions::PyBaseExceptionRef; use crate::pyobject::{ - PyCallable, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, - TypeProtocol, + PyCallable, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, + TryFromObject, TypeProtocol, }; use crate::vm::VirtualMachine; @@ -141,6 +141,7 @@ pub struct PySequenceIterator { pub obj: PyObjectRef, pub reversed: bool, } +impl ThreadSafe for PySequenceIterator {} impl PyValue for PySequenceIterator { fn class(vm: &VirtualMachine) -> PyClassRef { @@ -214,6 +215,7 @@ pub struct PyCallableIterator { sentinel: PyObjectRef, done: AtomicCell, } +impl ThreadSafe for PyCallableIterator {} impl PyValue for PyCallableIterator { fn class(vm: &VirtualMachine) -> PyClassRef { diff --git a/vm/src/obj/objmap.rs b/vm/src/obj/objmap.rs index a90b1c025..9b6e25ada 100644 --- a/vm/src/obj/objmap.rs +++ b/vm/src/obj/objmap.rs @@ -1,7 +1,7 @@ use super::objiter; use super::objtype::PyClassRef; use crate::function::Args; -use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe}; use crate::vm::VirtualMachine; /// map(func, *iterables) --> map object @@ -15,6 +15,7 @@ pub struct PyMap { iterators: Vec, } type PyMapRef = PyRef; +impl ThreadSafe for PyMap {} impl PyValue for PyMap { fn class(vm: &VirtualMachine) -> PyClassRef { diff --git a/vm/src/obj/objmemory.rs b/vm/src/obj/objmemory.rs index c4f9c4776..6ae3507de 100644 --- a/vm/src/obj/objmemory.rs +++ b/vm/src/obj/objmemory.rs @@ -1,7 +1,7 @@ use super::objbyteinner::try_as_byte; use super::objtype::{issubclass, PyClassRef}; use crate::pyobject::{ - PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, + PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, TypeProtocol, }; use crate::stdlib::array::PyArray; use crate::vm::VirtualMachine; @@ -11,6 +11,7 @@ use crate::vm::VirtualMachine; pub struct PyMemoryView { obj_ref: PyObjectRef, } +impl ThreadSafe for PyMemoryView {} pub type PyMemoryViewRef = PyRef; diff --git a/vm/src/obj/objproperty.rs b/vm/src/obj/objproperty.rs index f97c9a8ed..535196860 100644 --- a/vm/src/obj/objproperty.rs +++ b/vm/src/obj/objproperty.rs @@ -1,12 +1,13 @@ /*! Python `property` descriptor class. */ -use std::cell::RefCell; +use std::sync::RwLock; use super::objtype::PyClassRef; use crate::function::OptionalArg; use crate::pyobject::{ - IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, + IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, + TypeProtocol, }; use crate::slots::SlotDescriptor; use crate::vm::VirtualMachine; @@ -49,8 +50,9 @@ pub struct PyProperty { getter: Option, setter: Option, deleter: Option, - doc: RefCell>, + doc: RwLock>, } +impl ThreadSafe for PyProperty {} impl PyValue for PyProperty { fn class(vm: &VirtualMachine) -> PyClassRef { @@ -73,7 +75,6 @@ struct PropertyArgs { } impl SlotDescriptor for PyProperty { - #[allow(clippy::collapsible_if)] fn descr_get( vm: &VirtualMachine, zelf: PyObjectRef, @@ -83,12 +84,10 @@ impl SlotDescriptor for PyProperty { let (zelf, obj) = Self::_unwrap(zelf, obj, vm)?; if vm.is_none(&obj) { Ok(zelf.into_object()) + } else if let Some(getter) = zelf.getter.as_ref() { + vm.invoke(&getter, obj) } else { - if let Some(getter) = zelf.getter.as_ref() { - vm.invoke(&getter, obj) - } else { - Err(vm.new_attribute_error("unreadable attribute".to_string())) - } + Err(vm.new_attribute_error("unreadable attribute".to_string())) } } } @@ -101,7 +100,7 @@ impl PyProperty { getter: args.fget, setter: args.fset, deleter: args.fdel, - doc: RefCell::new(args.doc), + doc: RwLock::new(args.doc), } .into_ref_with_type(vm, cls) } @@ -144,11 +143,11 @@ impl PyProperty { } fn doc_getter(&self) -> Option { - self.doc.borrow().clone() + self.doc.read().unwrap().clone() } fn doc_setter(&self, value: PyObjectRef, vm: &VirtualMachine) { - self.doc.replace(py_none_to_option(vm, &value)); + *self.doc.write().unwrap() = py_none_to_option(vm, &value); } // Python builder functions @@ -163,7 +162,7 @@ impl PyProperty { getter: getter.or_else(|| zelf.getter.clone()), setter: zelf.setter.clone(), deleter: zelf.deleter.clone(), - doc: RefCell::new(None), + doc: RwLock::new(None), } .into_ref_with_type(vm, TypeProtocol::class(&zelf)) } @@ -178,7 +177,7 @@ impl PyProperty { getter: zelf.getter.clone(), setter: setter.or_else(|| zelf.setter.clone()), deleter: zelf.deleter.clone(), - doc: RefCell::new(None), + doc: RwLock::new(None), } .into_ref_with_type(vm, TypeProtocol::class(&zelf)) } @@ -193,7 +192,7 @@ impl PyProperty { getter: zelf.getter.clone(), setter: zelf.setter.clone(), deleter: deleter.or_else(|| zelf.deleter.clone()), - doc: RefCell::new(None), + doc: RwLock::new(None), } .into_ref_with_type(vm, TypeProtocol::class(&zelf)) } diff --git a/vm/src/obj/objrange.rs b/vm/src/obj/objrange.rs index a0dc60d27..68bc41dcd 100644 --- a/vm/src/obj/objrange.rs +++ b/vm/src/obj/objrange.rs @@ -1,5 +1,4 @@ -use std::cell::Cell; - +use crossbeam_utils::atomic::AtomicCell; use num_bigint::{BigInt, Sign}; use num_integer::Integer; use num_traits::{One, Signed, Zero}; @@ -13,7 +12,8 @@ use super::objtype::PyClassRef; use crate::function::{OptionalArg, PyFuncArgs}; use crate::pyhash; use crate::pyobject::{ - PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, + PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, TryFromObject, + TypeProtocol, }; use crate::vm::VirtualMachine; @@ -32,6 +32,7 @@ pub struct PyRange { pub stop: PyIntRef, pub step: PyIntRef, } +impl ThreadSafe for PyRange {} impl PyValue for PyRange { fn class(vm: &VirtualMachine) -> PyClassRef { @@ -171,7 +172,7 @@ impl PyRange { #[pymethod(name = "__iter__")] fn iter(zelf: PyRef) -> PyRangeIterator { PyRangeIterator { - position: Cell::new(0), + position: AtomicCell::new(0), range: zelf, } } @@ -204,7 +205,7 @@ impl PyRange { }; PyRangeIterator { - position: Cell::new(0), + position: AtomicCell::new(0), range: reversed.into_ref(vm), } } @@ -401,9 +402,10 @@ impl PyRange { #[pyclass] #[derive(Debug)] pub struct PyRangeIterator { - position: Cell, + position: AtomicCell, range: PyRangeRef, } +impl ThreadSafe for PyRangeIterator {} impl PyValue for PyRangeIterator { fn class(vm: &VirtualMachine) -> PyClassRef { @@ -417,9 +419,8 @@ type PyRangeIteratorRef = PyRef; impl PyRangeIterator { #[pymethod(name = "__next__")] fn next(&self, vm: &VirtualMachine) -> PyResult { - let position = BigInt::from(self.position.get()); + let position = BigInt::from(self.position.fetch_add(1)); if let Some(int) = self.range.get(&position) { - self.position.set(self.position.get() + 1); Ok(int) } else { Err(objiter::new_stop_iteration(vm)) diff --git a/vm/src/obj/objset.rs b/vm/src/obj/objset.rs index 7e65da207..f46724117 100644 --- a/vm/src/obj/objset.rs +++ b/vm/src/obj/objset.rs @@ -1,8 +1,6 @@ /* * Builtin set type with a sequence of unique items. */ - -use std::cell::RefCell; use std::fmt; use super::objlist::PyListIterator; @@ -11,8 +9,8 @@ use crate::dictdatatype; use crate::function::{Args, OptionalArg}; use crate::pyhash; use crate::pyobject::{ - PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, - TypeProtocol, + PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, + TryFromObject, TypeProtocol, }; use crate::vm::{ReprGuard, VirtualMachine}; @@ -25,9 +23,10 @@ pub type SetContentType = dictdatatype::Dict<()>; #[pyclass] #[derive(Default)] pub struct PySet { - inner: RefCell, + inner: PySetInner, } pub type PySetRef = PyRef; +impl ThreadSafe for PySet {} /// frozenset() -> empty frozenset object /// frozenset(iterable) -> frozenset object @@ -39,6 +38,7 @@ pub struct PyFrozenSet { inner: PySetInner, } pub type PyFrozenSetRef = PyRef; +impl ThreadSafe for PyFrozenSet {} impl fmt::Debug for PySet { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { @@ -73,7 +73,7 @@ struct PySetInner { impl PySetInner { fn new(iterable: PyIterable, vm: &VirtualMachine) -> PyResult { - let mut set = PySetInner::default(); + let set = PySetInner::default(); for item in iterable.iter(vm)? { set.add(&item?, vm)?; } @@ -177,7 +177,7 @@ impl PySetInner { } fn union(&self, other: PyIterable, vm: &VirtualMachine) -> PyResult { - let mut set = self.clone(); + let set = self.clone(); for item in other.iter(vm)? { set.add(&item?, vm)?; } @@ -186,7 +186,7 @@ impl PySetInner { } fn intersection(&self, other: PyIterable, vm: &VirtualMachine) -> PyResult { - let mut set = PySetInner::default(); + let set = PySetInner::default(); for item in other.iter(vm)? { let obj = item?; if self.contains(&obj, vm)? { @@ -197,7 +197,7 @@ impl PySetInner { } fn difference(&self, other: PyIterable, vm: &VirtualMachine) -> PyResult { - let mut set = self.copy(); + let set = self.copy(); for item in other.iter(vm)? { set.content.delete_if_exists(vm, &item?)?; } @@ -205,7 +205,7 @@ impl PySetInner { } fn symmetric_difference(&self, other: PyIterable, vm: &VirtualMachine) -> PyResult { - let mut new_inner = self.clone(); + let new_inner = self.clone(); // We want to remove duplicates in other let other_set = Self::new(other, vm)?; @@ -241,8 +241,7 @@ impl PySetInner { } fn iter(&self, vm: &VirtualMachine) -> PyListIterator { - let items = self.content.keys().collect(); - let set_list = vm.ctx.new_list(items); + let set_list = vm.ctx.new_list(self.content.keys()); PyListIterator { position: crossbeam_utils::atomic::AtomicCell::new(0), list: set_list.downcast().unwrap(), @@ -259,23 +258,23 @@ impl PySetInner { Ok(format!("{{{}}}", str_parts.join(", "))) } - fn add(&mut self, item: &PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + fn add(&self, item: &PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { self.content.insert(vm, item, ()) } - fn remove(&mut self, item: &PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + fn remove(&self, item: &PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { self.content.delete(vm, item) } - fn discard(&mut self, item: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + fn discard(&self, item: &PyObjectRef, vm: &VirtualMachine) -> PyResult { self.content.delete_if_exists(vm, item) } - fn clear(&mut self) { + fn clear(&self) { self.content.clear() } - fn pop(&mut self, vm: &VirtualMachine) -> PyResult { + fn pop(&self, vm: &VirtualMachine) -> PyResult { if let Some((key, _)) = self.content.pop_front() { Ok(key) } else { @@ -284,7 +283,7 @@ impl PySetInner { } } - fn update(&mut self, others: Args, vm: &VirtualMachine) -> PyResult<()> { + fn update(&self, others: Args, vm: &VirtualMachine) -> PyResult<()> { for iterable in others { for item in iterable.iter(vm)? { self.add(&item?, vm)?; @@ -293,11 +292,7 @@ impl PySetInner { Ok(()) } - fn intersection_update( - &mut self, - others: Args, - vm: &VirtualMachine, - ) -> PyResult<()> { + fn intersection_update(&self, others: Args, vm: &VirtualMachine) -> PyResult<()> { let mut temp_inner = self.copy(); self.clear(); for iterable in others { @@ -312,7 +307,7 @@ impl PySetInner { Ok(()) } - fn difference_update(&mut self, others: Args, vm: &VirtualMachine) -> PyResult<()> { + fn difference_update(&self, others: Args, vm: &VirtualMachine) -> PyResult<()> { for iterable in others { for item in iterable.iter(vm)? { self.content.delete_if_exists(vm, &item?)?; @@ -322,7 +317,7 @@ impl PySetInner { } fn symmetric_difference_update( - &mut self, + &self, others: Args, vm: &VirtualMachine, ) -> PyResult<()> { @@ -337,14 +332,14 @@ impl PySetInner { } fn hash(&self, vm: &VirtualMachine) -> PyResult { - pyhash::hash_iter_unordered(self.content.keys(), vm) + pyhash::hash_iter_unordered(self.content.keys().iter(), vm) } } macro_rules! try_set_cmp { ($vm:expr, $other:expr, $op:expr) => { Ok(match_class!(match ($other) { - set @ PySet => ($vm.new_bool($op(&*set.inner.borrow())?)), + set @ PySet => ($vm.new_bool($op(&set.inner)?)), frozen @ PyFrozenSet => ($vm.new_bool($op(&frozen.inner)?)), _ => $vm.ctx.not_implemented(), })); @@ -353,13 +348,11 @@ macro_rules! try_set_cmp { macro_rules! multi_args_set { ($vm:expr, $others:expr, $zelf:expr, $op:tt) => {{ - let mut res = $zelf.inner.borrow().copy(); + let mut res = $zelf.inner.copy(); for other in $others { res = res.$op(other, $vm)? } - Ok(Self { - inner: RefCell::new(res), - }) + Ok(Self { inner: res }) }}; } @@ -372,61 +365,61 @@ impl PySet { vm: &VirtualMachine, ) -> PyResult { Self { - inner: RefCell::new(PySetInner::from_arg(iterable, vm)?), + inner: PySetInner::from_arg(iterable, vm)?, } .into_ref_with_type(vm, cls) } #[pymethod(name = "__len__")] fn len(&self) -> usize { - self.inner.borrow().len() + self.inner.len() } #[pymethod(name = "__sizeof__")] fn sizeof(&self) -> usize { - std::mem::size_of::() + self.inner.borrow().sizeof() + std::mem::size_of::() + self.inner.sizeof() } #[pymethod] fn copy(&self) -> Self { Self { - inner: RefCell::new(self.inner.borrow().copy()), + inner: self.inner.copy(), } } #[pymethod(name = "__contains__")] fn contains(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { - self.inner.borrow().contains(&needle, vm) + self.inner.contains(&needle, vm) } #[pymethod(name = "__eq__")] fn eq(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { - try_set_cmp!(vm, other, |other| self.inner.borrow().eq(other, vm)) + try_set_cmp!(vm, other, |other| self.inner.eq(other, vm)) } #[pymethod(name = "__ne__")] fn ne(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { - try_set_cmp!(vm, other, |other| self.inner.borrow().ne(other, vm)) + try_set_cmp!(vm, other, |other| self.inner.ne(other, vm)) } #[pymethod(name = "__ge__")] fn ge(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { - try_set_cmp!(vm, other, |other| self.inner.borrow().ge(other, vm)) + try_set_cmp!(vm, other, |other| self.inner.ge(other, vm)) } #[pymethod(name = "__gt__")] fn gt(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { - try_set_cmp!(vm, other, |other| self.inner.borrow().gt(other, vm)) + try_set_cmp!(vm, other, |other| self.inner.gt(other, vm)) } #[pymethod(name = "__le__")] fn le(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { - try_set_cmp!(vm, other, |other| self.inner.borrow().le(other, vm)) + try_set_cmp!(vm, other, |other| self.inner.le(other, vm)) } #[pymethod(name = "__lt__")] fn lt(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { - try_set_cmp!(vm, other, |other| self.inner.borrow().lt(other, vm)) + try_set_cmp!(vm, other, |other| self.inner.lt(other, vm)) } #[pymethod] @@ -455,17 +448,17 @@ impl PySet { #[pymethod] fn issubset(&self, other: PyIterable, vm: &VirtualMachine) -> PyResult { - self.inner.borrow().issubset(other, vm) + self.inner.issubset(other, vm) } #[pymethod] fn issuperset(&self, other: PyIterable, vm: &VirtualMachine) -> PyResult { - self.inner.borrow().issuperset(other, vm) + self.inner.issuperset(other, vm) } #[pymethod] fn isdisjoint(&self, other: PyIterable, vm: &VirtualMachine) -> PyResult { - self.inner.borrow().isdisjoint(other, vm) + self.inner.isdisjoint(other, vm) } #[pymethod(name = "__or__")] @@ -510,16 +503,15 @@ impl PySet { #[pymethod(name = "__iter__")] fn iter(&self, vm: &VirtualMachine) -> PyListIterator { - self.inner.borrow().iter(vm) + self.inner.iter(vm) } #[pymethod(name = "__repr__")] fn repr(zelf: PyRef, vm: &VirtualMachine) -> PyResult { - let inner = zelf.inner.borrow(); - let s = if inner.len() == 0 { + let s = if zelf.inner.len() == 0 { "set()".to_owned() } else if let Some(_guard) = ReprGuard::enter(zelf.as_object()) { - inner.repr(vm)? + zelf.inner.repr(vm)? } else { "set(...)".to_owned() }; @@ -528,68 +520,64 @@ impl PySet { #[pymethod] pub fn add(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - self.inner.borrow_mut().add(&item, vm)?; + self.inner.add(&item, vm)?; Ok(()) } #[pymethod] fn remove(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - self.inner.borrow_mut().remove(&item, vm) + self.inner.remove(&item, vm) } #[pymethod] fn discard(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - self.inner.borrow_mut().discard(&item, vm)?; + self.inner.discard(&item, vm)?; Ok(()) } #[pymethod] fn clear(&self) { - self.inner.borrow_mut().clear() + self.inner.clear() } #[pymethod] fn pop(&self, vm: &VirtualMachine) -> PyResult { - self.inner.borrow_mut().pop(vm) + self.inner.pop(vm) } #[pymethod(name = "__ior__")] fn ior(zelf: PyRef, iterable: SetIterable, vm: &VirtualMachine) -> PyResult { - zelf.inner.borrow_mut().update(iterable.iterable, vm)?; + zelf.inner.update(iterable.iterable, vm)?; Ok(zelf.as_object().clone()) } #[pymethod] fn update(&self, others: Args, vm: &VirtualMachine) -> PyResult { - self.inner.borrow_mut().update(others, vm)?; + self.inner.update(others, vm)?; Ok(vm.get_none()) } #[pymethod] fn intersection_update(&self, others: Args, vm: &VirtualMachine) -> PyResult { - self.inner.borrow_mut().intersection_update(others, vm)?; + self.inner.intersection_update(others, vm)?; Ok(vm.get_none()) } #[pymethod(name = "__iand__")] fn iand(zelf: PyRef, iterable: SetIterable, vm: &VirtualMachine) -> PyResult { - zelf.inner - .borrow_mut() - .intersection_update(iterable.iterable, vm)?; + zelf.inner.intersection_update(iterable.iterable, vm)?; Ok(zelf.as_object().clone()) } #[pymethod] fn difference_update(&self, others: Args, vm: &VirtualMachine) -> PyResult { - self.inner.borrow_mut().difference_update(others, vm)?; + self.inner.difference_update(others, vm)?; Ok(vm.get_none()) } #[pymethod(name = "__isub__")] fn isub(zelf: PyRef, iterable: SetIterable, vm: &VirtualMachine) -> PyResult { - zelf.inner - .borrow_mut() - .difference_update(iterable.iterable, vm)?; + zelf.inner.difference_update(iterable.iterable, vm)?; Ok(zelf.as_object().clone()) } @@ -599,16 +587,13 @@ impl PySet { others: Args, vm: &VirtualMachine, ) -> PyResult { - self.inner - .borrow_mut() - .symmetric_difference_update(others, vm)?; + self.inner.symmetric_difference_update(others, vm)?; Ok(vm.get_none()) } #[pymethod(name = "__ixor__")] fn ixor(zelf: PyRef, iterable: SetIterable, vm: &VirtualMachine) -> PyResult { zelf.inner - .borrow_mut() .symmetric_difference_update(iterable.iterable, vm)?; Ok(zelf.as_object().clone()) } @@ -635,7 +620,7 @@ impl PyFrozenSet { vm: &VirtualMachine, it: impl IntoIterator, ) -> PyResult { - let mut inner = PySetInner::default(); + let inner = PySetInner::default(); for elem in it { inner.add(&elem, vm)?; } diff --git a/vm/src/obj/objslice.rs b/vm/src/obj/objslice.rs index ff11cc17d..198b740ca 100644 --- a/vm/src/obj/objslice.rs +++ b/vm/src/obj/objslice.rs @@ -2,7 +2,8 @@ use super::objint::PyInt; use super::objtype::PyClassRef; use crate::function::{OptionalArg, PyFuncArgs}; use crate::pyobject::{ - IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryIntoRef, + IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, + TryIntoRef, }; use crate::vm::VirtualMachine; use num_bigint::{BigInt, ToBigInt}; @@ -15,6 +16,7 @@ pub struct PySlice { pub stop: PyObjectRef, pub step: Option, } +impl ThreadSafe for PySlice {} impl PyValue for PySlice { fn class(vm: &VirtualMachine) -> PyClassRef { diff --git a/vm/src/obj/objstaticmethod.rs b/vm/src/obj/objstaticmethod.rs index f78842b2c..7b3f1241a 100644 --- a/vm/src/obj/objstaticmethod.rs +++ b/vm/src/obj/objstaticmethod.rs @@ -1,6 +1,6 @@ use super::objtype::PyClassRef; use crate::function::OptionalArg; -use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue}; +use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe}; use crate::slots::SlotDescriptor; use crate::vm::VirtualMachine; @@ -10,6 +10,7 @@ pub struct PyStaticMethod { pub callable: PyObjectRef, } pub type PyStaticMethodRef = PyRef; +impl ThreadSafe for PyStaticMethod {} impl PyValue for PyStaticMethod { fn class(vm: &VirtualMachine) -> PyClassRef { diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs index a53603cc3..936673797 100644 --- a/vm/src/obj/objstr.rs +++ b/vm/src/obj/objstr.rs @@ -21,17 +21,17 @@ use super::objsequence::PySliceableSequence; use super::objslice::PySliceRef; use super::objtuple; use super::objtype::{self, PyClassRef}; -use super::pystr::{adjust_indices, PyCommonString, StringRange}; +use super::pystr::{self, adjust_indices, PyCommonString, PyCommonStringWrapper, StringRange}; use crate::cformat::{ CFormatPart, CFormatPreconversor, CFormatQuantity, CFormatSpec, CFormatString, CFormatType, CNumberType, }; use crate::format::{FormatParseError, FormatPart, FormatPreconversor, FormatSpec, FormatString}; -use crate::function::{OptionalArg, PyFuncArgs}; +use crate::function::{OptionalArg, OptionalOption, PyFuncArgs}; use crate::pyhash; use crate::pyobject::{ Either, IdProtocol, IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyIterable, - PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TryIntoRef, TypeProtocol, + PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, TryFromObject, TryIntoRef, TypeProtocol, }; use crate::vm::VirtualMachine; @@ -51,6 +51,7 @@ pub struct PyString { value: String, hash: AtomicCell>, } +impl ThreadSafe for PyString {} impl PyString { #[inline] @@ -100,6 +101,7 @@ pub struct PyStringIterator { pub string: PyStringRef, position: AtomicCell, } +impl ThreadSafe for PyStringIterator {} impl PyValue for PyStringIterator { fn class(vm: &VirtualMachine) -> PyClassRef { @@ -133,6 +135,7 @@ pub struct PyStringReverseIterator { pub position: AtomicCell, pub string: PyStringRef, } +impl ThreadSafe for PyStringReverseIterator {} impl PyValue for PyStringReverseIterator { fn class(vm: &VirtualMachine) -> PyClassRef { @@ -144,12 +147,13 @@ impl PyValue for PyStringReverseIterator { impl PyStringReverseIterator { #[pymethod(name = "__next__")] fn next(&self, vm: &VirtualMachine) -> PyResult { - let pos = self.position.fetch_sub(1) as usize - 1; - if let Some(c) = self.string.value.chars().nth(pos) { - Ok(c.to_string()) - } else { - Err(objiter::new_stop_iteration(vm)) + let pos = self.position.fetch_sub(1); + if pos > 0 { + if let Some(c) = self.string.value.chars().nth(pos as usize - 1) { + return Ok(c.to_string()); + } } + Err(objiter::new_stop_iteration(vm)) } #[pymethod(name = "__iter__")] @@ -168,18 +172,6 @@ struct StrArgs { errors: OptionalArg, } -#[derive(FromArgs)] -struct SplitLineArgs { - #[pyarg(positional_or_keyword, optional = true)] - keepends: OptionalArg, -} - -#[derive(FromArgs)] -struct ExpandtabsArgs { - #[pyarg(positional_or_keyword, optional = true)] - tabsize: OptionalArg, -} - #[pyimpl(flags(BASETYPE))] impl PyString { #[pyslot] @@ -457,26 +449,24 @@ impl PyString { #[pymethod] fn split(&self, args: SplitArgs, vm: &VirtualMachine) -> PyResult { let elements = self.value.py_split( - args.non_empty_sep(vm)?, - args.maxsplit, + args, vm, |v, s, vm| v.split(s).map(|s| vm.ctx.new_str(s)).collect(), |v, s, n, vm| v.splitn(n, s).map(|s| vm.ctx.new_str(s)).collect(), |v, n, vm| v.py_split_whitespace(n, |s| vm.ctx.new_str(s)), - ); + )?; Ok(vm.ctx.new_list(elements)) } #[pymethod] fn rsplit(&self, args: SplitArgs, vm: &VirtualMachine) -> PyResult { let mut elements = self.value.py_split( - args.non_empty_sep(vm)?, - args.maxsplit, + args, vm, |v, s, vm| v.rsplit(s).map(|s| vm.ctx.new_str(s)).collect(), |v, s, n, vm| v.rsplitn(n, s).map(|s| vm.ctx.new_str(s)).collect(), |v, n, vm| v.py_rsplit_whitespace(n, |s| vm.ctx.new_str(s)), - ); + )?; // Unlike Python rsplit, Rust rsplitn returns an iterator that // starts from the end of the string. elements.reverse(); @@ -484,51 +474,42 @@ impl PyString { } #[pymethod] - fn strip(&self, chars: OptionalArg>) -> String { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.value, - None => return self.value.trim().to_owned(), - }; - self.value.trim_matches(|c| chars.contains(c)).to_owned() - } - - #[pymethod] - fn lstrip(&self, chars: OptionalArg>) -> String { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.value, - None => return self.value.trim_start().to_owned(), - }; + fn strip(&self, chars: OptionalOption) -> String { self.value - .trim_start_matches(|c| chars.contains(c)) + .py_strip( + chars, + |s, chars| s.trim_matches(|c| chars.contains(c)), + |s| s.trim(), + ) .to_owned() } #[pymethod] - fn rstrip(&self, chars: OptionalArg>) -> String { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.value, - None => return self.value.trim_end().to_owned(), - }; + fn lstrip(&self, chars: OptionalOption) -> String { self.value - .trim_end_matches(|c| chars.contains(c)) + .py_strip( + chars, + |s, chars| s.trim_start_matches(|c| chars.contains(c)), + |s| s.trim_start(), + ) .to_owned() } #[pymethod] - fn endswith( - &self, - suffix: PyObjectRef, - start: OptionalArg>, - end: OptionalArg>, - vm: &VirtualMachine, - ) -> PyResult { + fn rstrip(&self, chars: OptionalOption) -> String { + self.value + .py_strip( + chars, + |s, chars| s.trim_end_matches(|c| chars.contains(c)), + |s| s.trim_end(), + ) + .to_owned() + } + + #[pymethod] + fn endswith(&self, args: pystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult { self.value.as_str().py_startsendswith( - suffix, - start, - end, + args, "endswith", "str", |s, x: &PyStringRef| s.ends_with(x.as_str()), @@ -537,17 +518,9 @@ impl PyString { } #[pymethod] - fn startswith( - &self, - prefix: PyObjectRef, - start: OptionalArg>, - end: OptionalArg>, - vm: &VirtualMachine, - ) -> PyResult { + fn startswith(&self, args: pystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult { self.value.as_str().py_startsendswith( - prefix, - start, - end, + args, "startswith", "str", |s, x: &PyStringRef| s.starts_with(x.as_str()), @@ -799,14 +772,13 @@ impl PyString { } #[pymethod] - fn splitlines(&self, args: SplitLineArgs, vm: &VirtualMachine) -> PyObjectRef { - let keepends = args.keepends.unwrap_or(false); + fn splitlines(&self, args: pystr::SplitLinesArgs, vm: &VirtualMachine) -> PyObjectRef { let mut elements = vec![]; let mut curr = "".to_owned(); let mut chars = self.value.chars().peekable(); while let Some(ch) = chars.next() { if ch == '\n' || ch == '\r' { - if keepends { + if args.keepends { curr.push(ch); } if ch == '\r' && chars.peek() == Some(&'\n') { @@ -1082,8 +1054,8 @@ impl PyString { } #[pymethod] - fn expandtabs(&self, args: ExpandtabsArgs) -> String { - let tab_stop = args.tabsize.unwrap_or(8); + fn expandtabs(&self, args: pystr::ExpandTabsArgs) -> String { + let tab_stop = args.tabsize(); let mut expanded_str = String::with_capacity(self.value.len()); let mut tab_size = tab_stop; let mut col_count = 0 as usize; @@ -1298,28 +1270,7 @@ impl TryFromObject for std::ffi::CString { } } -#[derive(FromArgs)] -struct SplitArgs { - #[pyarg(positional_or_keyword, default = "None")] - sep: Option, - #[pyarg(positional_or_keyword, default = "-1")] - maxsplit: isize, -} - -impl SplitArgs { - fn non_empty_sep<'a>(&'a self, vm: &VirtualMachine) -> PyResult> { - let sep = if let Some(s) = self.sep.as_ref() { - let sep = s.as_str(); - if sep.is_empty() { - return Err(vm.new_value_error("empty separator".to_owned())); - } - Some(sep) - } else { - None - }; - Ok(sep) - } -} +type SplitArgs = pystr::SplitArgs; pub fn init(ctx: &PyContext) { PyString::extend_class(ctx, &ctx.types.str_type); @@ -1799,11 +1750,21 @@ mod tests { } } -impl PyCommonString<'_, char> for str { +impl PyCommonStringWrapper for PyStringRef { + fn as_ref(&self) -> &str { + self.value.as_str() + } +} + +impl PyCommonString for str { fn get_slice(&self, range: std::ops::Range) -> &Self { &self[range] } + fn is_empty(&self) -> bool { + Self::is_empty(self) + } + fn len(&self) -> usize { Self::len(self) } diff --git a/vm/src/obj/objtuple.rs b/vm/src/obj/objtuple.rs index 622da4276..85427a189 100644 --- a/vm/src/obj/objtuple.rs +++ b/vm/src/obj/objtuple.rs @@ -11,7 +11,7 @@ use crate::pyobject::{ PyArithmaticValue::{self, *}, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, PyRef, PyResult, PyValue, ThreadSafe, }; -use crate::sequence::{self, SimpleSeq}; +use crate::sequence; use crate::vm::{ReprGuard, VirtualMachine}; /// tuple() -> empty tuple diff --git a/vm/src/obj/pystr.rs b/vm/src/obj/pystr.rs index ec2677854..c0c0d113e 100644 --- a/vm/src/obj/pystr.rs +++ b/vm/src/obj/pystr.rs @@ -1,6 +1,75 @@ use crate::function::{single_or_tuple_any, OptionalOption}; use crate::pyobject::{PyObjectRef, PyResult, TryFromObject, TypeProtocol}; use crate::vm::VirtualMachine; +use num_traits::cast::ToPrimitive; + +#[derive(FromArgs)] +pub struct SplitArgs +where + T: TryFromObject + PyCommonStringWrapper, + S: ?Sized + PyCommonString, +{ + #[pyarg(positional_or_keyword, default = "None")] + sep: Option, + #[pyarg(positional_or_keyword, default = "-1")] + maxsplit: isize, + _phantom1: std::marker::PhantomData, + _phantom2: std::marker::PhantomData, +} + +impl SplitArgs +where + T: TryFromObject + PyCommonStringWrapper, + S: ?Sized + PyCommonString, +{ + pub fn get_value(self, vm: &VirtualMachine) -> PyResult<(Option, isize)> { + let sep = if let Some(s) = self.sep { + let sep = s.as_ref(); + if sep.is_empty() { + return Err(vm.new_value_error("empty separator".to_owned())); + } + Some(s) + } else { + None + }; + Ok((sep, self.maxsplit)) + } +} + +#[derive(FromArgs)] +pub struct SplitLinesArgs { + #[pyarg(positional_or_keyword, default = "false")] + pub keepends: bool, +} + +#[derive(FromArgs)] +pub struct ExpandTabsArgs { + #[pyarg(positional_or_keyword, default = "8")] + tabsize: isize, +} + +impl ExpandTabsArgs { + pub fn tabsize(&self) -> usize { + self.tabsize.to_usize().unwrap_or(0) + } +} + +#[derive(FromArgs)] +pub struct StartsEndsWithArgs { + #[pyarg(positional_only, optional = false)] + affix: PyObjectRef, + #[pyarg(positional_only, optional = true)] + start: OptionalOption, + #[pyarg(positional_only, optional = true)] + end: OptionalOption, +} + +impl StartsEndsWithArgs { + fn get_value(self, len: usize) -> (PyObjectRef, std::ops::Range) { + let range = adjust_indices(self.start, self.end, len); + (self.affix, range) + } +} // help get optional string indices pub fn adjust_indices( @@ -37,36 +106,43 @@ impl StringRange for std::ops::Range { } } -pub trait PyCommonString<'a, E> +pub trait PyCommonStringWrapper where - Self: 'a, + S: ?Sized, { + fn as_ref(&self) -> &S; +} + +pub trait PyCommonString { fn get_slice(&self, range: std::ops::Range) -> &Self; fn len(&self) -> usize; + fn is_empty(&self) -> bool; - fn py_split( + fn py_split( &self, - sep: Option<&Self>, - maxsplit: isize, + args: SplitArgs, vm: &VirtualMachine, split: SP, splitn: SN, splitw: SW, - ) -> Vec + ) -> PyResult> where + T: TryFromObject + PyCommonStringWrapper, SP: Fn(&Self, &Self, &VirtualMachine) -> Vec, SN: Fn(&Self, &Self, usize, &VirtualMachine) -> Vec, SW: Fn(&Self, isize, &VirtualMachine) -> Vec, { - if let Some(pattern) = sep { + let (sep, maxsplit) = args.get_value(vm)?; + let splited = if let Some(pattern) = sep { if maxsplit < 0 { - split(self, pattern, vm) + split(self, pattern.as_ref(), vm) } else { - splitn(self, pattern, (maxsplit + 1) as usize, vm) + splitn(self, pattern.as_ref(), (maxsplit + 1) as usize, vm) } } else { splitw(self, maxsplit, vm) - } + }; + Ok(splited) } fn py_split_whitespace(&self, maxsplit: isize, convert: F) -> Vec where @@ -75,13 +151,10 @@ where where F: Fn(&Self) -> PyObjectRef; - #[allow(clippy::too_many_arguments)] #[inline] fn py_startsendswith( &self, - affix: PyObjectRef, - start: OptionalOption, - end: OptionalOption, + args: StartsEndsWithArgs, func_name: &str, py_type_name: &str, func: F, @@ -91,7 +164,7 @@ where T: TryFromObject, F: Fn(&Self, &T) -> bool, { - let range = adjust_indices(start, end, self.len()); + let (affix, range) = args.get_value(self.len()); if range.is_normal() { let value = self.get_slice(range); single_or_tuple_any( @@ -112,4 +185,22 @@ where Ok(false) } } + + fn py_strip<'a, S, FC, FD>( + &'a self, + chars: OptionalOption, + func_chars: FC, + func_default: FD, + ) -> &'a Self + where + S: PyCommonStringWrapper, + FC: Fn(&'a Self, &Self) -> &'a Self, + FD: Fn(&'a Self) -> &'a Self, + { + let chars = chars.flat_option(); + match chars { + Some(chars) => func_chars(self, chars.as_ref()), + None => func_default(self), + } + } } diff --git a/vm/src/py_serde.rs b/vm/src/py_serde.rs index 47520bb13..38e0bc933 100644 --- a/vm/src/py_serde.rs +++ b/vm/src/py_serde.rs @@ -1,6 +1,5 @@ -use std::fmt; - -use serde; +use num_traits::cast::ToPrimitive; +use num_traits::sign::Signed; use serde::de::{DeserializeSeed, Visitor}; use serde::ser::{Serialize, SerializeMap, SerializeSeq}; @@ -10,8 +9,6 @@ use crate::obj::{ }; use crate::pyobject::{IdProtocol, ItemProtocol, PyObjectRef, TypeProtocol}; use crate::VirtualMachine; -use num_traits::cast::ToPrimitive; -use num_traits::sign::Signed; #[inline] pub fn serialize( @@ -137,7 +134,7 @@ impl<'de> DeserializeSeed<'de> for PyObjectDeserializer<'de> { impl<'de> Visitor<'de> for PyObjectDeserializer<'de> { type Value = PyObjectRef; - fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { formatter.write_str("a type that can deserialise in Python") } diff --git a/vm/src/pyhash.rs b/vm/src/pyhash.rs index 7c650517b..60a6d132e 100644 --- a/vm/src/pyhash.rs +++ b/vm/src/pyhash.rs @@ -84,13 +84,13 @@ pub fn hash_iter<'a, I: std::iter::Iterator>( Ok(hasher.finish() as PyHash) } -pub fn hash_iter_unordered>( +pub fn hash_iter_unordered<'a, I: std::iter::Iterator>( iter: I, vm: &VirtualMachine, ) -> PyResult { let mut hash: PyHash = 0; for element in iter { - let item_hash = vm._hash(&element)?; + let item_hash = vm._hash(element)?; // xor is commutative and hash should be independent of order hash ^= item_hash; } diff --git a/vm/src/stdlib/binascii.rs b/vm/src/stdlib/binascii.rs index 18db3e7d9..9803463e8 100644 --- a/vm/src/stdlib/binascii.rs +++ b/vm/src/stdlib/binascii.rs @@ -1,147 +1,144 @@ -use crate::function::OptionalArg; -use crate::obj::objbytearray::{PyByteArray, PyByteArrayRef}; -use crate::obj::objbyteinner::PyBytesLike; -use crate::obj::objbytes::{PyBytes, PyBytesRef}; -use crate::obj::objstr::{PyString, PyStringRef}; -use crate::pyobject::{PyObjectRef, PyResult, TryFromObject, TypeProtocol}; -use crate::vm::VirtualMachine; +pub(crate) use decl::make_module; -use crc::{crc32, Hasher32}; -use itertools::Itertools; +#[pymodule(name = "binascii")] +mod decl { + use crate::function::OptionalArg; + use crate::obj::objbytearray::{PyByteArray, PyByteArrayRef}; + use crate::obj::objbyteinner::PyBytesLike; + use crate::obj::objbytes::{PyBytes, PyBytesRef}; + use crate::obj::objstr::{PyString, PyStringRef}; + use crate::pyobject::{PyObjectRef, PyResult, TryFromObject, TypeProtocol}; + use crate::vm::VirtualMachine; + use crc::{crc32, Hasher32}; + use itertools::Itertools; -enum SerializedData { - Bytes(PyBytesRef), - Buffer(PyByteArrayRef), - Ascii(PyStringRef), -} + enum SerializedData { + Bytes(PyBytesRef), + Buffer(PyByteArrayRef), + Ascii(PyStringRef), + } -impl TryFromObject for SerializedData { - fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { - match_class!(match obj { - b @ PyBytes => Ok(SerializedData::Bytes(b)), - b @ PyByteArray => Ok(SerializedData::Buffer(b)), - a @ PyString => { - if a.as_str().is_ascii() { - Ok(SerializedData::Ascii(a)) - } else { - Err(vm.new_value_error( - "string argument should contain only ASCII characters".to_owned(), - )) + impl TryFromObject for SerializedData { + fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { + match_class!(match obj { + b @ PyBytes => Ok(SerializedData::Bytes(b)), + b @ PyByteArray => Ok(SerializedData::Buffer(b)), + a @ PyString => { + if a.as_str().is_ascii() { + Ok(SerializedData::Ascii(a)) + } else { + Err(vm.new_value_error( + "string argument should contain only ASCII characters".to_owned(), + )) + } } + obj => Err(vm.new_type_error(format!( + "argument should be bytes, buffer or ASCII string, not '{}'", + obj.class().name, + ))), + }) + } + } + + impl SerializedData { + #[inline] + pub fn with_ref(&self, f: impl FnOnce(&[u8]) -> R) -> R { + match self { + SerializedData::Bytes(b) => f(b.get_value()), + SerializedData::Buffer(b) => f(&b.borrow_value().elements), + SerializedData::Ascii(a) => f(a.as_str().as_bytes()), } - obj => Err(vm.new_type_error(format!( - "argument should be bytes, buffer or ASCII string, not '{}'", - obj.class().name, - ))), + } + } + + fn hex_nibble(n: u8) -> u8 { + match n { + 0..=9 => b'0' + n, + 10..=15 => b'a' + n, + _ => unreachable!(), + } + } + + #[pyfunction(name = "b2a_hex")] + #[pyfunction] + fn hexlify(data: PyBytesLike) -> Vec { + data.with_ref(|bytes| { + let mut hex = Vec::::with_capacity(bytes.len() * 2); + for b in bytes.iter() { + hex.push(hex_nibble(b >> 4)); + hex.push(hex_nibble(b & 0xf)); + } + hex }) } -} -impl SerializedData { - #[inline] - pub fn with_ref(&self, f: impl FnOnce(&[u8]) -> R) -> R { - match self { - SerializedData::Bytes(b) => f(b.get_value()), - SerializedData::Buffer(b) => f(&b.borrow_value().elements), - SerializedData::Ascii(a) => f(a.as_str().as_bytes()), + fn unhex_nibble(c: u8) -> Option { + match c { + b'0'..=b'9' => Some(c - b'0'), + b'a'..=b'f' => Some(c - b'a' + 10), + b'A'..=b'F' => Some(c - b'A' + 10), + _ => None, } } -} -fn hex_nibble(n: u8) -> u8 { - match n { - 0..=9 => b'0' + n, - 10..=15 => b'a' + n, - _ => unreachable!(), - } -} - -fn binascii_hexlify(data: PyBytesLike) -> Vec { - data.with_ref(|bytes| { - let mut hex = Vec::::with_capacity(bytes.len() * 2); - for b in bytes.iter() { - hex.push(hex_nibble(b >> 4)); - hex.push(hex_nibble(b & 0xf)); - } - hex - }) -} - -fn unhex_nibble(c: u8) -> Option { - match c { - b'0'..=b'9' => Some(c - b'0'), - b'a'..=b'f' => Some(c - b'a' + 10), - b'A'..=b'F' => Some(c - b'A' + 10), - _ => None, - } -} - -fn binascii_unhexlify(data: SerializedData, vm: &VirtualMachine) -> PyResult> { - data.with_ref(|hex_bytes| { - if hex_bytes.len() % 2 != 0 { - return Err(vm.new_value_error("Odd-length string".to_owned())); - } - - let mut unhex = Vec::::with_capacity(hex_bytes.len() / 2); - for (n1, n2) in hex_bytes.iter().tuples() { - if let (Some(n1), Some(n2)) = (unhex_nibble(*n1), unhex_nibble(*n2)) { - unhex.push(n1 << 4 | n2); - } else { - return Err(vm.new_value_error("Non-hexadecimal digit found".to_owned())); + #[pyfunction(name = "a2b_hex")] + #[pyfunction] + fn unhexlify(data: SerializedData, vm: &VirtualMachine) -> PyResult> { + data.with_ref(|hex_bytes| { + if hex_bytes.len() % 2 != 0 { + return Err(vm.new_value_error("Odd-length string".to_owned())); } + + let mut unhex = Vec::::with_capacity(hex_bytes.len() / 2); + for (n1, n2) in hex_bytes.iter().tuples() { + if let (Some(n1), Some(n2)) = (unhex_nibble(*n1), unhex_nibble(*n2)) { + unhex.push(n1 << 4 | n2); + } else { + return Err(vm.new_value_error("Non-hexadecimal digit found".to_owned())); + } + } + + Ok(unhex) + }) + } + + #[pyfunction] + fn crc32(data: SerializedData, value: OptionalArg, vm: &VirtualMachine) -> PyResult { + let crc = value.unwrap_or(0); + + let mut digest = crc32::Digest::new_with_initial(crc32::IEEE, crc); + data.with_ref(|bytes| digest.write(&bytes)); + + Ok(vm.ctx.new_int(digest.sum32())) + } + + #[derive(FromArgs)] + struct NewlineArg { + #[pyarg(keyword_only, default = "true")] + newline: bool, + } + + /// trim a newline from the end of the bytestring, if it exists + fn trim_newline(b: &[u8]) -> &[u8] { + if b.ends_with(b"\n") { + &b[..b.len() - 1] + } else { + b } + } - Ok(unhex) - }) -} + #[pyfunction] + fn a2b_base64(s: SerializedData, vm: &VirtualMachine) -> PyResult> { + s.with_ref(|b| base64::decode(trim_newline(b))) + .map_err(|err| vm.new_value_error(format!("error decoding base64: {}", err))) + } -fn binascii_crc32(data: SerializedData, value: OptionalArg, vm: &VirtualMachine) -> PyResult { - let crc = value.unwrap_or(0); - - let mut digest = crc32::Digest::new_with_initial(crc32::IEEE, crc); - data.with_ref(|bytes| digest.write(&bytes)); - - Ok(vm.ctx.new_int(digest.sum32())) -} - -#[derive(FromArgs)] -struct NewlineArg { - #[pyarg(keyword_only, default = "true")] - newline: bool, -} - -/// trim a newline from the end of the bytestring, if it exists -fn trim_newline(b: &[u8]) -> &[u8] { - if b.ends_with(b"\n") { - &b[..b.len() - 1] - } else { - b + #[pyfunction] + fn b2a_base64(data: PyBytesLike, NewlineArg { newline }: NewlineArg) -> Vec { + let mut encoded = data.with_ref(base64::encode).into_bytes(); + if newline { + encoded.push(b'\n'); + } + encoded } } - -fn binascii_a2b_base64(s: SerializedData, vm: &VirtualMachine) -> PyResult> { - s.with_ref(|b| base64::decode(trim_newline(b))) - .map_err(|err| vm.new_value_error(format!("error decoding base64: {}", err))) -} - -fn binascii_b2a_base64(data: PyBytesLike, NewlineArg { newline }: NewlineArg) -> Vec { - let mut encoded = data.with_ref(base64::encode).into_bytes(); - if newline { - encoded.push(b'\n'); - } - encoded -} - -pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { - let ctx = &vm.ctx; - - py_module!(vm, "binascii", { - "hexlify" => ctx.new_function(binascii_hexlify), - "b2a_hex" => ctx.new_function(binascii_hexlify), - "unhexlify" => ctx.new_function(binascii_unhexlify), - "a2b_hex" => ctx.new_function(binascii_unhexlify), - "crc32" => ctx.new_function(binascii_crc32), - "a2b_base64" => ctx.new_function(binascii_a2b_base64), - "b2a_base64" => ctx.new_function(binascii_b2a_base64), - }) -} diff --git a/vm/src/stdlib/collections.rs b/vm/src/stdlib/collections.rs index 5a5c6b9ea..4ec3633b5 100644 --- a/vm/src/stdlib/collections.rs +++ b/vm/src/stdlib/collections.rs @@ -1,395 +1,398 @@ -use crate::function::OptionalArg; -use crate::obj::{objiter, objtype::PyClassRef}; -use crate::pyobject::{ - IdProtocol, PyArithmaticValue::*, PyClassImpl, PyComparisonValue, PyIterable, PyObjectRef, - PyRef, PyResult, PyValue, -}; -use crate::sequence::{self, SimpleSeq}; -use crate::vm::ReprGuard; -use crate::VirtualMachine; -use itertools::Itertools; -use std::cell::{Cell, RefCell}; -use std::collections::VecDeque; +pub(crate) use _collections::make_module; -#[pyclass(name = "deque")] -#[derive(Debug, Clone)] -struct PyDeque { - deque: RefCell>, - maxlen: Cell>, -} -type PyDequeRef = PyRef; +#[pymodule] +mod _collections { + use crate::function::OptionalArg; + use crate::obj::{objiter, objtype::PyClassRef}; + use crate::pyobject::{ + IdProtocol, PyArithmaticValue::*, PyClassImpl, PyComparisonValue, PyIterable, PyObjectRef, + PyRef, PyResult, PyValue, + }; + use crate::sequence; + use crate::vm::ReprGuard; + use crate::VirtualMachine; + use itertools::Itertools; + use std::cell::{Cell, RefCell}; + use std::collections::VecDeque; -impl PyValue for PyDeque { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("_collections", "deque") + #[pyclass(name = "deque")] + #[derive(Debug, Clone)] + struct PyDeque { + deque: RefCell>, + maxlen: Cell>, } -} + type PyDequeRef = PyRef; -#[derive(FromArgs)] -struct PyDequeOptions { - #[pyarg(positional_or_keyword, default = "None")] - maxlen: Option, -} - -impl PyDeque { - fn borrow_deque<'a>(&'a self) -> impl std::ops::Deref> + 'a { - self.deque.borrow() - } -} - -#[pyimpl(flags(BASETYPE))] -impl PyDeque { - #[pyslot] - fn tp_new( - cls: PyClassRef, - iter: OptionalArg, - PyDequeOptions { maxlen }: PyDequeOptions, - vm: &VirtualMachine, - ) -> PyResult> { - let py_deque = PyDeque { - deque: RefCell::default(), - maxlen: maxlen.into(), - }; - if let OptionalArg::Present(iter) = iter { - py_deque.extend(iter, vm)?; + impl PyValue for PyDeque { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("_collections", "deque") } - py_deque.into_ref_with_type(vm, cls) } - #[pymethod] - fn append(&self, obj: PyObjectRef) { - let mut deque = self.deque.borrow_mut(); - if self.maxlen.get() == Some(deque.len()) { - deque.pop_front(); + #[derive(FromArgs)] + struct PyDequeOptions { + #[pyarg(positional_or_keyword, default = "None")] + maxlen: Option, + } + + impl PyDeque { + fn borrow_deque<'a>(&'a self) -> impl std::ops::Deref> + 'a { + self.deque.borrow() } - deque.push_back(obj); } - #[pymethod] - fn appendleft(&self, obj: PyObjectRef) { - let mut deque = self.deque.borrow_mut(); - if self.maxlen.get() == Some(deque.len()) { - deque.pop_back(); - } - deque.push_front(obj); - } - - #[pymethod] - fn clear(&self) { - self.deque.borrow_mut().clear() - } - - #[pymethod] - fn copy(&self) -> Self { - self.clone() - } - - #[pymethod] - fn count(&self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult { - let mut count = 0; - for elem in self.deque.borrow().iter() { - if vm.identical_or_equal(elem, &obj)? { - count += 1; + #[pyimpl(flags(BASETYPE))] + impl PyDeque { + #[pyslot] + fn tp_new( + cls: PyClassRef, + iter: OptionalArg, + PyDequeOptions { maxlen }: PyDequeOptions, + vm: &VirtualMachine, + ) -> PyResult> { + let py_deque = PyDeque { + deque: RefCell::default(), + maxlen: maxlen.into(), + }; + if let OptionalArg::Present(iter) = iter { + py_deque.extend(iter, vm)?; } + py_deque.into_ref_with_type(vm, cls) } - Ok(count) - } - #[pymethod] - fn extend(&self, iter: PyIterable, vm: &VirtualMachine) -> PyResult<()> { - // TODO: use length_hint here and for extendleft - for elem in iter.iter(vm)? { - self.append(elem?); - } - Ok(()) - } - - #[pymethod] - fn extendleft(&self, iter: PyIterable, vm: &VirtualMachine) -> PyResult<()> { - for elem in iter.iter(vm)? { - self.appendleft(elem?); - } - Ok(()) - } - - #[pymethod] - fn index( - &self, - obj: PyObjectRef, - start: OptionalArg, - stop: OptionalArg, - vm: &VirtualMachine, - ) -> PyResult { - let deque = self.deque.borrow(); - let start = start.unwrap_or(0); - let stop = stop.unwrap_or_else(|| deque.len()); - for (i, elem) in deque.iter().skip(start).take(stop - start).enumerate() { - if vm.identical_or_equal(elem, &obj)? { - return Ok(i); + #[pymethod] + fn append(&self, obj: PyObjectRef) { + let mut deque = self.deque.borrow_mut(); + if self.maxlen.get() == Some(deque.len()) { + deque.pop_front(); } - } - Err(vm.new_value_error( - vm.to_repr(&obj) - .map(|repr| format!("{} is not in deque", repr)) - .unwrap_or_else(|_| String::new()), - )) - } - - #[pymethod] - fn insert(&self, idx: i32, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - let mut deque = self.deque.borrow_mut(); - - if self.maxlen.get() == Some(deque.len()) { - return Err(vm.new_index_error("deque already at its maximum size".to_owned())); + deque.push_back(obj); } - let idx = if idx < 0 { - if -idx as usize > deque.len() { - 0 + #[pymethod] + fn appendleft(&self, obj: PyObjectRef) { + let mut deque = self.deque.borrow_mut(); + if self.maxlen.get() == Some(deque.len()) { + deque.pop_back(); + } + deque.push_front(obj); + } + + #[pymethod] + fn clear(&self) { + self.deque.borrow_mut().clear() + } + + #[pymethod] + fn copy(&self) -> Self { + self.clone() + } + + #[pymethod] + fn count(&self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let mut count = 0; + for elem in self.deque.borrow().iter() { + if vm.identical_or_equal(elem, &obj)? { + count += 1; + } + } + Ok(count) + } + + #[pymethod] + fn extend(&self, iter: PyIterable, vm: &VirtualMachine) -> PyResult<()> { + // TODO: use length_hint here and for extendleft + for elem in iter.iter(vm)? { + self.append(elem?); + } + Ok(()) + } + + #[pymethod] + fn extendleft(&self, iter: PyIterable, vm: &VirtualMachine) -> PyResult<()> { + for elem in iter.iter(vm)? { + self.appendleft(elem?); + } + Ok(()) + } + + #[pymethod] + fn index( + &self, + obj: PyObjectRef, + start: OptionalArg, + stop: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + let deque = self.deque.borrow(); + let start = start.unwrap_or(0); + let stop = stop.unwrap_or_else(|| deque.len()); + for (i, elem) in deque.iter().skip(start).take(stop - start).enumerate() { + if vm.identical_or_equal(elem, &obj)? { + return Ok(i); + } + } + Err(vm.new_value_error( + vm.to_repr(&obj) + .map(|repr| format!("{} is not in deque", repr)) + .unwrap_or_else(|_| String::new()), + )) + } + + #[pymethod] + fn insert(&self, idx: i32, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let mut deque = self.deque.borrow_mut(); + + if self.maxlen.get() == Some(deque.len()) { + return Err(vm.new_index_error("deque already at its maximum size".to_owned())); + } + + let idx = if idx < 0 { + if -idx as usize > deque.len() { + 0 + } else { + deque.len() - ((-idx) as usize) + } + } else if idx as usize >= deque.len() { + deque.len() - 1 } else { - deque.len() - ((-idx) as usize) + idx as usize + }; + + deque.insert(idx, obj); + + Ok(()) + } + + #[pymethod] + fn pop(&self, vm: &VirtualMachine) -> PyResult { + self.deque + .borrow_mut() + .pop_back() + .ok_or_else(|| vm.new_index_error("pop from an empty deque".to_owned())) + } + + #[pymethod] + fn popleft(&self, vm: &VirtualMachine) -> PyResult { + self.deque + .borrow_mut() + .pop_front() + .ok_or_else(|| vm.new_index_error("pop from an empty deque".to_owned())) + } + + #[pymethod] + fn remove(&self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let mut deque = self.deque.borrow_mut(); + let mut idx = None; + for (i, elem) in deque.iter().enumerate() { + if vm.identical_or_equal(elem, &obj)? { + idx = Some(i); + break; + } } - } else if idx as usize >= deque.len() { - deque.len() - 1 - } else { - idx as usize - }; + idx.map(|idx| deque.remove(idx).unwrap()) + .ok_or_else(|| vm.new_value_error("deque.remove(x): x not in deque".to_owned())) + } - deque.insert(idx, obj); + #[pymethod] + fn reverse(&self) { + self.deque + .replace_with(|deque| deque.iter().cloned().rev().collect()); + } - Ok(()) - } - - #[pymethod] - fn pop(&self, vm: &VirtualMachine) -> PyResult { - self.deque - .borrow_mut() - .pop_back() - .ok_or_else(|| vm.new_index_error("pop from an empty deque".to_owned())) - } - - #[pymethod] - fn popleft(&self, vm: &VirtualMachine) -> PyResult { - self.deque - .borrow_mut() - .pop_front() - .ok_or_else(|| vm.new_index_error("pop from an empty deque".to_owned())) - } - - #[pymethod] - fn remove(&self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult { - let mut deque = self.deque.borrow_mut(); - let mut idx = None; - for (i, elem) in deque.iter().enumerate() { - if vm.identical_or_equal(elem, &obj)? { - idx = Some(i); - break; + #[pymethod] + fn rotate(&self, mid: OptionalArg) { + let mut deque = self.deque.borrow_mut(); + let mid = mid.unwrap_or(1); + if mid < 0 { + deque.rotate_left(-mid as usize); + } else { + deque.rotate_right(mid as usize); } } - idx.map(|idx| deque.remove(idx).unwrap()) - .ok_or_else(|| vm.new_value_error("deque.remove(x): x not in deque".to_owned())) - } - #[pymethod] - fn reverse(&self) { - self.deque - .replace_with(|deque| deque.iter().cloned().rev().collect()); - } + #[pyproperty] + fn maxlen(&self) -> Option { + self.maxlen.get() + } - #[pymethod] - fn rotate(&self, mid: OptionalArg) { - let mut deque = self.deque.borrow_mut(); - let mid = mid.unwrap_or(1); - if mid < 0 { - deque.rotate_left(-mid as usize); - } else { - deque.rotate_right(mid as usize); + #[pyproperty(setter)] + fn set_maxlen(&self, maxlen: Option) { + self.maxlen.set(maxlen); + } + + #[pymethod(name = "__repr__")] + fn repr(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + let repr = if let Some(_guard) = ReprGuard::enter(zelf.as_object()) { + let elements = zelf + .deque + .borrow() + .iter() + .map(|obj| vm.to_repr(obj)) + .collect::, _>>()?; + let maxlen = zelf + .maxlen + .get() + .map(|maxlen| format!(", maxlen={}", maxlen)) + .unwrap_or_default(); + format!("deque([{}]{})", elements.into_iter().format(", "), maxlen) + } else { + "[...]".to_owned() + }; + Ok(repr) + } + + #[inline] + fn cmp( + &self, + other: PyObjectRef, + op: F, + vm: &VirtualMachine, + ) -> PyResult + where + F: Fn(&VecDeque, &VecDeque) -> PyResult, + { + let r = if let Some(other) = other.payload_if_subclass::(vm) { + Implemented(op(&*self.borrow_deque(), &*other.borrow_deque())?) + } else { + NotImplemented + }; + Ok(r) + } + + #[pymethod(name = "__eq__")] + fn eq( + zelf: PyRef, + other: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + if zelf.as_object().is(&other) { + Ok(Implemented(true)) + } else { + zelf.cmp(other, |a, b| sequence::eq(vm, a, b), vm) + } + } + + #[pymethod(name = "__ne__")] + fn ne( + zelf: PyRef, + other: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + Ok(PyDeque::eq(zelf, other, vm)?.map(|v| !v)) + } + + #[pymethod(name = "__lt__")] + fn lt( + zelf: PyRef, + other: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + if zelf.as_object().is(&other) { + Ok(Implemented(false)) + } else { + zelf.cmp(other, |a, b| sequence::lt(vm, a, b), vm) + } + } + + #[pymethod(name = "__gt__")] + fn gt( + zelf: PyRef, + other: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + if zelf.as_object().is(&other) { + Ok(Implemented(false)) + } else { + zelf.cmp(other, |a, b| sequence::gt(vm, a, b), vm) + } + } + + #[pymethod(name = "__le__")] + fn le( + zelf: PyRef, + other: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + if zelf.as_object().is(&other) { + Ok(Implemented(true)) + } else { + zelf.cmp(other, |a, b| sequence::le(vm, a, b), vm) + } + } + + #[pymethod(name = "__ge__")] + fn ge( + zelf: PyRef, + other: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + if zelf.as_object().is(&other) { + Ok(Implemented(true)) + } else { + zelf.cmp(other, |a, b| sequence::ge(vm, a, b), vm) + } + } + + #[pymethod(name = "__mul__")] + fn mul(&self, n: isize) -> Self { + let deque: &VecDeque<_> = &self.deque.borrow(); + let mul = sequence::seq_mul(deque, n); + let skipped = if let Some(maxlen) = self.maxlen.get() { + mul.len() - maxlen + } else { + 0 + }; + let deque = mul.skip(skipped).cloned().collect(); + PyDeque { + deque: RefCell::new(deque), + maxlen: self.maxlen.clone(), + } + } + + #[pymethod(name = "__len__")] + fn len(&self) -> usize { + self.deque.borrow().len() + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyDequeIterator { + PyDequeIterator { + position: Cell::new(0), + deque: zelf, + } } } - #[pyproperty] - fn maxlen(&self) -> Option { - self.maxlen.get() + #[pyclass(name = "_deque_iterator")] + #[derive(Debug)] + struct PyDequeIterator { + position: Cell, + deque: PyDequeRef, } - #[pyproperty(setter)] - fn set_maxlen(&self, maxlen: Option) { - self.maxlen.set(maxlen); - } - - #[pymethod(name = "__repr__")] - fn repr(zelf: PyRef, vm: &VirtualMachine) -> PyResult { - let repr = if let Some(_guard) = ReprGuard::enter(zelf.as_object()) { - let elements = zelf - .deque - .borrow() - .iter() - .map(|obj| vm.to_repr(obj)) - .collect::, _>>()?; - let maxlen = zelf - .maxlen - .get() - .map(|maxlen| format!(", maxlen={}", maxlen)) - .unwrap_or_default(); - format!("deque([{}]{})", elements.into_iter().format(", "), maxlen) - } else { - "[...]".to_owned() - }; - Ok(repr) - } - - #[inline] - fn cmp(&self, other: PyObjectRef, op: F, vm: &VirtualMachine) -> PyResult - where - F: Fn(&VecDeque, &VecDeque) -> PyResult, - { - let r = if let Some(other) = other.payload_if_subclass::(vm) { - Implemented(op(&*self.borrow_deque(), &*other.borrow_deque())?) - } else { - NotImplemented - }; - Ok(r) - } - - #[pymethod(name = "__eq__")] - fn eq( - zelf: PyRef, - other: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult { - if zelf.as_object().is(&other) { - Ok(Implemented(true)) - } else { - zelf.cmp(other, |a, b| sequence::eq(vm, a, b), vm) + impl PyValue for PyDequeIterator { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("_collections", "_deque_iterator") } } - #[pymethod(name = "__ne__")] - fn ne( - zelf: PyRef, - other: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult { - Ok(PyDeque::eq(zelf, other, vm)?.map(|v| !v)) - } - - #[pymethod(name = "__lt__")] - fn lt( - zelf: PyRef, - other: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult { - if zelf.as_object().is(&other) { - Ok(Implemented(false)) - } else { - zelf.cmp(other, |a, b| sequence::lt(vm, a, b), vm) + #[pyimpl] + impl PyDequeIterator { + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + if self.position.get() < self.deque.deque.borrow().len() { + let ret = self.deque.deque.borrow()[self.position.get()].clone(); + self.position.set(self.position.get() + 1); + Ok(ret) + } else { + Err(objiter::new_stop_iteration(vm)) + } } - } - #[pymethod(name = "__gt__")] - fn gt( - zelf: PyRef, - other: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult { - if zelf.as_object().is(&other) { - Ok(Implemented(false)) - } else { - zelf.cmp(other, |a, b| sequence::gt(vm, a, b), vm) - } - } - - #[pymethod(name = "__le__")] - fn le( - zelf: PyRef, - other: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult { - if zelf.as_object().is(&other) { - Ok(Implemented(true)) - } else { - zelf.cmp(other, |a, b| sequence::le(vm, a, b), vm) - } - } - - #[pymethod(name = "__ge__")] - fn ge( - zelf: PyRef, - other: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult { - if zelf.as_object().is(&other) { - Ok(Implemented(true)) - } else { - zelf.cmp(other, |a, b| sequence::ge(vm, a, b), vm) - } - } - - #[pymethod(name = "__mul__")] - fn mul(&self, n: isize) -> Self { - let deque: &VecDeque<_> = &self.deque.borrow(); - let mul = sequence::seq_mul(deque, n); - let skipped = if let Some(maxlen) = self.maxlen.get() { - mul.len() - maxlen - } else { - 0 - }; - let deque = mul.skip(skipped).cloned().collect(); - PyDeque { - deque: RefCell::new(deque), - maxlen: self.maxlen.clone(), - } - } - - #[pymethod(name = "__len__")] - fn len(&self) -> usize { - self.deque.borrow().len() - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyDequeIterator { - PyDequeIterator { - position: Cell::new(0), - deque: zelf, + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf } } } - -#[pyclass(name = "_deque_iterator")] -#[derive(Debug)] -struct PyDequeIterator { - position: Cell, - deque: PyDequeRef, -} - -impl PyValue for PyDequeIterator { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("_collections", "_deque_iterator") - } -} - -#[pyimpl] -impl PyDequeIterator { - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if self.position.get() < self.deque.deque.borrow().len() { - let ret = self.deque.deque.borrow()[self.position.get()].clone(); - self.position.set(self.position.get() + 1); - Ok(ret) - } else { - Err(objiter::new_stop_iteration(vm)) - } - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} - -pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { - py_module!(vm, "_collections", { - "deque" => PyDeque::make_class(&vm.ctx), - "_deque_iterator" => PyDequeIterator::make_class(&vm.ctx), - }) -} diff --git a/vm/src/stdlib/imp.rs b/vm/src/stdlib/imp.rs index ee0d78de6..43e8fa932 100644 --- a/vm/src/stdlib/imp.rs +++ b/vm/src/stdlib/imp.rs @@ -86,7 +86,7 @@ fn imp_fix_co_filename(_code: PyObjectRef, _path: PyStringRef) { pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let ctx = &vm.ctx; - let module = py_module!(vm, "_imp", { + py_module!(vm, "_imp", { "extension_suffixes" => ctx.new_function(imp_extension_suffixes), "acquire_lock" => ctx.new_function(imp_acquire_lock), "release_lock" => ctx.new_function(imp_release_lock), @@ -99,7 +99,5 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { "init_frozen" => ctx.new_function(imp_init_frozen), "is_frozen_package" => ctx.new_function(imp_is_frozen_package), "_fix_co_filename" => ctx.new_function(imp_fix_co_filename), - }); - - module + }) } diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs index 097ef72c0..41245e640 100644 --- a/vm/src/stdlib/io.rs +++ b/vm/src/stdlib/io.rs @@ -911,7 +911,7 @@ pub fn io_open(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { _ => unimplemented!("'a' mode is not yet implemented"), }; - let io_obj = match typ.chars().next().unwrap() { + match typ.chars().next().unwrap() { // If the mode is text this buffer type is consumed on construction of // a TextIOWrapper which is subsequently returned. 't' => { @@ -923,8 +923,7 @@ pub fn io_open(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult { // For Buffered class construct "raw" IO class e.g. FileIO and pass this into corresponding field 'b' => buffered, _ => unreachable!(), - }; - io_obj + } } pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index 2b4de0e25..ea7767fb2 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -1,1411 +1,1345 @@ -use std::cell::{Cell, RefCell}; -use std::iter; -use std::rc::Rc; +pub(crate) use decl::make_module; -use num_bigint::BigInt; -use num_traits::{One, Signed, ToPrimitive, Zero}; +#[pymodule(name = "itertools")] +mod decl { + use num_bigint::BigInt; + use num_traits::{One, Signed, ToPrimitive, Zero}; + use std::cell::{Cell, RefCell}; + use std::iter; + use std::rc::Rc; -use crate::function::{Args, OptionalArg, OptionalOption, PyFuncArgs}; -use crate::obj::objbool; -use crate::obj::objint::{self, PyInt, PyIntRef}; -use crate::obj::objiter::{call_next, get_all, get_iter, get_next_object, new_stop_iteration}; -use crate::obj::objtuple::PyTuple; -use crate::obj::objtype::{self, PyClassRef}; -use crate::pyobject::{ - IdProtocol, PyCallable, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, -}; -use crate::vm::VirtualMachine; + use crate::function::{Args, OptionalArg, OptionalOption, PyFuncArgs}; + use crate::obj::objbool; + use crate::obj::objint::{self, PyInt, PyIntRef}; + use crate::obj::objiter::{call_next, get_all, get_iter, get_next_object, new_stop_iteration}; + use crate::obj::objtuple::PyTuple; + use crate::obj::objtype::{self, PyClassRef}; + use crate::pyobject::{ + IdProtocol, PyCallable, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol, + }; + use crate::vm::VirtualMachine; -#[pyclass(name = "chain")] -#[derive(Debug)] -struct PyItertoolsChain { - iterables: Vec, - cur: RefCell<(usize, Option)>, -} - -impl PyValue for PyItertoolsChain { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "chain") + #[pyclass(name = "chain")] + #[derive(Debug)] + struct PyItertoolsChain { + iterables: Vec, + cur: RefCell<(usize, Option)>, } -} -#[pyimpl] -impl PyItertoolsChain { - #[pyslot] - fn tp_new(cls: PyClassRef, args: PyFuncArgs, vm: &VirtualMachine) -> PyResult> { - PyItertoolsChain { - iterables: args.args, - cur: RefCell::new((0, None)), + impl PyValue for PyItertoolsChain { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "chain") } - .into_ref_with_type(vm, cls) } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let (ref mut cur_idx, ref mut cur_iter) = *self.cur.borrow_mut(); - while *cur_idx < self.iterables.len() { - if cur_iter.is_none() { - *cur_iter = Some(get_iter(vm, &self.iterables[*cur_idx])?); + #[pyimpl] + impl PyItertoolsChain { + #[pyslot] + fn tp_new(cls: PyClassRef, args: PyFuncArgs, vm: &VirtualMachine) -> PyResult> { + PyItertoolsChain { + iterables: args.args, + cur: RefCell::new((0, None)), } + .into_ref_with_type(vm, cls) + } - // can't be directly inside the 'match' clause, otherwise the borrows collide. - let obj = call_next(vm, cur_iter.as_ref().unwrap()); - match obj { - Ok(ok) => return Ok(ok), - Err(err) => { - if objtype::isinstance(&err, &vm.ctx.exceptions.stop_iteration) { - *cur_idx += 1; - *cur_iter = None; - } else { - return Err(err); + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + let (ref mut cur_idx, ref mut cur_iter) = *self.cur.borrow_mut(); + while *cur_idx < self.iterables.len() { + if cur_iter.is_none() { + *cur_iter = Some(get_iter(vm, &self.iterables[*cur_idx])?); + } + + // can't be directly inside the 'match' clause, otherwise the borrows collide. + let obj = call_next(vm, cur_iter.as_ref().unwrap()); + match obj { + Ok(ok) => return Ok(ok), + Err(err) => { + if objtype::isinstance(&err, &vm.ctx.exceptions.stop_iteration) { + *cur_idx += 1; + *cur_iter = None; + } else { + return Err(err); + } } } } - } - Err(new_stop_iteration(vm)) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pyclassmethod(name = "from_iterable")] - fn from_iterable( - cls: PyClassRef, - iterable: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult> { - let it = get_iter(vm, &iterable)?; - let iterables = get_all(vm, &it)?; - - PyItertoolsChain { - iterables, - cur: RefCell::new((0, None)), - } - .into_ref_with_type(vm, cls) - } -} - -#[pyclass(name = "compress")] -#[derive(Debug)] -struct PyItertoolsCompress { - data: PyObjectRef, - selector: PyObjectRef, -} - -impl PyValue for PyItertoolsCompress { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "compress") - } -} - -#[pyimpl] -impl PyItertoolsCompress { - #[pyslot] - fn tp_new( - cls: PyClassRef, - data: PyObjectRef, - selector: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult> { - let data_iter = get_iter(vm, &data)?; - let selector_iter = get_iter(vm, &selector)?; - - PyItertoolsCompress { - data: data_iter, - selector: selector_iter, - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - loop { - let sel_obj = call_next(vm, &self.selector)?; - let verdict = objbool::boolval(vm, sel_obj.clone())?; - let data_obj = call_next(vm, &self.data)?; - - if verdict { - return Ok(data_obj); - } - } - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsCount { - cur: RefCell, - step: BigInt, -} - -impl PyValue for PyItertoolsCount { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "count") - } -} - -#[pyimpl] -impl PyItertoolsCount { - #[pyslot] - fn tp_new( - cls: PyClassRef, - start: OptionalArg, - step: OptionalArg, - vm: &VirtualMachine, - ) -> PyResult> { - let start = match start.into_option() { - Some(int) => int.as_bigint().clone(), - None => BigInt::zero(), - }; - let step = match step.into_option() { - Some(int) => int.as_bigint().clone(), - None => BigInt::one(), - }; - - PyItertoolsCount { - cur: RefCell::new(start), - step, - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__next__")] - fn next(&self) -> PyResult { - let result = self.cur.borrow().clone(); - *self.cur.borrow_mut() += &self.step; - Ok(PyInt::new(result)) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsCycle { - iter: RefCell, - saved: RefCell>, - index: Cell, - first_pass: Cell, -} - -impl PyValue for PyItertoolsCycle { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "cycle") - } -} - -#[pyimpl] -impl PyItertoolsCycle { - #[pyslot] - fn tp_new( - cls: PyClassRef, - iterable: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; - - PyItertoolsCycle { - iter: RefCell::new(iter.clone()), - saved: RefCell::new(Vec::new()), - index: Cell::new(0), - first_pass: Cell::new(false), - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let item = if let Some(item) = get_next_object(vm, &self.iter.borrow())? { - if self.first_pass.get() { - return Ok(item); - } - - self.saved.borrow_mut().push(item.clone()); - item - } else { - if self.saved.borrow().len() == 0 { - return Err(new_stop_iteration(vm)); - } - - let last_index = self.index.get(); - self.index.set(self.index.get() + 1); - - if self.index.get() >= self.saved.borrow().len() { - self.index.set(0); - } - - self.saved.borrow()[last_index].clone() - }; - - Ok(item) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsRepeat { - object: PyObjectRef, - times: Option>, -} - -impl PyValue for PyItertoolsRepeat { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "repeat") - } -} - -#[pyimpl] -impl PyItertoolsRepeat { - #[pyslot] - fn tp_new( - cls: PyClassRef, - object: PyObjectRef, - times: OptionalArg, - vm: &VirtualMachine, - ) -> PyResult> { - let times = match times.into_option() { - Some(int) => Some(RefCell::new(int.as_bigint().clone())), - None => None, - }; - - PyItertoolsRepeat { - object: object.clone(), - times, - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if let Some(ref times) = self.times { - if *times.borrow() <= BigInt::zero() { - return Err(new_stop_iteration(vm)); - } - *times.borrow_mut() -= 1; - } - - Ok(self.object.clone()) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__length_hint__")] - fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef { - match self.times { - Some(ref times) => vm.new_int(times.borrow().clone()), - None => vm.new_int(0), - } - } -} - -#[pyclass(name = "starmap")] -#[derive(Debug)] -struct PyItertoolsStarmap { - function: PyObjectRef, - iter: PyObjectRef, -} - -impl PyValue for PyItertoolsStarmap { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "starmap") - } -} - -#[pyimpl] -impl PyItertoolsStarmap { - #[pyslot] - fn tp_new( - cls: PyClassRef, - function: PyObjectRef, - iterable: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; - - PyItertoolsStarmap { function, iter }.into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let obj = call_next(vm, &self.iter)?; - let function = &self.function; - - vm.invoke(function, vm.extract_elements(&obj)?) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsTakewhile { - predicate: PyObjectRef, - iterable: PyObjectRef, - stop_flag: RefCell, -} - -impl PyValue for PyItertoolsTakewhile { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "takewhile") - } -} - -#[pyimpl] -impl PyItertoolsTakewhile { - #[pyslot] - fn tp_new( - cls: PyClassRef, - predicate: PyObjectRef, - iterable: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; - - PyItertoolsTakewhile { - predicate, - iterable: iter, - stop_flag: RefCell::new(false), - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if *self.stop_flag.borrow() { - return Err(new_stop_iteration(vm)); - } - - // might be StopIteration or anything else, which is propagated upwards - let obj = call_next(vm, &self.iterable)?; - let predicate = &self.predicate; - - let verdict = vm.invoke(predicate, vec![obj.clone()])?; - let verdict = objbool::boolval(vm, verdict)?; - if verdict { - Ok(obj) - } else { - *self.stop_flag.borrow_mut() = true; Err(new_stop_iteration(vm)) } - } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } -#[pyclass] -#[derive(Debug)] -struct PyItertoolsDropwhile { - predicate: PyCallable, - iterable: PyObjectRef, - start_flag: Cell, -} + #[pyclassmethod(name = "from_iterable")] + fn from_iterable( + cls: PyClassRef, + iterable: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult> { + let it = get_iter(vm, &iterable)?; + let iterables = get_all(vm, &it)?; -impl PyValue for PyItertoolsDropwhile { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "dropwhile") - } -} - -#[pyimpl] -impl PyItertoolsDropwhile { - #[pyslot] - fn tp_new( - cls: PyClassRef, - predicate: PyCallable, - iterable: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; - - PyItertoolsDropwhile { - predicate, - iterable: iter, - start_flag: Cell::new(false), + PyItertoolsChain { + iterables, + cur: RefCell::new((0, None)), + } + .into_ref_with_type(vm, cls) } - .into_ref_with_type(vm, cls) } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let predicate = &self.predicate; - let iterable = &self.iterable; + #[pyclass(name = "compress")] + #[derive(Debug)] + struct PyItertoolsCompress { + data: PyObjectRef, + selector: PyObjectRef, + } - if !self.start_flag.get() { + impl PyValue for PyItertoolsCompress { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "compress") + } + } + + #[pyimpl] + impl PyItertoolsCompress { + #[pyslot] + fn tp_new( + cls: PyClassRef, + data: PyObjectRef, + selector: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult> { + let data_iter = get_iter(vm, &data)?; + let selector_iter = get_iter(vm, &selector)?; + + PyItertoolsCompress { + data: data_iter, + selector: selector_iter, + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { loop { - let obj = call_next(vm, iterable)?; - let pred = predicate.clone(); - let pred_value = vm.invoke(&pred.into_object(), vec![obj.clone()])?; - if !objbool::boolval(vm, pred_value)? { - self.start_flag.set(true); - return Ok(obj); + let sel_obj = call_next(vm, &self.selector)?; + let verdict = objbool::boolval(vm, sel_obj.clone())?; + let data_obj = call_next(vm, &self.data)?; + + if verdict { + return Ok(data_obj); } } } - call_next(vm, iterable) + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf + #[pyclass(name = "count")] + #[derive(Debug)] + struct PyItertoolsCount { + cur: RefCell, + step: BigInt, } -} -#[pyclass(name = "islice")] -#[derive(Debug)] -struct PyItertoolsIslice { - iterable: PyObjectRef, - cur: RefCell, - next: RefCell, - stop: Option, - step: usize, -} - -impl PyValue for PyItertoolsIslice { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "islice") + impl PyValue for PyItertoolsCount { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "count") + } } -} -fn pyobject_to_opt_usize(obj: PyObjectRef, vm: &VirtualMachine) -> Option { - let is_int = objtype::isinstance(&obj, &vm.ctx.int_type()); - if is_int { - objint::get_value(&obj).to_usize() - } else { - None + #[pyimpl] + impl PyItertoolsCount { + #[pyslot] + fn tp_new( + cls: PyClassRef, + start: OptionalArg, + step: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult> { + let start = match start.into_option() { + Some(int) => int.as_bigint().clone(), + None => BigInt::zero(), + }; + let step = match step.into_option() { + Some(int) => int.as_bigint().clone(), + None => BigInt::one(), + }; + + PyItertoolsCount { + cur: RefCell::new(start), + step, + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self) -> PyResult { + let result = self.cur.borrow().clone(); + *self.cur.borrow_mut() += &self.step; + Ok(PyInt::new(result)) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } } -} -#[pyimpl] -impl PyItertoolsIslice { - #[pyslot] - fn tp_new(cls: PyClassRef, args: PyFuncArgs, vm: &VirtualMachine) -> PyResult> { - let (iter, start, stop, step) = match args.args.len() { - 0 | 1 => { - return Err(vm.new_type_error(format!( - "islice expected at least 2 arguments, got {}", - args.args.len() - ))); + #[pyclass(name = "cycle")] + #[derive(Debug)] + struct PyItertoolsCycle { + iter: RefCell, + saved: RefCell>, + index: Cell, + first_pass: Cell, + } + + impl PyValue for PyItertoolsCycle { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "cycle") + } + } + + #[pyimpl] + impl PyItertoolsCycle { + #[pyslot] + fn tp_new( + cls: PyClassRef, + iterable: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult> { + let iter = get_iter(vm, &iterable)?; + + PyItertoolsCycle { + iter: RefCell::new(iter.clone()), + saved: RefCell::new(Vec::new()), + index: Cell::new(0), + first_pass: Cell::new(false), + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + let item = if let Some(item) = get_next_object(vm, &self.iter.borrow())? { + if self.first_pass.get() { + return Ok(item); + } + + self.saved.borrow_mut().push(item.clone()); + item + } else { + if self.saved.borrow().len() == 0 { + return Err(new_stop_iteration(vm)); + } + + let last_index = self.index.get(); + self.index.set(self.index.get() + 1); + + if self.index.get() >= self.saved.borrow().len() { + self.index.set(0); + } + + self.saved.borrow()[last_index].clone() + }; + + Ok(item) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + } + + #[pyclass(name = "repeat")] + #[derive(Debug)] + struct PyItertoolsRepeat { + object: PyObjectRef, + times: Option>, + } + + impl PyValue for PyItertoolsRepeat { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "repeat") + } + } + + #[pyimpl] + impl PyItertoolsRepeat { + #[pyslot] + fn tp_new( + cls: PyClassRef, + object: PyObjectRef, + times: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult> { + let times = match times.into_option() { + Some(int) => Some(RefCell::new(int.as_bigint().clone())), + None => None, + }; + + PyItertoolsRepeat { + object: object.clone(), + times, + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + if let Some(ref times) = self.times { + if *times.borrow() <= BigInt::zero() { + return Err(new_stop_iteration(vm)); + } + *times.borrow_mut() -= 1; } - 2 => { - let (iter, stop): (PyObjectRef, PyObjectRef) = args.bind(vm)?; + Ok(self.object.clone()) + } - (iter, 0usize, stop, 1usize) + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + + #[pymethod(name = "__length_hint__")] + fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef { + match self.times { + Some(ref times) => vm.new_int(times.borrow().clone()), + None => vm.new_int(0), } - _ => { - let (iter, start, stop, step): ( - PyObjectRef, - PyObjectRef, - PyObjectRef, - PyObjectRef, - ) = args.bind(vm)?; + } + } - let start = if !start.is(&vm.get_none()) { - pyobject_to_opt_usize(start, &vm).ok_or_else(|| { + #[pyclass(name = "starmap")] + #[derive(Debug)] + struct PyItertoolsStarmap { + function: PyObjectRef, + iter: PyObjectRef, + } + + impl PyValue for PyItertoolsStarmap { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "starmap") + } + } + + #[pyimpl] + impl PyItertoolsStarmap { + #[pyslot] + fn tp_new( + cls: PyClassRef, + function: PyObjectRef, + iterable: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult> { + let iter = get_iter(vm, &iterable)?; + + PyItertoolsStarmap { function, iter }.into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + let obj = call_next(vm, &self.iter)?; + let function = &self.function; + + vm.invoke(function, vm.extract_elements(&obj)?) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + } + + #[pyclass(name = "takewhile")] + #[derive(Debug)] + struct PyItertoolsTakewhile { + predicate: PyObjectRef, + iterable: PyObjectRef, + stop_flag: RefCell, + } + + impl PyValue for PyItertoolsTakewhile { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "takewhile") + } + } + + #[pyimpl] + impl PyItertoolsTakewhile { + #[pyslot] + fn tp_new( + cls: PyClassRef, + predicate: PyObjectRef, + iterable: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult> { + let iter = get_iter(vm, &iterable)?; + + PyItertoolsTakewhile { + predicate, + iterable: iter, + stop_flag: RefCell::new(false), + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + if *self.stop_flag.borrow() { + return Err(new_stop_iteration(vm)); + } + + // might be StopIteration or anything else, which is propagated upwards + let obj = call_next(vm, &self.iterable)?; + let predicate = &self.predicate; + + let verdict = vm.invoke(predicate, vec![obj.clone()])?; + let verdict = objbool::boolval(vm, verdict)?; + if verdict { + Ok(obj) + } else { + *self.stop_flag.borrow_mut() = true; + Err(new_stop_iteration(vm)) + } + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + } + + #[pyclass(name = "dropwhile")] + #[derive(Debug)] + struct PyItertoolsDropwhile { + predicate: PyCallable, + iterable: PyObjectRef, + start_flag: Cell, + } + + impl PyValue for PyItertoolsDropwhile { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "dropwhile") + } + } + + #[pyimpl] + impl PyItertoolsDropwhile { + #[pyslot] + fn tp_new( + cls: PyClassRef, + predicate: PyCallable, + iterable: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult> { + let iter = get_iter(vm, &iterable)?; + + PyItertoolsDropwhile { + predicate, + iterable: iter, + start_flag: Cell::new(false), + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + let predicate = &self.predicate; + let iterable = &self.iterable; + + if !self.start_flag.get() { + loop { + let obj = call_next(vm, iterable)?; + let pred = predicate.clone(); + let pred_value = vm.invoke(&pred.into_object(), vec![obj.clone()])?; + if !objbool::boolval(vm, pred_value)? { + self.start_flag.set(true); + return Ok(obj); + } + } + } + call_next(vm, iterable) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + } + + #[pyclass(name = "islice")] + #[derive(Debug)] + struct PyItertoolsIslice { + iterable: PyObjectRef, + cur: RefCell, + next: RefCell, + stop: Option, + step: usize, + } + + impl PyValue for PyItertoolsIslice { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "islice") + } + } + + fn pyobject_to_opt_usize(obj: PyObjectRef, vm: &VirtualMachine) -> Option { + let is_int = objtype::isinstance(&obj, &vm.ctx.int_type()); + if is_int { + objint::get_value(&obj).to_usize() + } else { + None + } + } + + #[pyimpl] + impl PyItertoolsIslice { + #[pyslot] + fn tp_new(cls: PyClassRef, args: PyFuncArgs, vm: &VirtualMachine) -> PyResult> { + let (iter, start, stop, step) = match args.args.len() { + 0 | 1 => { + return Err(vm.new_type_error(format!( + "islice expected at least 2 arguments, got {}", + args.args.len() + ))); + } + + 2 => { + let (iter, stop): (PyObjectRef, PyObjectRef) = args.bind(vm)?; + + (iter, 0usize, stop, 1usize) + } + _ => { + let (iter, start, stop, step): ( + PyObjectRef, + PyObjectRef, + PyObjectRef, + PyObjectRef, + ) = args.bind(vm)?; + + let start = if !start.is(&vm.get_none()) { + pyobject_to_opt_usize(start, &vm).ok_or_else(|| { vm.new_value_error( "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.".to_owned(), ) })? - } else { - 0usize - }; + } else { + 0usize + }; - let step = if !step.is(&vm.get_none()) { - pyobject_to_opt_usize(step, &vm).ok_or_else(|| { - vm.new_value_error( - "Step for islice() must be a positive integer or None.".to_owned(), - ) - })? - } else { - 1usize - }; + let step = if !step.is(&vm.get_none()) { + pyobject_to_opt_usize(step, &vm).ok_or_else(|| { + vm.new_value_error( + "Step for islice() must be a positive integer or None.".to_owned(), + ) + })? + } else { + 1usize + }; - (iter, start, stop, step) - } - }; + (iter, start, stop, step) + } + }; - let stop = if !stop.is(&vm.get_none()) { - Some(pyobject_to_opt_usize(stop, &vm).ok_or_else(|| { - vm.new_value_error( + let stop = if !stop.is(&vm.get_none()) { + Some(pyobject_to_opt_usize(stop, &vm).ok_or_else(|| { + vm.new_value_error( "Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize." .to_owned(), ) - })?) - } else { - None - }; - - let iter = get_iter(vm, &iter)?; - - PyItertoolsIslice { - iterable: iter, - cur: RefCell::new(0), - next: RefCell::new(start), - stop, - step, - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - while *self.cur.borrow() < *self.next.borrow() { - call_next(vm, &self.iterable)?; - *self.cur.borrow_mut() += 1; - } - - if let Some(stop) = self.stop { - if *self.cur.borrow() >= stop { - return Err(new_stop_iteration(vm)); - } - } - - let obj = call_next(vm, &self.iterable)?; - *self.cur.borrow_mut() += 1; - - // TODO is this overflow check required? attempts to copy CPython. - let (next, ovf) = (*self.next.borrow()).overflowing_add(self.step); - *self.next.borrow_mut() = if ovf { self.stop.unwrap() } else { next }; - - Ok(obj) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsFilterFalse { - predicate: PyObjectRef, - iterable: PyObjectRef, -} - -impl PyValue for PyItertoolsFilterFalse { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "filterfalse") - } -} - -#[pyimpl] -impl PyItertoolsFilterFalse { - #[pyslot] - fn tp_new( - cls: PyClassRef, - predicate: PyObjectRef, - iterable: PyObjectRef, - vm: &VirtualMachine, - ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; - - PyItertoolsFilterFalse { - predicate, - iterable: iter, - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let predicate = &self.predicate; - let iterable = &self.iterable; - - loop { - let obj = call_next(vm, iterable)?; - let pred_value = if predicate.is(&vm.get_none()) { - obj.clone() + })?) } else { - vm.invoke(predicate, vec![obj.clone()])? + None }; - if !objbool::boolval(vm, pred_value)? { - return Ok(obj); + let iter = get_iter(vm, &iter)?; + + PyItertoolsIslice { + iterable: iter, + cur: RefCell::new(0), + next: RefCell::new(start), + stop, + step, } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + while *self.cur.borrow() < *self.next.borrow() { + call_next(vm, &self.iterable)?; + *self.cur.borrow_mut() += 1; + } + + if let Some(stop) = self.stop { + if *self.cur.borrow() >= stop { + return Err(new_stop_iteration(vm)); + } + } + + let obj = call_next(vm, &self.iterable)?; + *self.cur.borrow_mut() += 1; + + // TODO is this overflow check required? attempts to copy CPython. + let (next, ovf) = (*self.next.borrow()).overflowing_add(self.step); + *self.next.borrow_mut() = if ovf { self.stop.unwrap() } else { next }; + + Ok(obj) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf } } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsAccumulate { - iterable: PyObjectRef, - binop: PyObjectRef, - acc_value: RefCell>, -} - -impl PyValue for PyItertoolsAccumulate { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "accumulate") - } -} - -#[pyimpl] -impl PyItertoolsAccumulate { - #[pyslot] - fn tp_new( - cls: PyClassRef, + #[pyclass(name = "filterfalse")] + #[derive(Debug)] + struct PyItertoolsFilterFalse { + predicate: PyObjectRef, iterable: PyObjectRef, - binop: OptionalArg, - vm: &VirtualMachine, - ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; - - PyItertoolsAccumulate { - iterable: iter, - binop: binop.unwrap_or_else(|| vm.get_none()), - acc_value: RefCell::from(Option::None), - } - .into_ref_with_type(vm, cls) } - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let iterable = &self.iterable; - let obj = call_next(vm, iterable)?; + impl PyValue for PyItertoolsFilterFalse { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "filterfalse") + } + } - let next_acc_value = match &*self.acc_value.borrow() { - None => obj.clone(), - Some(value) => { - if self.binop.is(&vm.get_none()) { - vm._add(value.clone(), obj.clone())? + #[pyimpl] + impl PyItertoolsFilterFalse { + #[pyslot] + fn tp_new( + cls: PyClassRef, + predicate: PyObjectRef, + iterable: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult> { + let iter = get_iter(vm, &iterable)?; + + PyItertoolsFilterFalse { + predicate, + iterable: iter, + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + let predicate = &self.predicate; + let iterable = &self.iterable; + + loop { + let obj = call_next(vm, iterable)?; + let pred_value = if predicate.is(&vm.get_none()) { + obj.clone() } else { - vm.invoke(&self.binop, vec![value.clone(), obj.clone()])? - } - } - }; - self.acc_value.replace(Option::from(next_acc_value.clone())); - - Ok(next_acc_value) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} - -#[derive(Debug)] -struct PyItertoolsTeeData { - iterable: PyObjectRef, - values: RefCell>, -} - -impl PyItertoolsTeeData { - fn new(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult> { - Ok(Rc::new(PyItertoolsTeeData { - iterable: get_iter(vm, &iterable)?, - values: RefCell::new(vec![]), - })) - } - - fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult { - if self.values.borrow().len() == index { - let result = call_next(vm, &self.iterable)?; - self.values.borrow_mut().push(result); - } - Ok(self.values.borrow()[index].clone()) - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsTee { - tee_data: Rc, - index: Cell, -} - -impl PyValue for PyItertoolsTee { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "tee") - } -} - -#[pyimpl] -impl PyItertoolsTee { - fn from_iter(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult { - let it = get_iter(vm, &iterable)?; - if it.class().is(&PyItertoolsTee::class(vm)) { - return vm.call_method(&it, "__copy__", PyFuncArgs::from(vec![])); - } - Ok(PyItertoolsTee { - tee_data: PyItertoolsTeeData::new(it, vm)?, - index: Cell::from(0), - } - .into_ref_with_type(vm, PyItertoolsTee::class(vm))? - .into_object()) - } - - #[pyslot] - #[allow(clippy::new_ret_no_self)] - fn tp_new( - _cls: PyClassRef, - iterable: PyObjectRef, - n: OptionalArg, - vm: &VirtualMachine, - ) -> PyResult> { - let n = n.unwrap_or(2); - - let copyable = if iterable.class().has_attr("__copy__") { - vm.call_method(&iterable, "__copy__", PyFuncArgs::from(vec![]))? - } else { - PyItertoolsTee::from_iter(iterable, vm)? - }; - - let mut tee_vec: Vec = Vec::with_capacity(n); - for _ in 0..n { - let no_args = PyFuncArgs::from(vec![]); - tee_vec.push(vm.call_method(©able, "__copy__", no_args)?); - } - - Ok(PyTuple::from(tee_vec).into_ref(vm)) - } - - #[pymethod(name = "__copy__")] - fn copy(&self, vm: &VirtualMachine) -> PyResult { - Ok(PyItertoolsTee { - tee_data: Rc::clone(&self.tee_data), - index: self.index.clone(), - } - .into_ref_with_type(vm, Self::class(vm))? - .into_object()) - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - let value = self.tee_data.get_item(vm, self.index.get())?; - self.index.set(self.index.get() + 1); - Ok(value) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsProduct { - pools: Vec>, - idxs: RefCell>, - cur: Cell, - stop: Cell, -} - -impl PyValue for PyItertoolsProduct { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "product") - } -} - -#[derive(FromArgs)] -struct ProductArgs { - #[pyarg(keyword_only, optional = true)] - repeat: OptionalArg, -} - -#[pyimpl] -impl PyItertoolsProduct { - #[pyslot] - fn tp_new( - cls: PyClassRef, - iterables: Args, - args: ProductArgs, - vm: &VirtualMachine, - ) -> PyResult> { - let repeat = match args.repeat.into_option() { - Some(i) => i, - None => 1, - }; - - let mut pools = Vec::new(); - for arg in iterables.into_iter() { - let it = get_iter(vm, &arg)?; - let pool = get_all(vm, &it)?; - - pools.push(pool); - } - let pools = iter::repeat(pools) - .take(repeat) - .flatten() - .collect::>>(); - - let l = pools.len(); - - PyItertoolsProduct { - pools, - idxs: RefCell::new(vec![0; l]), - cur: Cell::new(l - 1), - stop: Cell::new(false), - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - // stop signal - if self.stop.get() { - return Err(new_stop_iteration(vm)); - } - - let pools = &self.pools; - - for p in pools { - if p.is_empty() { - return Err(new_stop_iteration(vm)); - } - } - - let res = PyTuple::from( - pools - .iter() - .zip(self.idxs.borrow().iter()) - .map(|(pool, idx)| pool[*idx].clone()) - .collect::>(), - ); - - self.update_idxs(); - - if self.is_end() { - self.stop.set(true); - } - - Ok(res.into_ref(vm).into_object()) - } - - fn is_end(&self) -> bool { - let cur = self.cur.get(); - self.idxs.borrow()[cur] == self.pools[cur].len() - 1 && cur == 0 - } - - fn update_idxs(&self) { - let lst_idx = &self.pools[self.cur.get()].len() - 1; - - if self.idxs.borrow()[self.cur.get()] == lst_idx { - if self.is_end() { - return; - } - self.idxs.borrow_mut()[self.cur.get()] = 0; - self.cur.set(self.cur.get() - 1); - self.update_idxs(); - } else { - self.idxs.borrow_mut()[self.cur.get()] += 1; - self.cur.set(self.idxs.borrow().len() - 1); - } - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsCombinations { - pool: Vec, - indices: RefCell>, - r: Cell, - exhausted: Cell, -} - -impl PyValue for PyItertoolsCombinations { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "combinations") - } -} - -#[pyimpl] -impl PyItertoolsCombinations { - #[pyslot] - fn tp_new( - cls: PyClassRef, - iterable: PyObjectRef, - r: PyIntRef, - vm: &VirtualMachine, - ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; - let pool = get_all(vm, &iter)?; - - let r = r.as_bigint(); - if r.is_negative() { - return Err(vm.new_value_error("r must be non-negative".to_owned())); - } - let r = r.to_usize().unwrap(); - - let n = pool.len(); - - PyItertoolsCombinations { - pool, - indices: RefCell::new((0..r).collect()), - r: Cell::new(r), - exhausted: Cell::new(r > n), - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - // stop signal - if self.exhausted.get() { - return Err(new_stop_iteration(vm)); - } - - let n = self.pool.len(); - let r = self.r.get(); - - if r == 0 { - self.exhausted.set(true); - return Ok(vm.ctx.new_tuple(vec![])); - } - - let res = PyTuple::from( - self.indices - .borrow() - .iter() - .map(|&i| self.pool[i].clone()) - .collect::>(), - ); - - let mut indices = self.indices.borrow_mut(); - - // Scan indices right-to-left until finding one that is not at its maximum (i + n - r). - let mut idx = r as isize - 1; - while idx >= 0 && indices[idx as usize] == idx as usize + n - r { - idx -= 1; - } - - // If no suitable index is found, then the indices are all at - // their maximum value and we're done. - if idx < 0 { - self.exhausted.set(true); - } else { - // Increment the current index which we know is not at its - // maximum. Then move back to the right setting each index - // to its lowest possible value (one higher than the index - // to its left -- this maintains the sort order invariant). - indices[idx as usize] += 1; - for j in idx as usize + 1..r { - indices[j] = indices[j - 1] + 1; - } - } - - Ok(res.into_ref(vm).into_object()) - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsCombinationsWithReplacement { - pool: Vec, - indices: RefCell>, - r: Cell, - exhausted: Cell, -} - -impl PyValue for PyItertoolsCombinationsWithReplacement { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "combinations_with_replacement") - } -} - -#[pyimpl] -impl PyItertoolsCombinationsWithReplacement { - #[pyslot] - fn tp_new( - cls: PyClassRef, - iterable: PyObjectRef, - r: PyIntRef, - vm: &VirtualMachine, - ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; - let pool = get_all(vm, &iter)?; - - let r = r.as_bigint(); - if r.is_negative() { - return Err(vm.new_value_error("r must be non-negative".to_owned())); - } - let r = r.to_usize().unwrap(); - - let n = pool.len(); - - PyItertoolsCombinationsWithReplacement { - pool, - indices: RefCell::new(vec![0; r]), - r: Cell::new(r), - exhausted: Cell::new(n == 0 && r > 0), - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - // stop signal - if self.exhausted.get() { - return Err(new_stop_iteration(vm)); - } - - let n = self.pool.len(); - let r = self.r.get(); - - if r == 0 { - self.exhausted.set(true); - return Ok(vm.ctx.new_tuple(vec![])); - } - - let mut indices = self.indices.borrow_mut(); - - let res = vm - .ctx - .new_tuple(indices.iter().map(|&i| self.pool[i].clone()).collect()); - - // Scan indices right-to-left until finding one that is not at its maximum (i + n - r). - let mut idx = r as isize - 1; - while idx >= 0 && indices[idx as usize] == n - 1 { - idx -= 1; - } - - // If no suitable index is found, then the indices are all at - // their maximum value and we're done. - if idx < 0 { - self.exhausted.set(true); - } else { - let index = indices[idx as usize] + 1; - - // Increment the current index which we know is not at its - // maximum. Then set all to the right to the same value. - for j in idx as usize..r { - indices[j as usize] = index as usize; - } - } - - Ok(res) - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsPermutations { - pool: Vec, // Collected input iterable - indices: RefCell>, // One index per element in pool - cycles: RefCell>, // One rollover counter per element in the result - result: RefCell>>, // Indexes of the most recently returned result - r: Cell, // Size of result tuple - exhausted: Cell, // Set when the iterator is exhausted -} - -impl PyValue for PyItertoolsPermutations { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "permutations") - } -} - -#[pyimpl] -impl PyItertoolsPermutations { - #[pyslot] - fn tp_new( - cls: PyClassRef, - iterable: PyObjectRef, - r: OptionalOption, - vm: &VirtualMachine, - ) -> PyResult> { - let iter = get_iter(vm, &iterable)?; - let pool = get_all(vm, &iter)?; - - let n = pool.len(); - // If r is not provided, r == n. If provided, r must be a positive integer, or None. - // If None, it behaves the same as if it was not provided. - let r = match r.flat_option() { - Some(r) => { - let val = r - .payload::() - .ok_or_else(|| vm.new_type_error("Expected int as r".to_owned()))? - .as_bigint(); - - if val.is_negative() { - return Err(vm.new_value_error("r must be non-negative".to_owned())); - } - val.to_usize().unwrap() - } - None => n, - }; - - PyItertoolsPermutations { - pool, - indices: RefCell::new((0..n).collect()), - cycles: RefCell::new((0..r).map(|i| n - i).collect()), - result: RefCell::new(None), - r: Cell::new(r), - exhausted: Cell::new(r > n), - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - // stop signal - if self.exhausted.get() { - return Err(new_stop_iteration(vm)); - } - - let n = self.pool.len(); - let r = self.r.get(); - - if n == 0 { - self.exhausted.set(true); - return Ok(vm.ctx.new_tuple(vec![])); - } - - let result = &mut *self.result.borrow_mut(); - - if let Some(ref mut result) = result { - let mut indices = self.indices.borrow_mut(); - let mut cycles = self.cycles.borrow_mut(); - let mut sentinel = false; - - // Decrement rightmost cycle, moving leftward upon zero rollover - for i in (0..r).rev() { - cycles[i] -= 1; - - if cycles[i] == 0 { - // rotation: indices[i:] = indices[i+1:] + indices[i:i+1] - let index = indices[i]; - for j in i..n - 1 { - indices[j] = indices[j + i]; - } - indices[n - 1] = index; - cycles[i] = n - i; - } else { - let j = cycles[i]; - indices.swap(i, n - j); - - for k in i..r { - // start with i, the leftmost element that changed - // yield tuple(pool[k] for k in indices[:r]) - result[k] = indices[k]; - } - sentinel = true; - break; - } - } - if !sentinel { - self.exhausted.set(true); - return Err(new_stop_iteration(vm)); - } - } else { - // On the first pass, initialize result tuple using the indices - *result = Some((0..r).collect()); - } - - Ok(vm.ctx.new_tuple( - result - .as_ref() - .unwrap() - .iter() - .map(|&i| self.pool[i].clone()) - .collect(), - )) - } -} - -#[pyclass] -#[derive(Debug)] -struct PyItertoolsZiplongest { - iterators: Vec, - fillvalue: PyObjectRef, - numactive: Cell, -} - -impl PyValue for PyItertoolsZiplongest { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("itertools", "zip_longest") - } -} - -#[derive(FromArgs)] -struct ZiplongestArgs { - #[pyarg(keyword_only, optional = true)] - fillvalue: OptionalArg, -} - -#[pyimpl] -impl PyItertoolsZiplongest { - #[pyslot] - fn tp_new( - cls: PyClassRef, - iterables: Args, - args: ZiplongestArgs, - vm: &VirtualMachine, - ) -> PyResult> { - let fillvalue = match args.fillvalue.into_option() { - Some(i) => i, - None => vm.get_none(), - }; - - let iterators = iterables - .into_iter() - .map(|iterable| get_iter(vm, &iterable)) - .collect::, _>>()?; - - let numactive = Cell::new(iterators.len()); - - PyItertoolsZiplongest { - iterators, - fillvalue, - numactive, - } - .into_ref_with_type(vm, cls) - } - - #[pymethod(name = "__next__")] - fn next(&self, vm: &VirtualMachine) -> PyResult { - if self.iterators.is_empty() { - Err(new_stop_iteration(vm)) - } else { - let mut result: Vec = Vec::new(); - let mut numactive = self.numactive.get(); - - for idx in 0..self.iterators.len() { - let next_obj = match call_next(vm, &self.iterators[idx]) { - Ok(obj) => obj, - Err(err) => { - if !objtype::isinstance(&err, &vm.ctx.exceptions.stop_iteration) { - return Err(err); - } - numactive -= 1; - if numactive == 0 { - return Err(new_stop_iteration(vm)); - } - self.fillvalue.clone() - } + vm.invoke(predicate, vec![obj.clone()])? }; - result.push(next_obj); + + if !objbool::boolval(vm, pred_value)? { + return Ok(obj); + } } - Ok(vm.ctx.new_tuple(result)) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf } } - #[pymethod(name = "__iter__")] - fn iter(zelf: PyRef) -> PyRef { - zelf + #[pyclass(name = "accumulate")] + #[derive(Debug)] + struct PyItertoolsAccumulate { + iterable: PyObjectRef, + binop: PyObjectRef, + acc_value: RefCell>, + } + + impl PyValue for PyItertoolsAccumulate { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "accumulate") + } + } + + #[pyimpl] + impl PyItertoolsAccumulate { + #[pyslot] + fn tp_new( + cls: PyClassRef, + iterable: PyObjectRef, + binop: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult> { + let iter = get_iter(vm, &iterable)?; + + PyItertoolsAccumulate { + iterable: iter, + binop: binop.unwrap_or_else(|| vm.get_none()), + acc_value: RefCell::from(Option::None), + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + let iterable = &self.iterable; + let obj = call_next(vm, iterable)?; + + let next_acc_value = match &*self.acc_value.borrow() { + None => obj.clone(), + Some(value) => { + if self.binop.is(&vm.get_none()) { + vm._add(value.clone(), obj.clone())? + } else { + vm.invoke(&self.binop, vec![value.clone(), obj.clone()])? + } + } + }; + self.acc_value.replace(Option::from(next_acc_value.clone())); + + Ok(next_acc_value) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + } + + #[derive(Debug)] + struct PyItertoolsTeeData { + iterable: PyObjectRef, + values: RefCell>, + } + + impl PyItertoolsTeeData { + fn new(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult> { + Ok(Rc::new(PyItertoolsTeeData { + iterable: get_iter(vm, &iterable)?, + values: RefCell::new(vec![]), + })) + } + + fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult { + if self.values.borrow().len() == index { + let result = call_next(vm, &self.iterable)?; + self.values.borrow_mut().push(result); + } + Ok(self.values.borrow()[index].clone()) + } + } + + #[pyclass(name = "tee")] + #[derive(Debug)] + struct PyItertoolsTee { + tee_data: Rc, + index: Cell, + } + + impl PyValue for PyItertoolsTee { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "tee") + } + } + + #[pyimpl] + impl PyItertoolsTee { + fn from_iter(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let it = get_iter(vm, &iterable)?; + if it.class().is(&PyItertoolsTee::class(vm)) { + return vm.call_method(&it, "__copy__", PyFuncArgs::from(vec![])); + } + Ok(PyItertoolsTee { + tee_data: PyItertoolsTeeData::new(it, vm)?, + index: Cell::from(0), + } + .into_ref_with_type(vm, PyItertoolsTee::class(vm))? + .into_object()) + } + + // TODO: make tee() a function, rename this class to itertools._tee and make + // teedata a python class + #[pyslot] + #[allow(clippy::new_ret_no_self)] + fn tp_new( + _cls: PyClassRef, + iterable: PyObjectRef, + n: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult> { + let n = n.unwrap_or(2); + + let copyable = if iterable.class().has_attr("__copy__") { + vm.call_method(&iterable, "__copy__", PyFuncArgs::from(vec![]))? + } else { + PyItertoolsTee::from_iter(iterable, vm)? + }; + + let mut tee_vec: Vec = Vec::with_capacity(n); + for _ in 0..n { + let no_args = PyFuncArgs::from(vec![]); + tee_vec.push(vm.call_method(©able, "__copy__", no_args)?); + } + + Ok(PyTuple::from(tee_vec).into_ref(vm)) + } + + #[pymethod(name = "__copy__")] + fn copy(&self, vm: &VirtualMachine) -> PyResult { + Ok(PyItertoolsTee { + tee_data: Rc::clone(&self.tee_data), + index: self.index.clone(), + } + .into_ref_with_type(vm, Self::class(vm))? + .into_object()) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + let value = self.tee_data.get_item(vm, self.index.get())?; + self.index.set(self.index.get() + 1); + Ok(value) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + } + + #[pyclass(name = "product")] + #[derive(Debug)] + struct PyItertoolsProduct { + pools: Vec>, + idxs: RefCell>, + cur: Cell, + stop: Cell, + } + + impl PyValue for PyItertoolsProduct { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "product") + } + } + + #[derive(FromArgs)] + struct ProductArgs { + #[pyarg(keyword_only, optional = true)] + repeat: OptionalArg, + } + + #[pyimpl] + impl PyItertoolsProduct { + #[pyslot] + fn tp_new( + cls: PyClassRef, + iterables: Args, + args: ProductArgs, + vm: &VirtualMachine, + ) -> PyResult> { + let repeat = match args.repeat.into_option() { + Some(i) => i, + None => 1, + }; + + let mut pools = Vec::new(); + for arg in iterables.into_iter() { + let it = get_iter(vm, &arg)?; + let pool = get_all(vm, &it)?; + + pools.push(pool); + } + let pools = iter::repeat(pools) + .take(repeat) + .flatten() + .collect::>>(); + + let l = pools.len(); + + PyItertoolsProduct { + pools, + idxs: RefCell::new(vec![0; l]), + cur: Cell::new(l - 1), + stop: Cell::new(false), + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + // stop signal + if self.stop.get() { + return Err(new_stop_iteration(vm)); + } + + let pools = &self.pools; + + for p in pools { + if p.is_empty() { + return Err(new_stop_iteration(vm)); + } + } + + let res = PyTuple::from( + pools + .iter() + .zip(self.idxs.borrow().iter()) + .map(|(pool, idx)| pool[*idx].clone()) + .collect::>(), + ); + + self.update_idxs(); + + if self.is_end() { + self.stop.set(true); + } + + Ok(res.into_ref(vm).into_object()) + } + + fn is_end(&self) -> bool { + let cur = self.cur.get(); + self.idxs.borrow()[cur] == self.pools[cur].len() - 1 && cur == 0 + } + + fn update_idxs(&self) { + let lst_idx = &self.pools[self.cur.get()].len() - 1; + + if self.idxs.borrow()[self.cur.get()] == lst_idx { + if self.is_end() { + return; + } + self.idxs.borrow_mut()[self.cur.get()] = 0; + self.cur.set(self.cur.get() - 1); + self.update_idxs(); + } else { + self.idxs.borrow_mut()[self.cur.get()] += 1; + self.cur.set(self.idxs.borrow().len() - 1); + } + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + } + + #[pyclass(name = "combinations")] + #[derive(Debug)] + struct PyItertoolsCombinations { + pool: Vec, + indices: RefCell>, + r: Cell, + exhausted: Cell, + } + + impl PyValue for PyItertoolsCombinations { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "combinations") + } + } + + #[pyimpl] + impl PyItertoolsCombinations { + #[pyslot] + fn tp_new( + cls: PyClassRef, + iterable: PyObjectRef, + r: PyIntRef, + vm: &VirtualMachine, + ) -> PyResult> { + let iter = get_iter(vm, &iterable)?; + let pool = get_all(vm, &iter)?; + + let r = r.as_bigint(); + if r.is_negative() { + return Err(vm.new_value_error("r must be non-negative".to_owned())); + } + let r = r.to_usize().unwrap(); + + let n = pool.len(); + + PyItertoolsCombinations { + pool, + indices: RefCell::new((0..r).collect()), + r: Cell::new(r), + exhausted: Cell::new(r > n), + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + // stop signal + if self.exhausted.get() { + return Err(new_stop_iteration(vm)); + } + + let n = self.pool.len(); + let r = self.r.get(); + + if r == 0 { + self.exhausted.set(true); + return Ok(vm.ctx.new_tuple(vec![])); + } + + let res = PyTuple::from( + self.indices + .borrow() + .iter() + .map(|&i| self.pool[i].clone()) + .collect::>(), + ); + + let mut indices = self.indices.borrow_mut(); + + // Scan indices right-to-left until finding one that is not at its maximum (i + n - r). + let mut idx = r as isize - 1; + while idx >= 0 && indices[idx as usize] == idx as usize + n - r { + idx -= 1; + } + + // If no suitable index is found, then the indices are all at + // their maximum value and we're done. + if idx < 0 { + self.exhausted.set(true); + } else { + // Increment the current index which we know is not at its + // maximum. Then move back to the right setting each index + // to its lowest possible value (one higher than the index + // to its left -- this maintains the sort order invariant). + indices[idx as usize] += 1; + for j in idx as usize + 1..r { + indices[j] = indices[j - 1] + 1; + } + } + + Ok(res.into_ref(vm).into_object()) + } + } + + #[pyclass(name = "combinations_with_replacement")] + #[derive(Debug)] + struct PyItertoolsCombinationsWithReplacement { + pool: Vec, + indices: RefCell>, + r: Cell, + exhausted: Cell, + } + + impl PyValue for PyItertoolsCombinationsWithReplacement { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "combinations_with_replacement") + } + } + + #[pyimpl] + impl PyItertoolsCombinationsWithReplacement { + #[pyslot] + fn tp_new( + cls: PyClassRef, + iterable: PyObjectRef, + r: PyIntRef, + vm: &VirtualMachine, + ) -> PyResult> { + let iter = get_iter(vm, &iterable)?; + let pool = get_all(vm, &iter)?; + + let r = r.as_bigint(); + if r.is_negative() { + return Err(vm.new_value_error("r must be non-negative".to_owned())); + } + let r = r.to_usize().unwrap(); + + let n = pool.len(); + + PyItertoolsCombinationsWithReplacement { + pool, + indices: RefCell::new(vec![0; r]), + r: Cell::new(r), + exhausted: Cell::new(n == 0 && r > 0), + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + // stop signal + if self.exhausted.get() { + return Err(new_stop_iteration(vm)); + } + + let n = self.pool.len(); + let r = self.r.get(); + + if r == 0 { + self.exhausted.set(true); + return Ok(vm.ctx.new_tuple(vec![])); + } + + let mut indices = self.indices.borrow_mut(); + + let res = vm + .ctx + .new_tuple(indices.iter().map(|&i| self.pool[i].clone()).collect()); + + // Scan indices right-to-left until finding one that is not at its maximum (i + n - r). + let mut idx = r as isize - 1; + while idx >= 0 && indices[idx as usize] == n - 1 { + idx -= 1; + } + + // If no suitable index is found, then the indices are all at + // their maximum value and we're done. + if idx < 0 { + self.exhausted.set(true); + } else { + let index = indices[idx as usize] + 1; + + // Increment the current index which we know is not at its + // maximum. Then set all to the right to the same value. + for j in idx as usize..r { + indices[j as usize] = index as usize; + } + } + + Ok(res) + } + } + + #[pyclass(name = "permutations")] + #[derive(Debug)] + struct PyItertoolsPermutations { + pool: Vec, // Collected input iterable + indices: RefCell>, // One index per element in pool + cycles: RefCell>, // One rollover counter per element in the result + result: RefCell>>, // Indexes of the most recently returned result + r: Cell, // Size of result tuple + exhausted: Cell, // Set when the iterator is exhausted + } + + impl PyValue for PyItertoolsPermutations { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "permutations") + } + } + + #[pyimpl] + impl PyItertoolsPermutations { + #[pyslot] + fn tp_new( + cls: PyClassRef, + iterable: PyObjectRef, + r: OptionalOption, + vm: &VirtualMachine, + ) -> PyResult> { + let iter = get_iter(vm, &iterable)?; + let pool = get_all(vm, &iter)?; + + let n = pool.len(); + // If r is not provided, r == n. If provided, r must be a positive integer, or None. + // If None, it behaves the same as if it was not provided. + let r = match r.flat_option() { + Some(r) => { + let val = r + .payload::() + .ok_or_else(|| vm.new_type_error("Expected int as r".to_owned()))? + .as_bigint(); + + if val.is_negative() { + return Err(vm.new_value_error("r must be non-negative".to_owned())); + } + val.to_usize().unwrap() + } + None => n, + }; + + PyItertoolsPermutations { + pool, + indices: RefCell::new((0..n).collect()), + cycles: RefCell::new((0..r).map(|i| n - i).collect()), + result: RefCell::new(None), + r: Cell::new(r), + exhausted: Cell::new(r > n), + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + // stop signal + if self.exhausted.get() { + return Err(new_stop_iteration(vm)); + } + + let n = self.pool.len(); + let r = self.r.get(); + + if n == 0 { + self.exhausted.set(true); + return Ok(vm.ctx.new_tuple(vec![])); + } + + let result = &mut *self.result.borrow_mut(); + + if let Some(ref mut result) = result { + let mut indices = self.indices.borrow_mut(); + let mut cycles = self.cycles.borrow_mut(); + let mut sentinel = false; + + // Decrement rightmost cycle, moving leftward upon zero rollover + for i in (0..r).rev() { + cycles[i] -= 1; + + if cycles[i] == 0 { + // rotation: indices[i:] = indices[i+1:] + indices[i:i+1] + let index = indices[i]; + for j in i..n - 1 { + indices[j] = indices[j + i]; + } + indices[n - 1] = index; + cycles[i] = n - i; + } else { + let j = cycles[i]; + indices.swap(i, n - j); + + for k in i..r { + // start with i, the leftmost element that changed + // yield tuple(pool[k] for k in indices[:r]) + result[k] = indices[k]; + } + sentinel = true; + break; + } + } + if !sentinel { + self.exhausted.set(true); + return Err(new_stop_iteration(vm)); + } + } else { + // On the first pass, initialize result tuple using the indices + *result = Some((0..r).collect()); + } + + Ok(vm.ctx.new_tuple( + result + .as_ref() + .unwrap() + .iter() + .map(|&i| self.pool[i].clone()) + .collect(), + )) + } + } + + #[pyclass(name = "zip_longest")] + #[derive(Debug)] + struct PyItertoolsZipLongest { + iterators: Vec, + fillvalue: PyObjectRef, + numactive: Cell, + } + + impl PyValue for PyItertoolsZipLongest { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("itertools", "zip_longest") + } + } + + #[derive(FromArgs)] + struct ZiplongestArgs { + #[pyarg(keyword_only, optional = true)] + fillvalue: OptionalArg, + } + + #[pyimpl] + impl PyItertoolsZipLongest { + #[pyslot] + fn tp_new( + cls: PyClassRef, + iterables: Args, + args: ZiplongestArgs, + vm: &VirtualMachine, + ) -> PyResult> { + let fillvalue = match args.fillvalue.into_option() { + Some(i) => i, + None => vm.get_none(), + }; + + let iterators = iterables + .into_iter() + .map(|iterable| get_iter(vm, &iterable)) + .collect::, _>>()?; + + let numactive = Cell::new(iterators.len()); + + PyItertoolsZipLongest { + iterators, + fillvalue, + numactive, + } + .into_ref_with_type(vm, cls) + } + + #[pymethod(name = "__next__")] + fn next(&self, vm: &VirtualMachine) -> PyResult { + if self.iterators.is_empty() { + Err(new_stop_iteration(vm)) + } else { + let mut result: Vec = Vec::new(); + let mut numactive = self.numactive.get(); + + for idx in 0..self.iterators.len() { + let next_obj = match call_next(vm, &self.iterators[idx]) { + Ok(obj) => obj, + Err(err) => { + if !objtype::isinstance(&err, &vm.ctx.exceptions.stop_iteration) { + return Err(err); + } + numactive -= 1; + if numactive == 0 { + return Err(new_stop_iteration(vm)); + } + self.fillvalue.clone() + } + }; + result.push(next_obj); + } + Ok(vm.ctx.new_tuple(result)) + } + } + + #[pymethod(name = "__iter__")] + fn iter(zelf: PyRef) -> PyRef { + zelf + } } } - -pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { - let ctx = &vm.ctx; - - let accumulate = ctx.new_class("accumulate", ctx.object()); - PyItertoolsAccumulate::extend_class(ctx, &accumulate); - - let chain = PyItertoolsChain::make_class(ctx); - - let compress = PyItertoolsCompress::make_class(ctx); - - let combinations = ctx.new_class("combinations", ctx.object()); - PyItertoolsCombinations::extend_class(ctx, &combinations); - - let combinations_with_replacement = - ctx.new_class("combinations_with_replacement", ctx.object()); - PyItertoolsCombinationsWithReplacement::extend_class(ctx, &combinations_with_replacement); - - let count = ctx.new_class("count", ctx.object()); - PyItertoolsCount::extend_class(ctx, &count); - - let cycle = ctx.new_class("cycle", ctx.object()); - PyItertoolsCycle::extend_class(ctx, &cycle); - - let dropwhile = ctx.new_class("dropwhile", ctx.object()); - PyItertoolsDropwhile::extend_class(ctx, &dropwhile); - - let islice = PyItertoolsIslice::make_class(ctx); - - let filterfalse = ctx.new_class("filterfalse", ctx.object()); - PyItertoolsFilterFalse::extend_class(ctx, &filterfalse); - - let permutations = ctx.new_class("permutations", ctx.object()); - PyItertoolsPermutations::extend_class(ctx, &permutations); - - let product = ctx.new_class("product", ctx.object()); - PyItertoolsProduct::extend_class(ctx, &product); - - let repeat = ctx.new_class("repeat", ctx.object()); - PyItertoolsRepeat::extend_class(ctx, &repeat); - - let starmap = PyItertoolsStarmap::make_class(ctx); - - let takewhile = ctx.new_class("takewhile", ctx.object()); - PyItertoolsTakewhile::extend_class(ctx, &takewhile); - - let tee = ctx.new_class("tee", ctx.object()); - PyItertoolsTee::extend_class(ctx, &tee); - - let zip_longest = ctx.new_class("zip_longest", ctx.object()); - PyItertoolsZiplongest::extend_class(ctx, &zip_longest); - - py_module!(vm, "itertools", { - "accumulate" => accumulate, - "chain" => chain, - "compress" => compress, - "combinations" => combinations, - "combinations_with_replacement" => combinations_with_replacement, - "count" => count, - "cycle" => cycle, - "dropwhile" => dropwhile, - "islice" => islice, - "filterfalse" => filterfalse, - "repeat" => repeat, - "starmap" => starmap, - "takewhile" => takewhile, - "tee" => tee, - "permutations" => permutations, - "product" => product, - "zip_longest" => zip_longest, - }) -} diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs index 0de717689..3c6f97cd9 100644 --- a/vm/src/stdlib/os.rs +++ b/vm/src/stdlib/os.rs @@ -1,4 +1,3 @@ -use num_cpus; use std::cell::{Cell, RefCell}; use std::ffi; use std::fs::File; @@ -13,8 +12,6 @@ use std::{env, fs}; use bitflags::bitflags; #[cfg(unix)] -use exitcode; -#[cfg(unix)] use nix::errno::Errno; #[cfg(all(unix, not(target_os = "redox")))] use nix::pty::openpty; @@ -24,8 +21,6 @@ use nix::unistd::{self, Gid, Pid, Uid}; use std::os::unix::io::RawFd; #[cfg(windows)] use std::os::windows::io::RawHandle; -#[cfg(unix)] -use uname; use super::errno::errors; use crate::exceptions::PyBaseExceptionRef; diff --git a/vm/src/stdlib/platform.rs b/vm/src/stdlib/platform.rs index 26f494927..af22e5497 100644 --- a/vm/src/stdlib/platform.rs +++ b/vm/src/stdlib/platform.rs @@ -1,39 +1,37 @@ -use crate::pyobject::PyObjectRef; -use crate::version; -use crate::vm::VirtualMachine; +pub(crate) use decl::make_module; -pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { - let ctx = &vm.ctx; - py_module!(vm, "platform", { - "python_branch" => ctx.new_function(platform_python_branch), - "python_build" => ctx.new_function(platform_python_build), - "python_compiler" => ctx.new_function(platform_python_compiler), - "python_implementation" => ctx.new_function(platform_python_implementation), - "python_revision" => ctx.new_function(platform_python_revision), - "python_version" => ctx.new_function(platform_python_version), - }) -} +#[pymodule(name = "platform")] +mod decl { + use crate::version; + use crate::vm::VirtualMachine; -fn platform_python_implementation(_vm: &VirtualMachine) -> String { - "RustPython".to_owned() -} + #[pyfunction] + fn python_implementation(_vm: &VirtualMachine) -> String { + "RustPython".to_owned() + } -fn platform_python_version(_vm: &VirtualMachine) -> String { - version::get_version_number() -} + #[pyfunction] + fn python_version(_vm: &VirtualMachine) -> String { + version::get_version_number() + } -fn platform_python_compiler(_vm: &VirtualMachine) -> String { - version::get_compiler() -} + #[pyfunction] + fn python_compiler(_vm: &VirtualMachine) -> String { + version::get_compiler() + } -fn platform_python_build(_vm: &VirtualMachine) -> (String, String) { - (version::get_git_identifier(), version::get_git_datetime()) -} + #[pyfunction] + fn python_build(_vm: &VirtualMachine) -> (String, String) { + (version::get_git_identifier(), version::get_git_datetime()) + } -fn platform_python_branch(_vm: &VirtualMachine) -> String { - version::get_git_branch() -} + #[pyfunction] + fn python_branch(_vm: &VirtualMachine) -> String { + version::get_git_branch() + } -fn platform_python_revision(_vm: &VirtualMachine) -> String { - version::get_git_revision() + #[pyfunction] + fn python_revision(_vm: &VirtualMachine) -> String { + version::get_git_revision() + } } diff --git a/vm/src/stdlib/pystruct.rs b/vm/src/stdlib/pystruct.rs index f31aa36c7..320d9875a 100644 --- a/vm/src/stdlib/pystruct.rs +++ b/vm/src/stdlib/pystruct.rs @@ -9,236 +9,242 @@ * https://docs.rs/byteorder/1.2.6/byteorder/ */ -use byteorder::{ReadBytesExt, WriteBytesExt}; -use num_bigint::BigInt; -use num_traits::ToPrimitive; -use std::cmp; -use std::io::{Cursor, Read, Write}; -use std::iter::Peekable; - -use crate::exceptions::PyBaseExceptionRef; -use crate::function::Args; -use crate::obj::{ - objbool::IntoPyBool, objbytes::PyBytesRef, objstr::PyString, objstr::PyStringRef, - objtuple::PyTuple, objtype::PyClassRef, -}; -use crate::pyobject::{Either, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject}; +use crate::pyobject::PyObjectRef; use crate::VirtualMachine; -#[derive(Debug)] -enum Endianness { - Native, - Little, - Big, - Network, -} +#[pymodule] +mod _struct { + use byteorder::{ReadBytesExt, WriteBytesExt}; + use num_bigint::BigInt; + use num_traits::ToPrimitive; + use std::io::{Cursor, Read, Write}; + use std::iter::Peekable; -#[derive(Debug)] -struct FormatCode { - repeat: u32, - code: char, -} + use crate::exceptions::PyBaseExceptionRef; + use crate::function::Args; + use crate::obj::{ + objbool::IntoPyBool, objbytes::PyBytesRef, objstr::PyString, objstr::PyStringRef, + objtuple::PyTuple, objtype::PyClassRef, + }; + use crate::pyobject::{ + Either, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, + }; + use crate::VirtualMachine; -impl FormatCode { - fn unit_size(&self) -> usize { - match self.code { - 'x' | 'c' | 'b' | 'B' | '?' | 's' | 'p' => 1, - 'h' | 'H' => 2, - 'i' | 'l' | 'I' | 'L' | 'f' => 4, - 'q' | 'Q' | 'd' => 8, - 'n' | 'N' | 'P' => std::mem::size_of::(), - c => { - panic!("Unsupported format code {:?}", c); + #[derive(Debug)] + enum Endianness { + Native, + Little, + Big, + Network, + } + + #[derive(Debug)] + struct FormatCode { + repeat: u32, + code: char, + } + + impl FormatCode { + fn unit_size(&self) -> usize { + match self.code { + 'x' | 'c' | 'b' | 'B' | '?' | 's' | 'p' => 1, + 'h' | 'H' => 2, + 'i' | 'l' | 'I' | 'L' | 'f' => 4, + 'q' | 'Q' | 'd' => 8, + 'n' | 'N' | 'P' => std::mem::size_of::(), + c => { + panic!("Unsupported format code {:?}", c); + } + } + } + + fn size(&self) -> usize { + self.unit_size() * self.repeat as usize + } + + fn arg_count(&self) -> usize { + match self.code { + 'x' => 0, + 's' | 'p' => 1, + _ => self.repeat as usize, } } } - fn size(&self) -> usize { - self.unit_size() * self.repeat as usize + #[derive(Debug)] + struct FormatSpec { + endianness: Endianness, + codes: Vec, } - fn arg_count(&self) -> usize { - match self.code { - 'x' => 0, - 's' | 'p' => 1, - _ => self.repeat as usize, - } - } -} + impl FormatSpec { + fn parse(fmt: &str) -> Result { + let mut chars = fmt.chars().peekable(); -#[derive(Debug)] -struct FormatSpec { - endianness: Endianness, - codes: Vec, -} + // First determine "<", ">","!" or "=" + let endianness = parse_endiannes(&mut chars); -impl FormatSpec { - fn parse(fmt: &str) -> Result { - let mut chars = fmt.chars().peekable(); + // Now, analyze struct string furter: + let codes = parse_format_codes(&mut chars)?; - // First determine "<", ">","!" or "=" - let endianness = parse_endiannes(&mut chars); - - // Now, analyze struct string furter: - let codes = parse_format_codes(&mut chars)?; - - Ok(FormatSpec { endianness, codes }) - } - - fn pack(&self, args: &[PyObjectRef], vm: &VirtualMachine) -> PyResult> { - let arg_count: usize = self.codes.iter().map(|c| c.arg_count()).sum(); - if arg_count != args.len() { - return Err(new_struct_error( - vm, - format!( - "pack expected {} items for packing (got {})", - self.codes.len(), - args.len() - ), - )); + Ok(FormatSpec { endianness, codes }) } - // Create data vector: - let mut data = Vec::::new(); - let mut arg_idx = 0; - // Loop over all opcodes: - for code in self.codes.iter() { - debug!("code: {:?}", code); - let pack_item = match self.endianness { - Endianness::Little => pack_item::, - Endianness::Big => pack_item::, - Endianness::Network => pack_item::, - Endianness::Native => pack_item::, - }; - arg_idx += pack_item(vm, code, &args[arg_idx..], &mut data)?; + fn pack(&self, args: &[PyObjectRef], vm: &VirtualMachine) -> PyResult> { + let arg_count: usize = self.codes.iter().map(|c| c.arg_count()).sum(); + if arg_count != args.len() { + return Err(new_struct_error( + vm, + format!( + "pack expected {} items for packing (got {})", + self.codes.len(), + args.len() + ), + )); + } + + // Create data vector: + let mut data = Vec::::new(); + let mut arg_idx = 0; + // Loop over all opcodes: + for code in self.codes.iter() { + debug!("code: {:?}", code); + let pack_item = match self.endianness { + Endianness::Little => pack_item::, + Endianness::Big => pack_item::, + Endianness::Network => pack_item::, + Endianness::Native => pack_item::, + }; + arg_idx += pack_item(vm, code, &args[arg_idx..], &mut data)?; + } + + Ok(data) } - Ok(data) - } + fn unpack(&self, data: &[u8], vm: &VirtualMachine) -> PyResult { + if self.size() != data.len() { + return Err(new_struct_error( + vm, + format!("unpack requires a buffer of {} bytes", self.size()), + )); + } - fn unpack(&self, data: &[u8], vm: &VirtualMachine) -> PyResult { - if self.size() != data.len() { - return Err(new_struct_error( - vm, - format!("unpack requires a buffer of {} bytes", self.size()), - )); - } - - let mut rdr = Cursor::new(data); - let mut items = vec![]; - for code in &self.codes { - debug!("unpack code: {:?}", code); - match self.endianness { - Endianness::Little => { - unpack_code::(vm, &code, &mut rdr, &mut items)? - } - Endianness::Big => { - unpack_code::(vm, &code, &mut rdr, &mut items)? - } - Endianness::Network => { - unpack_code::(vm, &code, &mut rdr, &mut items)? - } - Endianness::Native => { - unpack_code::(vm, &code, &mut rdr, &mut items)? - } - }; - } - - Ok(PyTuple::from(items)) - } - - fn size(&self) -> usize { - self.codes.iter().map(FormatCode::size).sum() - } -} - -/// Parse endianness -/// See also: https://docs.python.org/3/library/struct.html?highlight=struct#byte-order-size-and-alignment -fn parse_endiannes(chars: &mut Peekable) -> Endianness -where - I: Sized + Iterator, -{ - match chars.peek() { - Some('@') => { - chars.next().unwrap(); - Endianness::Native - } - Some('=') => { - chars.next().unwrap(); - Endianness::Native - } - Some('<') => { - chars.next().unwrap(); - Endianness::Little - } - Some('>') => { - chars.next().unwrap(); - Endianness::Big - } - Some('!') => { - chars.next().unwrap(); - Endianness::Network - } - _ => Endianness::Native, - } -} - -fn parse_format_codes(chars: &mut Peekable) -> Result, String> -where - I: Sized + Iterator, -{ - let mut codes = vec![]; - while chars.peek().is_some() { - // determine repeat operator: - let repeat = match chars.peek() { - Some('0'..='9') => { - let mut repeat = 0; - while let Some('0'..='9') = chars.peek() { - if let Some(c) = chars.next() { - let current_digit = c.to_digit(10).unwrap(); - repeat = repeat * 10 + current_digit; + let mut rdr = Cursor::new(data); + let mut items = vec![]; + for code in &self.codes { + debug!("unpack code: {:?}", code); + match self.endianness { + Endianness::Little => { + unpack_code::(vm, &code, &mut rdr, &mut items)? } - } - repeat + Endianness::Big => { + unpack_code::(vm, &code, &mut rdr, &mut items)? + } + Endianness::Network => { + unpack_code::(vm, &code, &mut rdr, &mut items)? + } + Endianness::Native => { + unpack_code::(vm, &code, &mut rdr, &mut items)? + } + }; } - _ => 1, - }; - // determine format char: - let c = chars.next(); - match c { - Some(c) if is_supported_format_character(c) => { - codes.push(FormatCode { repeat, code: c }) - } - _ => return Err(format!("Illegal format code {:?}", c)), + Ok(PyTuple::from(items)) + } + + fn size(&self) -> usize { + self.codes.iter().map(FormatCode::size).sum() } } - Ok(codes) -} - -fn is_supported_format_character(c: char) -> bool { - match c { - 'x' | 'c' | 'b' | 'B' | '?' | 'h' | 'H' | 'i' | 'I' | 'l' | 'L' | 'q' | 'Q' | 'n' | 'N' - | 'f' | 'd' | 's' | 'p' | 'P' => true, - _ => false, + /// Parse endianness + /// See also: https://docs.python.org/3/library/struct.html?highlight=struct#byte-order-size-and-alignment + fn parse_endiannes(chars: &mut Peekable) -> Endianness + where + I: Sized + Iterator, + { + match chars.peek() { + Some('@') => { + chars.next().unwrap(); + Endianness::Native + } + Some('=') => { + chars.next().unwrap(); + Endianness::Native + } + Some('<') => { + chars.next().unwrap(); + Endianness::Little + } + Some('>') => { + chars.next().unwrap(); + Endianness::Big + } + Some('!') => { + chars.next().unwrap(); + Endianness::Network + } + _ => Endianness::Native, + } } -} -fn get_int_or_index(vm: &VirtualMachine, arg: &PyObjectRef) -> PyResult -where - T: TryFromObject, -{ - match vm.to_index(arg) { - Some(index) => Ok(T::try_from_object(vm, index?.into_object())?), - None => Err(new_struct_error( - vm, - "required argument is not an integer".to_owned(), - )), + fn parse_format_codes(chars: &mut Peekable) -> Result, String> + where + I: Sized + Iterator, + { + let mut codes = vec![]; + while chars.peek().is_some() { + // determine repeat operator: + let repeat = match chars.peek() { + Some('0'..='9') => { + let mut repeat = 0; + while let Some('0'..='9') = chars.peek() { + if let Some(c) = chars.next() { + let current_digit = c.to_digit(10).unwrap(); + repeat = repeat * 10 + current_digit; + } + } + repeat + } + _ => 1, + }; + + // determine format char: + let c = chars.next(); + match c { + Some(c) if is_supported_format_character(c) => { + codes.push(FormatCode { repeat, code: c }) + } + _ => return Err(format!("Illegal format code {:?}", c)), + } + } + + Ok(codes) } -} -macro_rules! make_pack_no_endianess { + fn is_supported_format_character(c: char) -> bool { + match c { + 'x' | 'c' | 'b' | 'B' | '?' | 'h' | 'H' | 'i' | 'I' | 'l' | 'L' | 'q' | 'Q' | 'n' + | 'N' | 'f' | 'd' | 's' | 'p' | 'P' => true, + _ => false, + } + } + + fn get_int_or_index(vm: &VirtualMachine, arg: &PyObjectRef) -> PyResult + where + T: TryFromObject, + { + match vm.to_index(arg) { + Some(index) => Ok(T::try_from_object(vm, index?.into_object())?), + None => Err(new_struct_error( + vm, + "required argument is not an integer".to_owned(), + )), + } + } + + macro_rules! make_pack_no_endianess { ($T:ty) => { paste::item! { fn [](vm: &VirtualMachine, arg: &PyObjectRef, data: &mut dyn Write) -> PyResult<()> { @@ -249,7 +255,7 @@ macro_rules! make_pack_no_endianess { }; } -macro_rules! make_pack_with_endianess_int { + macro_rules! make_pack_with_endianess_int { ($T:ty) => { paste::item! { fn [](vm: &VirtualMachine, arg: &PyObjectRef, data: &mut dyn Write) -> PyResult<()> @@ -263,7 +269,7 @@ macro_rules! make_pack_with_endianess_int { }; } -macro_rules! make_pack_with_endianess { + macro_rules! make_pack_with_endianess { ($T:ty) => { paste::item! { fn [](vm: &VirtualMachine, arg: &PyObjectRef, data: &mut dyn Write) -> PyResult<()> @@ -278,452 +284,455 @@ macro_rules! make_pack_with_endianess { }; } -make_pack_no_endianess!(i8); -make_pack_no_endianess!(u8); -make_pack_with_endianess_int!(i16); -make_pack_with_endianess_int!(u16); -make_pack_with_endianess_int!(i32); -make_pack_with_endianess_int!(u32); -make_pack_with_endianess_int!(i64); -make_pack_with_endianess_int!(u64); -make_pack_with_endianess!(f32); -make_pack_with_endianess!(f64); + make_pack_no_endianess!(i8); + make_pack_no_endianess!(u8); + make_pack_with_endianess_int!(i16); + make_pack_with_endianess_int!(u16); + make_pack_with_endianess_int!(i32); + make_pack_with_endianess_int!(u32); + make_pack_with_endianess_int!(i64); + make_pack_with_endianess_int!(u64); + make_pack_with_endianess!(f32); + make_pack_with_endianess!(f64); -fn pack_bool(vm: &VirtualMachine, arg: &PyObjectRef, data: &mut dyn Write) -> PyResult<()> { - let v = if IntoPyBool::try_from_object(vm, arg.clone())?.to_bool() { - 1 - } else { - 0 - }; - data.write_u8(v).unwrap(); - Ok(()) -} - -fn pack_isize( - vm: &VirtualMachine, - arg: &PyObjectRef, - data: &mut dyn Write, -) -> PyResult<()> -where - Endianness: byteorder::ByteOrder, -{ - let v: isize = get_int_or_index(vm, arg)?; - match std::mem::size_of::() { - 8 => data.write_i64::(v as i64).unwrap(), - 4 => data.write_i32::(v as i32).unwrap(), - _ => unreachable!("unexpected architecture"), - } - Ok(()) -} - -fn pack_usize( - vm: &VirtualMachine, - arg: &PyObjectRef, - data: &mut dyn Write, -) -> PyResult<()> -where - Endianness: byteorder::ByteOrder, -{ - let v: usize = get_int_or_index(vm, arg)?; - match std::mem::size_of::() { - 8 => data.write_u64::(v as u64).unwrap(), - 4 => data.write_u32::(v as u32).unwrap(), - _ => unreachable!("unexpected architecture"), - } - Ok(()) -} - -fn pack_string( - vm: &VirtualMachine, - arg: &PyObjectRef, - data: &mut dyn Write, - length: usize, -) -> PyResult<()> { - let mut v = PyBytesRef::try_from_object(vm, arg.clone())? - .get_value() - .to_vec(); - v.resize(length, 0); - match data.write_all(&v) { - Ok(_) => Ok(()), - Err(e) => Err(new_struct_error(vm, format!("{:?}", e))), - } -} - -fn pack_pascal( - vm: &VirtualMachine, - arg: &PyObjectRef, - data: &mut dyn Write, - length: usize, -) -> PyResult<()> { - let mut v = PyBytesRef::try_from_object(vm, arg.clone())? - .get_value() - .to_vec(); - let string_length = cmp::min(cmp::min(v.len(), 255), length - 1); - data.write_u8(string_length as u8).unwrap(); - v.resize(length - 1, 0); - match data.write_all(&v) { - Ok(_) => Ok(()), - Err(e) => Err(new_struct_error(vm, format!("{:?}", e))), - } -} - -fn pack_char(vm: &VirtualMachine, arg: &PyObjectRef, data: &mut dyn Write) -> PyResult<()> { - let v = PyBytesRef::try_from_object(vm, arg.clone())?; - if v.len() == 1 { - data.write_u8(v[0]).unwrap(); - Ok(()) - } else { - Err(new_struct_error( - vm, - "char format requires a bytes object of length 1".to_owned(), - )) - } -} - -fn pack_item( - vm: &VirtualMachine, - code: &FormatCode, - args: &[PyObjectRef], - data: &mut dyn Write, -) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - let pack = match code.code { - 'c' => pack_char, - 'b' => pack_i8, - 'B' => pack_u8, - '?' => pack_bool, - 'h' => pack_i16::, - 'H' => pack_u16::, - 'i' | 'l' => pack_i32::, - 'I' | 'L' => pack_u32::, - 'q' => pack_i64::, - 'Q' => pack_u64::, - 'n' => pack_isize::, - 'N' | 'P' => pack_usize::, - 'f' => pack_f32::, - 'd' => pack_f64::, - 's' => { - pack_string(vm, &args[0], data, code.repeat as usize)?; - return Ok(1); - } - 'p' => { - pack_pascal(vm, &args[0], data, code.repeat as usize)?; - return Ok(1); - } - 'x' => { - for _ in 0..code.repeat as usize { - data.write_u8(0).unwrap(); - } - return Ok(0); - } - c => { - panic!("Unsupported format code {:?}", c); - } - }; - - for arg in args.iter().take(code.repeat as usize) { - pack(vm, arg, data)?; - } - Ok(code.repeat as usize) -} - -fn struct_pack(fmt: PyStringRef, args: Args, vm: &VirtualMachine) -> PyResult> { - let format_spec = FormatSpec::parse(fmt.as_str()).map_err(|e| new_struct_error(vm, e))?; - format_spec.pack(args.as_ref(), vm) -} - -#[inline] -fn unpack(vm: &VirtualMachine, rdr: &mut dyn Read, read: F, transform: G) -> PyResult -where - F: Fn(&mut dyn Read) -> std::io::Result, - G: Fn(T) -> PyResult, -{ - match read(rdr) { - Ok(v) => transform(v), - Err(_) => Err(new_struct_error( - vm, - format!( - "unpack requires a buffer of {} bytes", - std::mem::size_of::() - ), - )), - } -} - -#[inline] -fn unpack_int(vm: &VirtualMachine, rdr: &mut dyn Read, read: F) -> PyResult -where - F: Fn(&mut dyn Read) -> std::io::Result, - T: Into + ToPrimitive, -{ - unpack(vm, rdr, read, |v| Ok(vm.ctx.new_int(v))) -} - -fn unpack_bool(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult { - unpack(vm, rdr, |rdr| rdr.read_u8(), |v| Ok(vm.ctx.new_bool(v > 0))) -} - -fn unpack_i8(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult { - unpack_int(vm, rdr, |rdr| rdr.read_i8()) -} - -fn unpack_u8(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult { - unpack_int(vm, rdr, |rdr| rdr.read_u8()) -} - -fn unpack_i16(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - unpack_int(vm, rdr, |rdr| rdr.read_i16::()) -} - -fn unpack_u16(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - unpack_int(vm, rdr, |rdr| rdr.read_u16::()) -} - -fn unpack_i32(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - unpack_int(vm, rdr, |rdr| rdr.read_i32::()) -} - -fn unpack_u32(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - unpack_int(vm, rdr, |rdr| rdr.read_u32::()) -} - -fn unpack_i64(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - unpack_int(vm, rdr, |rdr| rdr.read_i64::()) -} - -fn unpack_u64(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - unpack_int(vm, rdr, |rdr| rdr.read_u64::()) -} - -fn unpack_isize(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - match std::mem::size_of::() { - 8 => unpack_i64::(vm, rdr), - 4 => unpack_i32::(vm, rdr), - _ => unreachable!("unexpected architecture"), - } -} - -fn unpack_usize(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - match std::mem::size_of::() { - 8 => unpack_u64::(vm, rdr), - 4 => unpack_u32::(vm, rdr), - _ => unreachable!("unexpected architecture"), - } -} - -fn unpack_f32(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - unpack( - vm, - rdr, - |rdr| rdr.read_f32::(), - |v| Ok(vm.ctx.new_float(f64::from(v))), - ) -} - -fn unpack_f64(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult -where - Endianness: byteorder::ByteOrder, -{ - unpack( - vm, - rdr, - |rdr| rdr.read_f64::(), - |v| Ok(vm.ctx.new_float(v)), - ) -} - -fn unpack_empty(_vm: &VirtualMachine, rdr: &mut dyn Read, length: u32) { - let mut handle = rdr.take(length as u64); - let mut buf: Vec = Vec::new(); - let _ = handle.read_to_end(&mut buf); -} - -fn unpack_char(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult { - unpack_string(vm, rdr, 1) -} - -fn unpack_string(vm: &VirtualMachine, rdr: &mut dyn Read, length: u32) -> PyResult { - let mut handle = rdr.take(length as u64); - let mut buf: Vec = Vec::new(); - handle.read_to_end(&mut buf).map_err(|_| { - new_struct_error(vm, format!("unpack requires a buffer of {} bytes", length,)) - })?; - Ok(vm.ctx.new_bytes(buf)) -} - -fn unpack_pascal(vm: &VirtualMachine, rdr: &mut dyn Read, length: u32) -> PyResult { - let mut handle = rdr.take(length as u64); - let mut buf: Vec = Vec::new(); - handle.read_to_end(&mut buf).map_err(|_| { - new_struct_error(vm, format!("unpack requires a buffer of {} bytes", length,)) - })?; - let string_length = buf[0] as usize; - Ok(vm.ctx.new_bytes(buf[1..=string_length].to_vec())) -} - -fn struct_unpack(fmt: PyStringRef, buffer: PyBytesRef, vm: &VirtualMachine) -> PyResult { - let fmt_str = fmt.as_str(); - let format_spec = FormatSpec::parse(fmt_str).map_err(|e| new_struct_error(vm, e))?; - format_spec.unpack(buffer.get_value(), vm) -} - -fn unpack_code( - vm: &VirtualMachine, - code: &FormatCode, - rdr: &mut dyn Read, - items: &mut Vec, -) -> PyResult<()> -where - Endianness: byteorder::ByteOrder, -{ - let unpack = match code.code { - 'b' => unpack_i8, - 'B' => unpack_u8, - 'c' => unpack_char, - '?' => unpack_bool, - 'h' => unpack_i16::, - 'H' => unpack_u16::, - 'i' | 'l' => unpack_i32::, - 'I' | 'L' => unpack_u32::, - 'q' => unpack_i64::, - 'Q' => unpack_u64::, - 'n' => unpack_isize::, - 'N' => unpack_usize::, - 'P' => unpack_usize::, // FIXME: native-only - 'f' => unpack_f32::, - 'd' => unpack_f64::, - 'x' => { - unpack_empty(vm, rdr, code.repeat); - return Ok(()); - } - 's' => { - items.push(unpack_string(vm, rdr, code.repeat)?); - return Ok(()); - } - 'p' => { - items.push(unpack_pascal(vm, rdr, code.repeat)?); - return Ok(()); - } - c => { - panic!("Unsupported format code {:?}", c); - } - }; - for _ in 0..code.repeat { - items.push(unpack(vm, rdr)?); - } - Ok(()) -} - -fn struct_calcsize(fmt: Either, vm: &VirtualMachine) -> PyResult { - // FIXME: the given fmt must be parsed as ascii string - // https://github.com/RustPython/RustPython/pull/1792#discussion_r387340905 - let parsed = match fmt { - Either::A(string) => FormatSpec::parse(string.as_str()), - Either::B(bytes) => FormatSpec::parse(std::str::from_utf8(&bytes).unwrap()), - }; - let format_spec = parsed.map_err(|e| new_struct_error(vm, e))?; - Ok(format_spec.size()) -} - -#[pyclass(name = "Struct")] -#[derive(Debug)] -struct PyStruct { - spec: FormatSpec, - fmt_str: PyStringRef, -} - -impl PyValue for PyStruct { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("_struct", "Struct") - } -} - -#[pyimpl] -impl PyStruct { - #[pyslot] - fn tp_new( - cls: PyClassRef, - fmt: Either, - vm: &VirtualMachine, - ) -> PyResult> { - let fmt_str = match fmt { - Either::A(s) => s, - Either::B(b) => PyString::from(std::str::from_utf8(b.get_value()).unwrap()) - .into_ref_with_type(vm, vm.ctx.str_type())?, + fn pack_bool(vm: &VirtualMachine, arg: &PyObjectRef, data: &mut dyn Write) -> PyResult<()> { + let v = if IntoPyBool::try_from_object(vm, arg.clone())?.to_bool() { + 1 + } else { + 0 }; - let spec = FormatSpec::parse(fmt_str.as_str()).map_err(|e| new_struct_error(vm, e))?; - PyStruct { spec, fmt_str }.into_ref_with_type(vm, cls) + data.write_u8(v).unwrap(); + Ok(()) } - #[pyproperty] - fn format(&self) -> PyStringRef { - self.fmt_str.clone() - } - #[pyproperty] - fn size(&self) -> usize { - self.spec.size() + fn pack_isize( + vm: &VirtualMachine, + arg: &PyObjectRef, + data: &mut dyn Write, + ) -> PyResult<()> + where + Endianness: byteorder::ByteOrder, + { + let v: isize = get_int_or_index(vm, arg)?; + match std::mem::size_of::() { + 8 => data.write_i64::(v as i64).unwrap(), + 4 => data.write_i32::(v as i32).unwrap(), + _ => unreachable!("unexpected architecture"), + } + Ok(()) } - #[pymethod] - fn pack(&self, args: Args, vm: &VirtualMachine) -> PyResult> { - self.spec.pack(args.as_ref(), vm) + fn pack_usize( + vm: &VirtualMachine, + arg: &PyObjectRef, + data: &mut dyn Write, + ) -> PyResult<()> + where + Endianness: byteorder::ByteOrder, + { + let v: usize = get_int_or_index(vm, arg)?; + match std::mem::size_of::() { + 8 => data.write_u64::(v as u64).unwrap(), + 4 => data.write_u32::(v as u32).unwrap(), + _ => unreachable!("unexpected architecture"), + } + Ok(()) } - #[pymethod] - fn unpack(&self, data: PyBytesRef, vm: &VirtualMachine) -> PyResult { - self.spec.unpack(data.get_value(), vm) + + fn pack_string( + vm: &VirtualMachine, + arg: &PyObjectRef, + data: &mut dyn Write, + length: usize, + ) -> PyResult<()> { + let mut v = PyBytesRef::try_from_object(vm, arg.clone())? + .get_value() + .to_vec(); + v.resize(length, 0); + match data.write_all(&v) { + Ok(_) => Ok(()), + Err(e) => Err(new_struct_error(vm, format!("{:?}", e))), + } + } + + fn pack_pascal( + vm: &VirtualMachine, + arg: &PyObjectRef, + data: &mut dyn Write, + length: usize, + ) -> PyResult<()> { + let mut v = PyBytesRef::try_from_object(vm, arg.clone())? + .get_value() + .to_vec(); + let string_length = std::cmp::min(std::cmp::min(v.len(), 255), length - 1); + data.write_u8(string_length as u8).unwrap(); + v.resize(length - 1, 0); + match data.write_all(&v) { + Ok(_) => Ok(()), + Err(e) => Err(new_struct_error(vm, format!("{:?}", e))), + } + } + + fn pack_char(vm: &VirtualMachine, arg: &PyObjectRef, data: &mut dyn Write) -> PyResult<()> { + let v = PyBytesRef::try_from_object(vm, arg.clone())?; + if v.len() == 1 { + data.write_u8(v[0]).unwrap(); + Ok(()) + } else { + Err(new_struct_error( + vm, + "char format requires a bytes object of length 1".to_owned(), + )) + } + } + + fn pack_item( + vm: &VirtualMachine, + code: &FormatCode, + args: &[PyObjectRef], + data: &mut dyn Write, + ) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + let pack = match code.code { + 'c' => pack_char, + 'b' => pack_i8, + 'B' => pack_u8, + '?' => pack_bool, + 'h' => pack_i16::, + 'H' => pack_u16::, + 'i' | 'l' => pack_i32::, + 'I' | 'L' => pack_u32::, + 'q' => pack_i64::, + 'Q' => pack_u64::, + 'n' => pack_isize::, + 'N' | 'P' => pack_usize::, + 'f' => pack_f32::, + 'd' => pack_f64::, + 's' => { + pack_string(vm, &args[0], data, code.repeat as usize)?; + return Ok(1); + } + 'p' => { + pack_pascal(vm, &args[0], data, code.repeat as usize)?; + return Ok(1); + } + 'x' => { + for _ in 0..code.repeat as usize { + data.write_u8(0).unwrap(); + } + return Ok(0); + } + c => { + panic!("Unsupported format code {:?}", c); + } + }; + + for arg in args.iter().take(code.repeat as usize) { + pack(vm, arg, data)?; + } + Ok(code.repeat as usize) + } + + #[pyfunction] + fn pack(fmt: PyStringRef, args: Args, vm: &VirtualMachine) -> PyResult> { + let format_spec = FormatSpec::parse(fmt.as_str()).map_err(|e| new_struct_error(vm, e))?; + format_spec.pack(args.as_ref(), vm) + } + + #[inline] + fn _unpack(vm: &VirtualMachine, rdr: &mut dyn Read, read: F, transform: G) -> PyResult + where + F: Fn(&mut dyn Read) -> std::io::Result, + G: Fn(T) -> PyResult, + { + match read(rdr) { + Ok(v) => transform(v), + Err(_) => Err(new_struct_error( + vm, + format!( + "unpack requires a buffer of {} bytes", + std::mem::size_of::() + ), + )), + } + } + + #[inline] + fn unpack_int(vm: &VirtualMachine, rdr: &mut dyn Read, read: F) -> PyResult + where + F: Fn(&mut dyn Read) -> std::io::Result, + T: Into + ToPrimitive, + { + _unpack(vm, rdr, read, |v| Ok(vm.ctx.new_int(v))) + } + + fn unpack_bool(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult { + _unpack(vm, rdr, |rdr| rdr.read_u8(), |v| Ok(vm.ctx.new_bool(v > 0))) + } + + fn unpack_i8(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult { + unpack_int(vm, rdr, |rdr| rdr.read_i8()) + } + + fn unpack_u8(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult { + unpack_int(vm, rdr, |rdr| rdr.read_u8()) + } + + fn unpack_i16(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + unpack_int(vm, rdr, |rdr| rdr.read_i16::()) + } + + fn unpack_u16(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + unpack_int(vm, rdr, |rdr| rdr.read_u16::()) + } + + fn unpack_i32(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + unpack_int(vm, rdr, |rdr| rdr.read_i32::()) + } + + fn unpack_u32(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + unpack_int(vm, rdr, |rdr| rdr.read_u32::()) + } + + fn unpack_i64(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + unpack_int(vm, rdr, |rdr| rdr.read_i64::()) + } + + fn unpack_u64(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + unpack_int(vm, rdr, |rdr| rdr.read_u64::()) + } + + fn unpack_isize(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + match std::mem::size_of::() { + 8 => unpack_i64::(vm, rdr), + 4 => unpack_i32::(vm, rdr), + _ => unreachable!("unexpected architecture"), + } + } + + fn unpack_usize(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + match std::mem::size_of::() { + 8 => unpack_u64::(vm, rdr), + 4 => unpack_u32::(vm, rdr), + _ => unreachable!("unexpected architecture"), + } + } + + fn unpack_f32(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + _unpack( + vm, + rdr, + |rdr| rdr.read_f32::(), + |v| Ok(vm.ctx.new_float(f64::from(v))), + ) + } + + fn unpack_f64(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult + where + Endianness: byteorder::ByteOrder, + { + _unpack( + vm, + rdr, + |rdr| rdr.read_f64::(), + |v| Ok(vm.ctx.new_float(v)), + ) + } + + fn unpack_empty(_vm: &VirtualMachine, rdr: &mut dyn Read, length: u32) { + let mut handle = rdr.take(length as u64); + let mut buf: Vec = Vec::new(); + let _ = handle.read_to_end(&mut buf); + } + + fn unpack_char(vm: &VirtualMachine, rdr: &mut dyn Read) -> PyResult { + unpack_string(vm, rdr, 1) + } + + fn unpack_string(vm: &VirtualMachine, rdr: &mut dyn Read, length: u32) -> PyResult { + let mut handle = rdr.take(length as u64); + let mut buf: Vec = Vec::new(); + handle.read_to_end(&mut buf).map_err(|_| { + new_struct_error(vm, format!("unpack requires a buffer of {} bytes", length,)) + })?; + Ok(vm.ctx.new_bytes(buf)) + } + + fn unpack_pascal(vm: &VirtualMachine, rdr: &mut dyn Read, length: u32) -> PyResult { + let mut handle = rdr.take(length as u64); + let mut buf: Vec = Vec::new(); + handle.read_to_end(&mut buf).map_err(|_| { + new_struct_error(vm, format!("unpack requires a buffer of {} bytes", length,)) + })?; + let string_length = buf[0] as usize; + Ok(vm.ctx.new_bytes(buf[1..=string_length].to_vec())) + } + + #[pyfunction] + fn unpack(fmt: PyStringRef, buffer: PyBytesRef, vm: &VirtualMachine) -> PyResult { + let fmt_str = fmt.as_str(); + let format_spec = FormatSpec::parse(fmt_str).map_err(|e| new_struct_error(vm, e))?; + format_spec.unpack(buffer.get_value(), vm) + } + + fn unpack_code( + vm: &VirtualMachine, + code: &FormatCode, + rdr: &mut dyn Read, + items: &mut Vec, + ) -> PyResult<()> + where + Endianness: byteorder::ByteOrder, + { + let unpack = match code.code { + 'b' => unpack_i8, + 'B' => unpack_u8, + 'c' => unpack_char, + '?' => unpack_bool, + 'h' => unpack_i16::, + 'H' => unpack_u16::, + 'i' | 'l' => unpack_i32::, + 'I' | 'L' => unpack_u32::, + 'q' => unpack_i64::, + 'Q' => unpack_u64::, + 'n' => unpack_isize::, + 'N' => unpack_usize::, + 'P' => unpack_usize::, // FIXME: native-only + 'f' => unpack_f32::, + 'd' => unpack_f64::, + 'x' => { + unpack_empty(vm, rdr, code.repeat); + return Ok(()); + } + 's' => { + items.push(unpack_string(vm, rdr, code.repeat)?); + return Ok(()); + } + 'p' => { + items.push(unpack_pascal(vm, rdr, code.repeat)?); + return Ok(()); + } + c => { + panic!("Unsupported format code {:?}", c); + } + }; + for _ in 0..code.repeat { + items.push(unpack(vm, rdr)?); + } + Ok(()) + } + + #[pyfunction] + fn calcsize(fmt: Either, vm: &VirtualMachine) -> PyResult { + // FIXME: the given fmt must be parsed as ascii string + // https://github.com/RustPython/RustPython/pull/1792#discussion_r387340905 + let parsed = match fmt { + Either::A(string) => FormatSpec::parse(string.as_str()), + Either::B(bytes) => FormatSpec::parse(std::str::from_utf8(&bytes).unwrap()), + }; + let format_spec = parsed.map_err(|e| new_struct_error(vm, e))?; + Ok(format_spec.size()) + } + + #[pyclass(name = "Struct")] + #[derive(Debug)] + struct PyStruct { + spec: FormatSpec, + fmt_str: PyStringRef, + } + + impl PyValue for PyStruct { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("_struct", "Struct") + } + } + + #[pyimpl] + impl PyStruct { + #[pyslot] + fn tp_new( + cls: PyClassRef, + fmt: Either, + vm: &VirtualMachine, + ) -> PyResult> { + let fmt_str = match fmt { + Either::A(s) => s, + Either::B(b) => PyString::from(std::str::from_utf8(b.get_value()).unwrap()) + .into_ref_with_type(vm, vm.ctx.str_type())?, + }; + let spec = FormatSpec::parse(fmt_str.as_str()).map_err(|e| new_struct_error(vm, e))?; + PyStruct { spec, fmt_str }.into_ref_with_type(vm, cls) + } + + #[pyproperty] + fn format(&self) -> PyStringRef { + self.fmt_str.clone() + } + + #[pyproperty] + fn size(&self) -> usize { + self.spec.size() + } + + #[pymethod] + fn pack(&self, args: Args, vm: &VirtualMachine) -> PyResult> { + self.spec.pack(args.as_ref(), vm) + } + #[pymethod] + fn unpack(&self, data: PyBytesRef, vm: &VirtualMachine) -> PyResult { + self.spec.unpack(data.get_value(), vm) + } + } + + // seems weird that this is part of the "public" API, but whatever + // TODO: implement a format code->spec cache like CPython does? + #[pyfunction] + fn _clearcache() {} + + fn new_struct_error(vm: &VirtualMachine, msg: String) -> PyBaseExceptionRef { + // _struct.error must exist + let class = vm.try_class("_struct", "error").unwrap(); + vm.new_exception_msg(class, msg) } } -// seems weird that this is part of the "public" API, but whatever -// TODO: implement a format code->spec cache like CPython does? -fn clearcache() {} - -pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { +pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef { let ctx = &vm.ctx; let struct_error = ctx.new_class("struct.error", ctx.exceptions.exception_type.clone()); - py_module!(vm, "_struct", { - "_clearcache" => ctx.new_function(clearcache), - "pack" => ctx.new_function(struct_pack), - "unpack" => ctx.new_function(struct_unpack), - "calcsize" => ctx.new_function(struct_calcsize), + let module = _struct::make_module(vm); + extend_module!(vm, module, { "error" => struct_error, - "Struct" => PyStruct::make_class(ctx), - }) -} - -fn new_struct_error(vm: &VirtualMachine, msg: String) -> PyBaseExceptionRef { - // _struct.error must exist - let class = vm.try_class("_struct", "error").unwrap(); - vm.new_exception_msg(class, msg) + }); + module } diff --git a/vm/src/stdlib/random.rs b/vm/src/stdlib/random.rs index 7079ad564..63144451c 100644 --- a/vm/src/stdlib/random.rs +++ b/vm/src/stdlib/random.rs @@ -1,127 +1,123 @@ //! Random module. -use std::cell::RefCell; +pub(crate) use _random::make_module; -use num_bigint::{BigInt, Sign}; -use num_traits::Signed; -use rand::RngCore; +#[pymodule] +mod _random { + use crate::function::OptionalOption; + use crate::obj::objint::PyIntRef; + use crate::obj::objtype::PyClassRef; + use crate::pyobject::{PyClassImpl, PyRef, PyResult, PyValue}; + use crate::VirtualMachine; + use num_bigint::{BigInt, Sign}; + use num_traits::Signed; + use rand::RngCore; + use std::cell::RefCell; -use crate::function::OptionalOption; -use crate::obj::objint::PyIntRef; -use crate::obj::objtype::PyClassRef; -use crate::pyobject::{PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue}; -use crate::VirtualMachine; - -#[derive(Debug)] -enum PyRng { - Std(rand::rngs::ThreadRng), - MT(Box), -} - -impl Default for PyRng { - fn default() -> Self { - PyRng::Std(rand::thread_rng()) + #[derive(Debug)] + enum PyRng { + Std(rand::rngs::ThreadRng), + MT(Box), } -} -impl RngCore for PyRng { - fn next_u32(&mut self) -> u32 { - match self { - Self::Std(s) => s.next_u32(), - Self::MT(m) => m.next_u32(), + impl Default for PyRng { + fn default() -> Self { + PyRng::Std(rand::thread_rng()) } } - fn next_u64(&mut self) -> u64 { - match self { - Self::Std(s) => s.next_u64(), - Self::MT(m) => m.next_u64(), - } - } - fn fill_bytes(&mut self, dest: &mut [u8]) { - match self { - Self::Std(s) => s.fill_bytes(dest), - Self::MT(m) => m.fill_bytes(dest), - } - } - fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> { - match self { - Self::Std(s) => s.try_fill_bytes(dest), - Self::MT(m) => m.try_fill_bytes(dest), - } - } -} -#[pyclass(name = "Random")] -#[derive(Debug)] -struct PyRandom { - rng: RefCell, -} - -impl PyValue for PyRandom { - fn class(vm: &VirtualMachine) -> PyClassRef { - vm.class("_random", "Random") - } -} - -#[pyimpl(flags(BASETYPE))] -impl PyRandom { - #[pyslot(new)] - fn new(cls: PyClassRef, vm: &VirtualMachine) -> PyResult> { - PyRandom { - rng: RefCell::new(PyRng::default()), - } - .into_ref_with_type(vm, cls) - } - - #[pymethod] - fn random(&self) -> f64 { - mt19937::gen_res53(&mut *self.rng.borrow_mut()) - } - - #[pymethod] - fn seed(&self, n: OptionalOption) { - let new_rng = match n.flat_option() { - None => PyRng::default(), - Some(n) => { - let (_, mut key) = n.as_bigint().abs().to_u32_digits(); - if cfg!(target_endian = "big") { - key.reverse(); - } - PyRng::MT(Box::new(mt19937::MT19937::new_with_slice_seed(&key))) + impl RngCore for PyRng { + fn next_u32(&mut self) -> u32 { + match self { + Self::Std(s) => s.next_u32(), + Self::MT(m) => m.next_u32(), } - }; - - *self.rng.borrow_mut() = new_rng; + } + fn next_u64(&mut self) -> u64 { + match self { + Self::Std(s) => s.next_u64(), + Self::MT(m) => m.next_u64(), + } + } + fn fill_bytes(&mut self, dest: &mut [u8]) { + match self { + Self::Std(s) => s.fill_bytes(dest), + Self::MT(m) => m.fill_bytes(dest), + } + } + fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> { + match self { + Self::Std(s) => s.try_fill_bytes(dest), + Self::MT(m) => m.try_fill_bytes(dest), + } + } } - #[pymethod] - fn getrandbits(&self, mut k: usize) -> BigInt { - let mut rng = self.rng.borrow_mut(); + #[pyclass(name = "Random")] + #[derive(Debug)] + struct PyRandom { + rng: RefCell, + } - let mut gen_u32 = |k| rng.next_u32() >> (32 - k) as u32; + impl PyValue for PyRandom { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("_random", "Random") + } + } - if k <= 32 { - return gen_u32(k).into(); + #[pyimpl(flags(BASETYPE))] + impl PyRandom { + #[pyslot(new)] + fn new(cls: PyClassRef, vm: &VirtualMachine) -> PyResult> { + PyRandom { + rng: RefCell::new(PyRng::default()), + } + .into_ref_with_type(vm, cls) } - let words = (k - 1) / 8 + 1; - let mut wordarray = vec![0u32; words]; - - let it = wordarray.iter_mut(); - #[cfg(target_endian = "big")] - let it = it.rev(); - for word in it { - *word = gen_u32(k); - k -= 32; + #[pymethod] + fn random(&self) -> f64 { + mt19937::gen_res53(&mut *self.rng.borrow_mut()) } - BigInt::from_slice(Sign::NoSign, &wordarray) + #[pymethod] + fn seed(&self, n: OptionalOption) { + let new_rng = match n.flat_option() { + None => PyRng::default(), + Some(n) => { + let (_, mut key) = n.as_bigint().abs().to_u32_digits(); + if cfg!(target_endian = "big") { + key.reverse(); + } + PyRng::MT(Box::new(mt19937::MT19937::new_with_slice_seed(&key))) + } + }; + + *self.rng.borrow_mut() = new_rng; + } + + #[pymethod] + fn getrandbits(&self, mut k: usize) -> BigInt { + let mut rng = self.rng.borrow_mut(); + + let mut gen_u32 = |k| rng.next_u32() >> (32 - k) as u32; + + if k <= 32 { + return gen_u32(k).into(); + } + + let words = (k - 1) / 8 + 1; + let mut wordarray = vec![0u32; words]; + + let it = wordarray.iter_mut(); + #[cfg(target_endian = "big")] + let it = it.rev(); + for word in it { + *word = gen_u32(k); + k -= 32; + } + + BigInt::from_slice(Sign::NoSign, &wordarray) + } } } - -pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { - let ctx = &vm.ctx; - py_module!(vm, "_random", { - "Random" => PyRandom::make_class(ctx), - }) -} diff --git a/vm/src/stdlib/signal.rs b/vm/src/stdlib/signal.rs index e8ae72991..b5c9baabf 100644 --- a/vm/src/stdlib/signal.rs +++ b/vm/src/stdlib/signal.rs @@ -8,7 +8,6 @@ use arr_macro::arr; #[cfg(unix)] use nix::unistd::alarm as sig_alarm; -use libc; #[cfg(not(windows))] use libc::{SIG_DFL, SIG_ERR, SIG_IGN}; diff --git a/vm/src/stdlib/subprocess.rs b/vm/src/stdlib/subprocess.rs index 67cffccb3..400c30838 100644 --- a/vm/src/stdlib/subprocess.rs +++ b/vm/src/stdlib/subprocess.rs @@ -4,8 +4,6 @@ use std::fs::File; use std::io::ErrorKind; use std::time::Duration; -use subprocess; - use crate::function::OptionalArg; use crate::obj::objbytes::PyBytesRef; use crate::obj::objlist::PyListRef; diff --git a/vm/src/stdlib/warnings.rs b/vm/src/stdlib/warnings.rs index 175f5a118..3a872bf98 100644 --- a/vm/src/stdlib/warnings.rs +++ b/vm/src/stdlib/warnings.rs @@ -34,9 +34,7 @@ fn warnings_warn(args: WarnArgs, vm: &VirtualMachine) -> PyResult<()> { pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let ctx = &vm.ctx; - let module = py_module!(vm, "_warnings", { + py_module!(vm, "_warnings", { "warn" => ctx.new_function(warnings_warn), - }); - - module + }) } diff --git a/vm/src/version.rs b/vm/src/version.rs index af82e9598..f8aa7bbda 100644 --- a/vm/src/version.rs +++ b/vm/src/version.rs @@ -102,9 +102,8 @@ fn get_git_timestamp_datetime() -> DateTime { let timestamp = timestamp.parse::().unwrap_or(0); let datetime = UNIX_EPOCH + Duration::from_secs(timestamp); - let datetime = DateTime::::from(datetime); - datetime + datetime.into() } pub fn get_git_date() -> String { diff --git a/vm/src/vm.rs b/vm/src/vm.rs index 330842e1c..a919f46fe 100644 --- a/vm/src/vm.rs +++ b/vm/src/vm.rs @@ -791,8 +791,7 @@ impl VirtualMachine { self.trace_event(TraceEvent::Return)?; result } else if class.has_attr("__call__") { - let result = self.call_method(&callable, "__call__", args); - result + self.call_method(&callable, "__call__", args) } else { Err(self.new_type_error(format!( "'{}' object is not callable", @@ -806,8 +805,7 @@ impl VirtualMachine { where T: Into, { - let res = self._invoke(func_ref, args.into()); - res + self._invoke(func_ref, args.into()) } /// Call registered trace function. diff --git a/wasm/lib/src/convert.rs b/wasm/lib/src/convert.rs index 916d36fb9..cb0f6c707 100644 --- a/wasm/lib/src/convert.rs +++ b/wasm/lib/src/convert.rs @@ -1,5 +1,4 @@ use js_sys::{Array, ArrayBuffer, Object, Promise, Reflect, SyntaxError, Uint8Array}; -use serde_wasm_bindgen; use wasm_bindgen::{closure::Closure, prelude::*, JsCast}; use rustpython_compiler::error::{CompileError, CompileErrorType};