[Python-checkins] r70625 - in python/branches/py3k/Lib: test/test_urllib_response.py urllib/response.py

[jeremy.hylton]

Author: jeremy.hylton
Date: Thu Mar 26 22:34:20 2009
New Revision: 70625

Log:
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


Added:
   python/branches/py3k/Lib/test/test_urllib_response.py
Modified:
   python/branches/py3k/Lib/urllib/response.py

Added: python/branches/py3k/Lib/test/test_urllib_response.py
==============================================================================
--- (empty file)
+++ python/branches/py3k/Lib/test/test_urllib_response.py	Thu Mar 26 22:34:20 2009
@@ -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()

Modified: python/branches/py3k/Lib/urllib/response.py
==============================================================================
--- python/branches/py3k/Lib/urllib/response.py	(original)
+++ python/branches/py3k/Lib/urllib/response.py	Thu Mar 26 22:34:20 2009
@@ -40,6 +40,14 @@
         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."""