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

jeremyhylton · jeremyhylton · commit b476d597921a · 2009-03-26T21:34:20.000Z
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
@@ -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
@@ -40,6 +40,14 @@ def close(self):
         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."""
 
