From 9ce5240598f186e66c5e5d765042294fe95b4efa Mon Sep 17 00:00:00 2001
From: Maxim Kurnikov <maxim.kurnikov@gmail.com>
Date: Wed, 3 Jul 2019 22:15:02 +0300
Subject: [PATCH] add name, path, msg attributes to ImportError

---
 tests/snippets/exceptions.py | 20 ++++++++++++++++++++
 tests/snippets/import.py     |  5 +++++
 vm/src/exceptions.rs         | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+)

diff --git a/tests/snippets/exceptions.py b/tests/snippets/exceptions.py
index 1349f46296..cb178fc066 100644
--- a/tests/snippets/exceptions.py
+++ b/tests/snippets/exceptions.py
@@ -1,3 +1,4 @@
+# KeyError
 empty_exc = KeyError()
 assert str(empty_exc) == ''
 assert repr(empty_exc) == 'KeyError()'
@@ -23,3 +24,22 @@ def __str__(self):
 exc = KeyError(A())
 assert str(exc) == 'repr'
 assert repr(exc) == 'KeyError(repr,)'
+
+# ImportError / ModuleNotFoundError
+exc = ImportError()
+assert exc.name is None
+assert exc.path is None
+assert exc.msg is None
+assert exc.args == ()
+
+exc = ImportError('hello')
+assert exc.name is None
+assert exc.path is None
+assert exc.msg == 'hello'
+assert exc.args == ('hello',)
+
+exc = ImportError('hello', name='name', path='path')
+assert exc.name == 'name'
+assert exc.path == 'path'
+assert exc.msg == 'hello'
+assert exc.args == ('hello',)
diff --git a/tests/snippets/import.py b/tests/snippets/import.py
index 99322a2ea9..7c730c2659 100644
--- a/tests/snippets/import.py
+++ b/tests/snippets/import.py
@@ -25,6 +25,11 @@
 except ImportError:
     pass
 
+try:
+    import mymodule
+except ModuleNotFoundError as exc:
+    assert exc.name == 'mymodule'
+
 
 test = __import__("import_target")
 assert test.X == import_target.X
diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs
index b3d03ea806..cfd611c3ff 100644
--- a/vm/src/exceptions.rs
+++ b/vm/src/exceptions.rs
@@ -312,6 +312,35 @@ impl ExceptionZoo {
     }
 }
 
+fn import_error_init(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    // TODO: call super().__init__(*args) instead
+    exception_init(vm, args.clone())?;
+
+    let exc_self = args.args[0].clone();
+    vm.set_attr(
+        &exc_self,
+        "name",
+        args.kwargs
+            .get("name")
+            .cloned()
+            .unwrap_or_else(|| vm.get_none()),
+    )?;
+    vm.set_attr(
+        &exc_self,
+        "path",
+        args.kwargs
+            .get("path")
+            .cloned()
+            .unwrap_or_else(|| vm.get_none()),
+    )?;
+    vm.set_attr(
+        &exc_self,
+        "msg",
+        args.args.get(1).cloned().unwrap_or_else(|| vm.get_none()),
+    )?;
+    Ok(vm.get_none())
+}
+
 pub fn init(context: &PyContext) {
     let base_exception_type = &context.exceptions.base_exception_type;
     extend_class!(context, base_exception_type, {
@@ -323,4 +352,9 @@ pub fn init(context: &PyContext) {
         "__str__" => context.new_rustfunc(exception_str),
         "__repr__" => context.new_rustfunc(exception_repr),
     });
+
+    let import_error_type = &context.exceptions.import_error;
+    extend_class!(context, import_error_type, {
+        "__init__" => context.new_rustfunc(import_error_init)
+    });
 }