Index: Modules/cStringIO.c =================================================================== --- Modules/cStringIO.c (revision 60906) +++ Modules/cStringIO.c (working copy) @@ -477,6 +477,35 @@ Py_RETURN_NONE; } + +PyDoc_STRVAR(IO_enter__doc__, + "__enter__() -> self."); +static PyObject * +IO_enter(IOobject *self) +{ + if (!IO__opencheck(self)) + return NULL; + Py_INCREF(self); + return (PyObject *)self; +} + +PyDoc_STRVAR(IO_exit__doc__, + "__exit__(*excinfo) -> None. Calls close()."); + +static PyObject * +IO_exit(PyObject *self, PyObject *args) +{ + PyObject *ret = PyObject_CallMethod(self, "close", NULL); + if (!ret) + /* If error occurred, pass through */ + return NULL; + Py_DECREF(ret); + /* We cannot return the result of close since a true + * value will be interpreted as "yes, swallow the + * exception if one was raised inside the with block". */ + Py_RETURN_NONE; +} + static struct PyMethodDef O_methods[] = { /* Common methods: */ {"flush", (PyCFunction)IO_flush, METH_NOARGS, IO_flush__doc__}, @@ -494,6 +523,11 @@ {"seek", (PyCFunction)O_seek, METH_VARARGS, O_seek__doc__}, {"write", (PyCFunction)O_write, METH_VARARGS, O_write__doc__}, {"writelines", (PyCFunction)O_writelines, METH_O, O_writelines__doc__}, + + /* Context management methods: */ + {"__enter__", (PyCFunction)IO_enter, METH_NOARGS, IO_enter__doc__}, + {"__exit__", (PyCFunction)IO_exit, METH_VARARGS, IO_exit__doc__}, + {NULL, NULL} /* sentinel */ }; Index: Lib/StringIO.py =================================================================== --- Lib/StringIO.py (revision 60906) +++ Lib/StringIO.py (working copy) @@ -62,6 +62,16 @@ self.closed = False self.softspace = 0 + # Context management protocol + def __enter__(self): + if self.closed: + raise ValueError("I/O operation on closed file") + return self + + def __exit__(self, exc, value, tb): + self.close() + + # file protocol def __iter__(self): return self Index: Lib/test/test_StringIO.py =================================================================== --- Lib/test/test_StringIO.py (revision 60906) +++ Lib/test/test_StringIO.py (working copy) @@ -118,6 +118,15 @@ self.assertEqual(s, unicode('abcuvwxyz!')) self.assertEqual(type(s), types.UnicodeType) + def test_context_management(self): + with self.MODULE.StringIO() as f: + self.failIf(f.closed) + self.failUnless(f.closed) + def use_closed(): + with f: + pass + self.failUnlessRaises(ValueError, use_closed) + class TestcStringIO(TestGenericStringIO): MODULE = cStringIO