[pypy-svn] pypy default: Fix using BZ2File as a contextmanager.

[alex_gaynor]

Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: 
Changeset: r41832:eeec2ae3a0f3
Date: 2011-02-11 23:56 -0500
http://bitbucket.org/pypy/pypy/changeset/eeec2ae3a0f3/

Log:	Fix using BZ2File as a contextmanager.

diff --git a/pypy/module/bz2/test/test_bz2_file.py b/pypy/module/bz2/test/test_bz2_file.py
--- a/pypy/module/bz2/test/test_bz2_file.py
+++ b/pypy/module/bz2/test/test_bz2_file.py
@@ -1,3 +1,5 @@
+from __future__ import with_statement
+
 import py
 from pypy.conftest import gettestobjspace
 from pypy.module.bz2.test.support import CheckAllocation
@@ -414,6 +416,19 @@
         bz2f.close()
         bz2f = BZ2File(self.temppath, 'r')
         assert bz2f.read() == self.random_data
+
+    def test_context_manager(self):
+        from bz2 import BZ2File
+
+        with BZ2File(self.temppath, 'w') as f:
+            assert not f.closed
+            f.write("abc")
+        assert f.closed
+        with BZ2File(self.temppath, 'r') as f:
+            data = f.read()
+            assert data == "abc"
+        assert f.closed
+
         
         
 # has_cmdline_bunzip2 = sys.platform not in ("win32", "os2emx", "riscos")

diff --git a/pypy/module/bz2/interp_bz2.py b/pypy/module/bz2/interp_bz2.py
--- a/pypy/module/bz2/interp_bz2.py
+++ b/pypy/module/bz2/interp_bz2.py
@@ -290,7 +290,7 @@
 same_attributes_as_in_file.remove('__init__')
 same_attributes_as_in_file.extend([
     'name', 'mode', 'encoding', 'closed', 'newlines', 'softspace',
-    'writelines', '__weakref__'])
+    'writelines', '__exit__', '__weakref__'])
 
 W_BZ2File.typedef = TypeDef(
     "BZ2File",