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 0000000..fdd3325
--- /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 52eeed0..9859642 100644
--- a/Lib/urllib/response.py
+++ b/Lib/urllib/response.py
@@ -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."""