RustPython · youknowone · Jun 3, 2020 · May 8, 2020 · May 8, 2020 · May 8, 2020
diff --git a/Lib/test/test_unpack.py b/Lib/test/test_unpack.py
@@ -0,0 +1,151 @@
+doctests = """
+
+Unpack tuple
+
+    >>> t = (1, 2, 3)
+    >>> a, b, c = t
+    >>> a == 1 and b == 2 and c == 3
+    True
+
+Unpack list
+
+    >>> l = [4, 5, 6]
+    >>> a, b, c = l
+    >>> a == 4 and b == 5 and c == 6
+    True
+
+Unpack implied tuple
+
+    >>> a, b, c = 7, 8, 9
+    >>> a == 7 and b == 8 and c == 9
+    True
+
+Unpack string... fun!
+
+    >>> a, b, c = 'one'
+    >>> a == 'o' and b == 'n' and c == 'e'
+    True
+
+Unpack generic sequence
+
+    >>> class Seq:
+    ...     def __getitem__(self, i):
+    ...         if i >= 0 and i < 3: return i
+    ...         raise IndexError
+    ...
+    >>> a, b, c = Seq()
+    >>> a == 0 and b == 1 and c == 2
+    True
+
+Single element unpacking, with extra syntax
+
+    >>> st = (99,)
+    >>> sl = [100]
+    >>> a, = st
+    >>> a
+    99
+    >>> b, = sl
+    >>> b
+    100
+
+Now for some failures
+
+Unpacking non-sequence
+
+    >>> a, b, c = 7
+    Traceback (most recent call last):
+      ...
+    TypeError: cannot unpack non-iterable int object
+
+Unpacking tuple of wrong size
+
+    >>> a, b = t
+    Traceback (most recent call last):
+      ...
+    ValueError: too many values to unpack (expected 2)
+
+Unpacking tuple of wrong size
+
+    >>> a, b = l
+    Traceback (most recent call last):
+      ...
+    ValueError: too many values to unpack (expected 2)
+
+Unpacking sequence too short
+
+    >>> a, b, c, d = Seq()
+    Traceback (most recent call last):
+      ...
+    ValueError: not enough values to unpack (expected 4, got 3)
+
+Unpacking sequence too long
+
+    >>> a, b = Seq()
+    Traceback (most recent call last):
+      ...
+    ValueError: too many values to unpack (expected 2)
+
+Unpacking a sequence where the test for too long raises a different kind of
+error
+
+    >>> class BozoError(Exception):
+    ...     pass
+    ...
+    >>> class BadSeq:
+    ...     def __getitem__(self, i):
+    ...         if i >= 0 and i < 3:
+    ...             return i
+    ...         elif i == 3:
+    ...             raise BozoError
+    ...         else:
+    ...             raise IndexError
+    ...
+
+Trigger code while not expecting an IndexError (unpack sequence too long, wrong
+error)
+
+    >>> a, b, c, d, e = BadSeq()
+    Traceback (most recent call last):
+      ...
+    test.test_unpack.BozoError
+
+Trigger code while expecting an IndexError (unpack sequence too short, wrong
+error)
+
+    >>> a, b, c = BadSeq()
+    Traceback (most recent call last):
+      ...
+    test.test_unpack.BozoError
+
+Allow unpacking empty iterables
+
+    >>> () = []
+    >>> [] = ()
+    >>> [] = []
+    >>> () = ()
+
+Unpacking non-iterables should raise TypeError
+
+    >>> () = 42
+    Traceback (most recent call last):
+      ...
+    TypeError: cannot unpack non-iterable int object
+
+Unpacking to an empty iterable should raise ValueError
+
+    >>> () = [42]
+    Traceback (most recent call last):
+      ...
+    ValueError: too many values to unpack (expected 0)
+
+"""
+
+__test__ = {'doctests' : doctests}
+
+def test_main(verbose=False):
+    from test import support
+    from test import test_unpack
+    support.run_doctest(test_unpack, verbose)
+
+if __name__ == "__main__":
+    test_main(verbose=True)
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -667,13 +667,35 @@ impl ExecutingFrame<'_> {
             }
             bytecode::Instruction::UnpackSequence { size } => {
                 let value = self.pop_value();
-                let elements = vm.extract_elements(&value)?;
-                if elements.len() != *size {
-                    Err(vm.new_value_error("Wrong number of values to unpack".to_owned()))
-                } else {
-                    for element in elements.into_iter().rev() {
-                        self.push_value(element);
+                let elements = vm.extract_elements(&value).map_err(|e| {
+                    if e.class().is(&vm.ctx.exceptions.type_error) {
+                        vm.new_type_error(format!(
+                            "cannot unpack non-iterable {} object",
+                            value.class().name
+                        ))
+                    } else {
+                        e
+                    }
+                })?;
+                let msg = match elements.len().cmp(size) {
+                    std::cmp::Ordering::Equal => {
+                        for element in elements.into_iter().rev() {
+                            self.push_value(element);
+                        }
+                        None
+                    }
+                    std::cmp::Ordering::Greater => {
+                        Some(format!("too many values to unpack (expected {})", size))
                     }
+                    std::cmp::Ordering::Less => Some(format!(
+                        "not enough values to unpack (expected {}, got {})",
+                        size,
+                        elements.len()
+                    )),
+                };
+                if let Some(msg) = msg {
+                    Err(vm.new_value_error(msg))
+                } else {
                     Ok(None)
                 }
             }

diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs
@@ -327,7 +327,7 @@ impl PyString {
     }
 
     #[pymethod(name = "__repr__")]
-    fn repr(&self, vm: &VirtualMachine) -> PyResult<String> {
+    pub(crate) fn repr(&self, vm: &VirtualMachine) -> PyResult<String> {
         let in_len = self.value.len();
         let mut out_len = 0usize;
         // let mut max = 127;

diff --git a/vm/src/stdlib/io.rs b/vm/src/stdlib/io.rs
@@ -16,7 +16,7 @@ use crate::obj::objbyteinner::PyBytesLike;
 use crate::obj::objbytes::PyBytesRef;
 use crate::obj::objint;
 use crate::obj::objiter;
-use crate::obj::objstr::{self, PyStringRef};
+use crate::obj::objstr::{self, PyString, PyStringRef};
 use crate::obj::objtype::{self, PyClassRef};
 use crate::pyobject::{
     BufferProtocol, Either, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
@@ -812,12 +812,69 @@ fn buffered_writer_seekable(_self: PyObjectRef) -> bool {
     true
 }
 
+#[derive(FromArgs)]
+struct TextIOWrapperArgs {
+    #[pyarg(positional_or_keyword, optional = false)]
+    buffer: PyObjectRef,
+    #[pyarg(positional_or_keyword, default = "None")]
+    encoding: Option<PyStringRef>,
+    #[pyarg(positional_or_keyword, default = "None")]
+    errors: Option<PyStringRef>,
+    #[pyarg(positional_or_keyword, default = "None")]
+    newline: Option<PyStringRef>,
+}
+
+impl TextIOWrapperArgs {
+    fn validate_newline(&self, vm: &VirtualMachine) -> PyResult<()> {
+        if let Some(pystr) = &self.newline {
+            match pystr.as_str() {
+                "" | "\n" | "\r" | "\r\n" => Ok(()),
+                _ => {
+                    Err(vm.new_value_error(format!("illegal newline value: '{}'", pystr.repr(vm)?)))
+                }
+            }
+        } else {
+            Ok(())
+        }
+    }
+}
+
 fn text_io_wrapper_init(
     instance: PyObjectRef,
-    buffer: PyObjectRef,
+    args: TextIOWrapperArgs,
     vm: &VirtualMachine,
 ) -> PyResult<()> {
-    vm.set_attr(&instance, "buffer", buffer.clone())?;
+    args.validate_newline(vm)?;
+
+    let mut encoding: Option<PyStringRef> = args.encoding.clone();
+    let mut self_encoding = None; // TODO: Try os.device_encoding(fileno)
+    if encoding.is_none() && self_encoding.is_none() {
+        // TODO: locale module
+        self_encoding = Some("utf-8");
+    }
+    if let Some(self_encoding) = self_encoding {
+        encoding = Some(PyString::from(self_encoding).into_ref(vm));
+    } else if let Some(ref encoding) = encoding {
+        self_encoding = Some(encoding.as_str())
+    } else {
+        return Err(vm.new_os_error("could not determine default encoding".to_owned()));
+    }
+    let _ = encoding; // TODO: check codec
+
+    let errors = args
+        .errors
+        .map_or_else(|| vm.ctx.new_str("strict"), |o| o.into_object());
+
+    // let readuniversal = args.newline.map_or_else(true, |s| s.as_str().is_empty());
+
+    vm.set_attr(
+        &instance,
+        "encoding",
+        self_encoding.map_or_else(|| vm.get_none(), |s| vm.ctx.new_str(s)),
+    )?;
+    vm.set_attr(&instance, "errors", errors)?;
+    vm.set_attr(&instance, "buffer", args.buffer.clone())?;
+
     Ok(())
 }