generic and efficiency way to consume iterable object

qingshi163 · qingshi163 · commit 2d3ea4c205e1 · 2020-10-11T20:28:06.000+02:00
diff --git a/vm/src/bytesinner.rs b/vm/src/bytesinner.rs
@@ -9,15 +9,13 @@ use crate::function::{OptionalArg, OptionalOption};
 use crate::obj::objbytearray::PyByteArray;
 use crate::obj::objbytes::{PyBytes, PyBytesRef};
 use crate::obj::objint::{self, PyInt, PyIntRef};
-use crate::obj::objlist::PyList;
 use crate::obj::objmemory::PyMemoryView;
 use crate::obj::objsingletons::PyNoneRef;
 use crate::obj::objslice::PySliceRef;
 use crate::obj::objstr::{self, PyStr, PyStrRef};
-use crate::obj::objtuple::PyTuple;
 use crate::pyobject::{
-    BorrowValue, Either, PyComparisonValue, PyIterable, PyIterator, PyObjectRef, PyResult,
-    TryFromObject, TypeProtocol,
+    try_iterable_map, BorrowValue, Either, PyComparisonValue, PyIterable, PyIterator, PyObjectRef,
+    PyResult, TryFromObject, TypeProtocol,
 };
 use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex};
 use crate::slots::PyComparisonOp;
@@ -41,23 +39,18 @@ impl TryFromObject for PyBytesInner {
             return Ok(zelf);
         }
 
-        match_class!(match obj {
-            // TODO: generic way from &[PyObjectRef]
-            l @ PyList => l.to_byte_inner(vm),
-            t @ PyTuple => t.to_bytes_inner(vm),
-            obj => {
-                let iter = vm.get_method_or_type_error(obj.clone(), "__iter__", || {
-                    format!(
-                        "a bytes-like object is required, not '{}'",
-                        obj.class().name
-                    )
-                })?;
-                let iter = PyIterable::from_method(iter);
-                Ok(PyBytesInner {
-                    elements: iter.iter(vm)?.collect::<PyResult<_>>()?,
-                })
-            }
-        })
+        if let Ok(elements) = try_iterable_map(vm, &obj, |elem: PyIntRef| {
+            elem.borrow_value()
+                .to_u8()
+                .ok_or_else(|| vm.new_value_error("bytes must be in range (0, 256)".to_owned()))
+        }) {
+            return Ok(Self::from(elements));
+        }
+
+        Err(vm.new_type_error(format!(
+            "a bytes-like object is required, not '{}'",
+            obj.class().name
+        )))
     }
 }
 
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -4,18 +4,20 @@ use std::fmt;
 use std::marker::PhantomData;
 use std::ops::Deref;
 
+use itertools::Itertools;
 use num_bigint::BigInt;
 use num_complex::Complex64;
 use num_traits::ToPrimitive;
 
 use crate::bytecode;
+use crate::common::lock::{PyRwLock, PyRwLockReadGuard};
+use crate::common::rc::PyRc;
 use crate::exceptions::{self, PyBaseExceptionRef};
 use crate::function::{IntoFuncArgs, IntoPyNativeFunc};
 use crate::obj::objbuiltinfunc::PyFuncDef;
 use crate::obj::objbytearray;
 use crate::obj::objbytes;
-use crate::obj::objcode;
-use crate::obj::objcode::PyCodeRef;
+use crate::obj::objcode::{self, PyCodeRef};
 use crate::obj::objcomplex::PyComplex;
 use crate::obj::objdict::{PyDict, PyDictRef};
 use crate::obj::objfloat::PyFloat;
@@ -33,16 +35,13 @@ use crate::obj::objstaticmethod::PyStaticMethod;
 use crate::obj::objstr;
 use crate::obj::objtuple::{PyTuple, PyTupleRef};
 use crate::obj::objtype::{self, PyType, PyTypeRef};
-pub use crate::pyobjectrc::{PyObjectRc, PyObjectWeak};
 use crate::scope::Scope;
 use crate::slots::{PyTpFlags, PyTypeSlots};
 use crate::types::{create_type_with_slots, initialize_types, TypeZoo};
 use crate::vm::VirtualMachine;
-use rustpython_common::lock::{PyRwLock, PyRwLockReadGuard};
-use rustpython_common::rc::PyRc;
 
 pub use crate::common::borrow::BorrowValue;
-
+pub use crate::pyobjectrc::{PyObjectRc, PyObjectWeak};
 /* Python objects and references.
 
 Okay, so each python object itself is an class itself (PyObject). Each
@@ -900,6 +899,34 @@ where
     }
 }
 
+pub fn try_iterable_map<F, T, R>(
+    vm: &VirtualMachine,
+    obj: &PyObjectRef,
+    mut f: F,
+) -> PyResult<Vec<R>>
+where
+    T: TryFromObject,
+    F: FnMut(T) -> PyResult<R>,
+{
+    match_class!(match obj {
+        ref l @ PyList => l
+            .borrow_value()
+            .iter()
+            .map(|x| f(T::try_from_object(vm, x.clone())?))
+            .try_collect(),
+        ref t @ PyTuple => t
+            .borrow_value()
+            .iter()
+            .map(|x| f(T::try_from_object(vm, x.clone())?))
+            .try_collect(),
+        // TODO: put internal iterable typ
+        obj => {
+            let iter = PyIterable::<T>::try_from_object(vm, obj.clone())?;
+            iter.iter(vm)?.map(|x| f(x?)).try_collect()
+        }
+    })
+}
+
 impl TryFromObject for PyObjectRef {
     #[inline]
     fn try_from_object(_vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
