Merge pull request RustPython#649 from jaroel/context-manager-open

Make IOBase a context manager

Author: rmliddle

tests/snippets/builtin_open.py

@@ -4,3 +4,7 @@
 assert 'RustPython' in fd.read()
 
 assert_raises(FileNotFoundError, lambda: open('DoesNotExist'))
+
+# Use open as a context manager
+with open('README.md') as fp:
+    fp.read()

vm/src/stdlib/io.rs

@@ -60,6 +60,26 @@ fn bytes_io_getvalue(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.get_none())
 }
 
+fn io_base_cm_enter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(instance, None)]);
+    Ok(instance.clone())
+}
+
+fn io_base_cm_exit(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        // The context manager protocol requires these, but we don't use them
+        required = [
+            (_instance, None),
+            (_exception_type, None),
+            (_exception_value, None),
+            (_traceback, None)
+        ]
+    );
+    Ok(vm.get_none())
+}
+
 fn buffered_io_base_init(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(buffered, None), (raw, None)]);
     vm.ctx.set_attr(&buffered, "raw", raw.clone());
@@ -347,7 +367,10 @@ pub fn io_open(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 
 pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
     //IOBase the abstract base class of the IO Module
-    let io_base = py_class!(ctx, "IOBase", ctx.object(), {});
+    let io_base = py_class!(ctx, "IOBase", ctx.object(), {
+        "__enter__" => ctx.new_rustfunc(io_base_cm_enter),
+        "__exit__" => ctx.new_rustfunc(io_base_cm_exit)
+    });
 
     // IOBase Subclasses
     let raw_io_base = py_class!(ctx, "RawIOBase", ctx.object(), {});