Add OSError.__reduce__

dvermd · dvermd · commit df5bca4c0644 · 2022-10-04T12:22:53.000+02:00
diff --git a/extra_tests/snippets/builtin_exceptions.py b/extra_tests/snippets/builtin_exceptions.py
@@ -273,6 +273,9 @@ class SubError(MyError):
 assert w.filename == None
 assert w.filename2 == None
 assert str(w) == ""
+x = pickle.loads(pickle.dumps(w, 4))
+assert type(w) == type(x)
+assert w == x
 
 w = OSError(0)
 assert w.errno == None
@@ -281,6 +284,9 @@ class SubError(MyError):
 assert w.filename == None
 assert w.filename2 == None
 assert str(w) == "0"
+x = pickle.loads(pickle.dumps(w, 4))
+assert type(w) == type(x)
+assert w == x
 
 w = OSError('foo')
 assert w.errno == None
@@ -289,6 +295,9 @@ class SubError(MyError):
 assert w.filename == None
 assert w.filename2 == None
 assert str(w) == "foo"
+x = pickle.loads(pickle.dumps(w, 4))
+assert type(w) == type(x)
+assert w == x
 
 w = OSError('a', 'b', 'c', 'd', 'e', 'f')
 assert w.errno == None
@@ -297,6 +306,9 @@ class SubError(MyError):
 assert w.filename == None
 assert w.filename2 == None
 assert str(w) == "('a', 'b', 'c', 'd', 'e', 'f')"
+x = pickle.loads(pickle.dumps(w, 4))
+assert type(w) == type(x)
+assert w == x
 
 # Custom `__new__` and `__init__`:
 assert ImportError.__init__.__qualname__ == 'ImportError.__init__'
diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs
@@ -1,4 +1,5 @@
 use self::types::{PyBaseException, PyBaseExceptionRef};
+use crate::builtins::tuple::IntoPyTuple;
 use crate::common::lock::PyRwLock;
 use crate::{
     builtins::{
@@ -773,6 +774,7 @@ impl ExceptionZoo {
             // second exception filename
             "filename2" => ctx.none(),
             "__str__" => ctx.new_method("__str__", excs.os_error, os_error_str),
+            "__reduce__" => ctx.new_method("__reduce__", excs.os_error, os_error_reduce),
         });
         // TODO: this isn't really accurate
         #[cfg(windows)]
@@ -907,6 +909,42 @@ fn os_error_str(exc: PyBaseExceptionRef, vm: &VirtualMachine) -> PyResult<PyStrR
     }
 }
 
+fn os_error_reduce(exc: PyBaseExceptionRef, vm: &VirtualMachine) -> PyTupleRef {
+    let args = exc.args();
+    let obj = exc.as_object().to_owned();
+    let mut result: Vec<PyObjectRef> = vec![obj.class().clone().into()];
+
+    if args.len() >= 2 && args.len() <= 5 {
+        // SAFETY: len() == 2 is checked so get_arg 1 or 2 won't panic
+        let errno = exc.get_arg(0).unwrap();
+        let msg = exc.get_arg(1).unwrap();
+
+        if let Ok(filename) = obj.get_attr("filename", vm) {
+            if !vm.is_none(&filename) {
+                let mut args_reduced: Vec<PyObjectRef> = vec![errno, msg, filename];
+
+                if let Ok(filename2) = obj.get_attr("filename2", vm) {
+                    if !vm.is_none(&filename2) {
+                        args_reduced.push(filename2);
+                    }
+                }
+                result.push(args_reduced.into_pytuple(vm).into());
+            } else {
+                result.push(vm.new_tuple((errno, msg)).into());
+            }
+        } else {
+            result.push(vm.new_tuple((errno, msg)).into());
+        }
+    } else {
+        result.push(args.into());
+    }
+
+    if let Some(dict) = obj.dict().filter(|x| !x.is_empty()) {
+        result.push(dict.into());
+    }
+    result.into_pytuple(vm)
+}
+
 fn system_exit_code(exc: PyBaseExceptionRef) -> Option<PyObjectRef> {
     exc.args.read().first().map(|code| {
         match_class!(match code {
