unix file: Implement context manager protocol (for "with" statement).

2 files changed, 19 insertions, 0 deletions

diff --git a/tests/io/file-with.py b/tests/io/file-with.py
new file mode 100644
index 0000000000..2350721641
--- /dev/null
+++ b/tests/io/file-with.py
@@ -0,0 +1,12 @@
+f = open("io/data/file1")
+
+with f as f2:
+    print(f2.read())
+
+# File should be closed
+try:
+    f.read()
+except:
+    # Note: CPython and us throw different exception trying to read from
+    # close file.
+    print("can't read file after with")
diff --git a/unix/file.c b/unix/file.c
index bd54e7482c..41eb39e025 100644
--- a/unix/file.c
+++ b/unix/file.c
@@ -49,6 +49,11 @@ static mp_obj_t fdfile_close(mp_obj_t self_in) {
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(fdfile_close_obj, fdfile_close);
 
+mp_obj_t fdfile___exit__(uint n_args, const mp_obj_t *args) {
+    return fdfile_close(args[0]);
+}
+static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(fdfile___exit___obj, 4, 4, fdfile___exit__);
+
 static mp_obj_t fdfile_fileno(mp_obj_t self_in) {
     mp_obj_fdfile_t *self = self_in;
     return MP_OBJ_NEW_SMALL_INT((machine_int_t)self->fd);
@@ -112,6 +117,8 @@ STATIC const mp_map_elem_t rawfile_locals_dict_table[] = {
     { MP_OBJ_NEW_QSTR(MP_QSTR_readline), (mp_obj_t)&mp_stream_unbuffered_readline_obj},
     { MP_OBJ_NEW_QSTR(MP_QSTR_write), (mp_obj_t)&mp_stream_write_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_close), (mp_obj_t)&fdfile_close_obj },
+    { MP_OBJ_NEW_QSTR(MP_QSTR___enter__), (mp_obj_t)&mp_identity_obj },
+    { MP_OBJ_NEW_QSTR(MP_QSTR___exit__), (mp_obj_t)&fdfile___exit___obj },
 };
 
 STATIC MP_DEFINE_CONST_DICT(rawfile_locals_dict, rawfile_locals_dict_table);