py_class!: add __enter__ and __exit__

src/py_class/py_class_impl.py

@@ -717,8 +717,8 @@ special_names = {
     '__coerce__': unimplemented(),
 
     # With statement context managers
-    '__enter__': unimplemented(),
-    '__exit__': unimplemented(),
+    '__enter__': normal_method(),
+    '__exit__': normal_method(),
 
     # Coroutines
     '__await__': unimplemented(),

src/py_class/py_class_impl2.rs

@@ -578,18 +578,10 @@ macro_rules! py_class_impl {
         py_error! { "__divmod__ is not supported by py_class! yet." }
     };
 
-    { { def __enter__ $($tail:tt)* } $( $stuff:tt )* } => {
-        py_error! { "__enter__ is not supported by py_class! yet." }
-    };
-
     { { def __eq__ $($tail:tt)* } $( $stuff:tt )* } => {
         py_error! { "__eq__ is not supported by py_class! yet." }
     };
 
-    { { def __exit__ $($tail:tt)* } $( $stuff:tt )* } => {
-        py_error! { "__exit__ is not supported by py_class! yet." }
-    };
-
     { { def __float__ $($tail:tt)* } $( $stuff:tt )* } => {
         py_error! { "__float__ is not supported by py_class! yet." }
     };

src/py_class/py_class_impl3.rs

@@ -578,18 +578,10 @@ macro_rules! py_class_impl {
         py_error! { "__divmod__ is not supported by py_class! yet." }
     };
 
-    { { def __enter__ $($tail:tt)* } $( $stuff:tt )* } => {
-        py_error! { "__enter__ is not supported by py_class! yet." }
-    };
-
     { { def __eq__ $($tail:tt)* } $( $stuff:tt )* } => {
         py_error! { "__eq__ is not supported by py_class! yet." }
     };
 
-    { { def __exit__ $($tail:tt)* } $( $stuff:tt )* } => {
-        py_error! { "__exit__ is not supported by py_class! yet." }
-    };
-
     { { def __float__ $($tail:tt)* } $( $stuff:tt )* } => {
         py_error! { "__float__ is not supported by py_class! yet." }
     };

tests/test_class.rs

@@ -2,8 +2,7 @@
 
 #[macro_use] extern crate cpython;
 
-use cpython::{PyObject, PythonObject, PyDrop, PyClone, PyResult, Python, NoArgs, ObjectProtocol,
-    PyDict, PyBytes, PyUnicode, PyErr, exc};
+use cpython::*;
 use std::{mem, isize, iter};
 use std::cell::{Cell, RefCell};
 use std::sync::Arc;
@@ -25,7 +24,7 @@ macro_rules! py_assert {
 macro_rules! py_expect_exception {
     ($py:expr, $val:ident, $code:expr, $err:ident) => {{
         let d = PyDict::new($py);
-        d.set_item($py, stringify!($val), $val).unwrap();
+        d.set_item($py, stringify!($val), &$val).unwrap();
         let res = $py.run($code, None, Some(&d));
         let err = res.unwrap_err();
         if !err.matches($py, $py.get_type::<exc::$err>()) {
@@ -666,3 +665,38 @@ fn binary_arithmetic() {
     py_run!(py, c, "assert 1 | c == '1 | BA'");
 }
 
+py_class!(class ContextManager |py| {
+    data exit_called : Cell<bool>;
+
+    def __enter__(&self) -> PyResult<i32> {
+        Ok(42)
+    }
+
+    def __exit__(&self, ty: PyObject, value: PyObject, traceback: PyObject) -> PyResult<bool> {
+        self.exit_called(py).set(true);
+        if ty == py.get_type::<exc::ValueError>().into_object() {
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+});
+
+#[test]
+fn context_manager() {
+    let gil = Python::acquire_gil();
+    let py = gil.python();
+
+    let c = ContextManager::create_instance(py, Cell::new(false)).unwrap();
+    py_run!(py, c, "with c as x:\n  assert x == 42");
+    assert!(c.exit_called(py).get());
+
+    c.exit_called(py).set(false);
+    py_run!(py, c, "with c as x:\n  raise ValueError");
+    assert!(c.exit_called(py).get());
+
+    c.exit_called(py).set(false);
+    py_expect_exception!(py, c, "with c as x:\n  raise NotImplementedError", NotImplementedError);
+    assert!(c.exit_called(py).get());
+}
+