Add attributes to OSError and BlockingIOError

ChJR · ChJR · commit ee4d72e80af1 · 2021-10-30T05:18:22.000+09:00
diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs
@@ -963,6 +963,7 @@ pub(super) mod types {
     use crate::{
         builtins::{traceback::PyTracebackRef, PyInt, PyTupleRef, PyTypeRef},
         function::{FuncArgs, IntoPyResult},
+        pyobject::{PyValue, TypeProtocol},
         PyObjectRef, PyRef, PyResult, VirtualMachine,
     };
     use crossbeam_utils::atomic::AtomicCell;
@@ -1160,8 +1161,9 @@ pub(super) mod types {
         os_error,
         "Base class for I/O related errors.",
         os_error_new,
-        base_exception_init,
+        os_error_init,
     }
+
     #[cfg(not(target_arch = "wasm32"))]
     fn os_error_optional_new(
         args: Vec<PyObjectRef>,
@@ -1217,29 +1219,80 @@ pub(super) mod types {
     }
     #[cfg(not(target_arch = "wasm32"))]
     fn os_error_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
-        // We need this method, because of how `CPython` copies `init`
-        // from `BaseException` in `SimpleExtendsException` macro.
-        // See: `BaseException_new`
-        if cls.name().deref() == vm.ctx.exceptions.os_error.name().deref() {
-            match os_error_optional_new(args.args.to_vec(), vm) {
-                Some(error) => error.unwrap().into_pyresult(vm),
-                None => PyBaseException::slot_new(cls, args, vm),
-            }
-        } else {
-            PyBaseException::slot_new(cls, args, vm)
-        }
+        let exc_self: PyBaseExceptionRef =
+            if cls.name().deref() == vm.ctx.exceptions.os_error.name().deref() {
+                match os_error_optional_new(args.args.to_vec(), vm) {
+                    Some(error) => error.unwrap(),
+                    None => PyBaseException::new(args.args.to_vec(), vm)
+                        .into_ref_with_type(vm, cls)
+                        .unwrap(),
+                }
+            } else {
+                PyBaseException::new(args.args.to_vec(), vm)
+                    .into_ref_with_type(vm, cls)
+                    .unwrap()
+            };
+
+        os_error_parse_args(exc_self.clone(), &args, vm)?;
+        exc_self.into_pyresult(vm)
     }
     #[cfg(target_arch = "wasm32")]
     fn os_error_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
-        PyBaseException::slot_new(cls, args, vm)
+        let exc_self = PyBaseException::new(args.args.to_vec(), vm)
+            .into_ref_with_type(vm, cls)
+            .unwrap();
+        os_error_parse_args(exc_self.clone(), &args, vm)?;
+        exc_self.into_pyresult(vm)
+    }
+
+    fn os_error_parse_args(
+        exc_self: PyBaseExceptionRef,
+        args: &FuncArgs,
+        vm: &VirtualMachine,
+    ) -> PyResult<()> {
+        os_error_set_attr(&exc_self, args, vm, "filename", 2)?;
+        #[cfg(windows)]
+        os_error_set_attr(&exc_self, args, vm, "winerror", 3)?;
+        os_error_set_attr(&exc_self, args, vm, "filename2", 4)?;
+        Ok(())
     }
 
-    fn base_exception_init(
+    fn os_error_set_attr(
+        exc_self: &PyBaseExceptionRef,
+        args: &FuncArgs,
+        vm: &VirtualMachine,
+        attr_name: &str,
+        arg_idx: usize,
+    ) -> PyResult<()> {
+        let exc_self = exc_self.as_ref();
+
+        if let Some(arg) = args.args.as_slice().get(arg_idx).cloned() {
+            exc_self.set_attr(attr_name, arg, vm)?;
+        }
+
+        Ok(())
+    }
+
+    fn os_error_init(
         zelf: PyRef<PyBaseException>,
         args: FuncArgs,
         vm: &VirtualMachine,
     ) -> PyResult<()> {
-        PyBaseException::init(zelf, args, vm)
+        PyBaseException::init(zelf.clone(), args.clone(), vm)?;
+
+        let exc_self = zelf.as_object();
+
+        if exc_self.class().name().deref() == "BlockingIOError" {
+            exc_self.set_attr(
+                "characters_written",
+                vm.unwrap_or_none(args.args.as_slice().get(2).cloned()),
+                vm,
+            )?;
+        } else {
+            os_error_parse_args(zelf, &args, vm)?;
+        }
+
+        Ok(())
     }
 
     define_exception! {
@@ -1248,6 +1301,7 @@ pub(super) mod types {
         blocking_io_error,
         "I/O operation would block."
     }
+
     define_exception! {
         PyChildProcessError,
         PyOSError,
