From: Jeremy Hylton <jeremy@alum.mit.edu>
Date: Thu, 26 Mar 2009 21:34:20 +0000 (+0000)
Subject: Add __enter__ and __exit__ methods to addbase() so that it supports with.
X-Git-Tag: v3.1a2~200
X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=b476d597921ab0ef52de1f56a5efa71c65727114;p=thirdparty%2FPython%2Fcpython.git

Add __enter__ and __exit__ methods to addbase() so that it supports with.

This change also adds a minimal unittest of urllib.response.addbase.
More are needed, but not to cover the small change being made here.

Addresses http://bugs.python.org/issue5418
---

diff --git a/Lib/test/test_urllib_response.py b/Lib/test/test_urllib_response.py
new file mode 100644
index 000000000000..fdd33259dd38
--- /dev/null
+++ b/Lib/test/test_urllib_response.py
@@ -0,0 +1,42 @@
+"""Unit tests for code in urllib.response."""
+
+import test.support
+import urllib.response
+import unittest
+
+class TestFile(object):
+
+    def __init__(self):
+        self.closed = False
+
+    def read(self, bytes):
+        pass
+
+    def readline(self):
+        pass
+
+    def close(self):
+        self.closed = True
+
+class Testaddbase(unittest.TestCase):
+
+    # TODO(jhylton): Write tests for other functionality of addbase()
+
+    def setUp(self):
+        self.fp = TestFile()
+        self.addbase = urllib.response.addbase(self.fp)
+
+    def test_with(self):
+        def f():
+            with self.addbase as spam:
+                pass
+        self.assertFalse(self.fp.closed)
+        f()
+        self.assertTrue(self.fp.closed)
+        self.assertRaises(ValueError, f)
+
+def test_main():
+    test.support.run_unittest(Testaddbase)
+
+if __name__ == '__main__':
+    test_main()
diff --git a/Lib/urllib/response.py b/Lib/urllib/response.py
index 52eeed0f70c7..985964231093 100644
--- a/Lib/urllib/response.py
+++ b/Lib/urllib/response.py
@@ -40,6 +40,14 @@ class addbase(object):
         if self.fp: self.fp.close()
         self.fp = None
 
+    def __enter__(self):
+        if self.fp is None:
+            raise ValueError("I/O operation on closed file")
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
+
 class addclosehook(addbase):
     """Class to add a close hook to an open file."""