svn commit: r1235069 - in /avro/trunk: CHANGES.txt lang/py/src/avro/datafile.py lang/py/test/test_datafile.py

[cu...@apache.org]

Author: cutting
Date: Mon Jan 23 23:43:26 2012
New Revision: 1235069

URL: http://svn.apache.org/viewvc?rev=1235069&view=rev
Log:
AVRO-854.  Python: Permit DataFileWriter and DataFileReader to be used as context managers in with statements.  Contributed by Harsh J.

Modified:
    avro/trunk/CHANGES.txt
    avro/trunk/lang/py/src/avro/datafile.py
    avro/trunk/lang/py/test/test_datafile.py

Modified: avro/trunk/CHANGES.txt
URL: http://svn.apache.org/viewvc/avro/trunk/CHANGES.txt?rev=1235069&r1=1235068&r2=1235069&view=diff
==============================================================================
--- avro/trunk/CHANGES.txt (original)
+++ avro/trunk/CHANGES.txt Mon Jan 23 23:43:26 2012
@@ -4,6 +4,9 @@ Avro 1.6.2 (unreleased)
 
   NEW FEATURES
 
+    AVRO-854.  Python: Permit DataFileWriter and DataFileReader to be
+    used as context managers in "with" statements. (Harsh J via cutting)
+
   OPTIMIZATIONS
 
   IMPROVEMENTS

Modified: avro/trunk/lang/py/src/avro/datafile.py
URL: http://svn.apache.org/viewvc/avro/trunk/lang/py/src/avro/datafile.py?rev=1235069&r1=1235068&r2=1235069&view=diff
==============================================================================
--- avro/trunk/lang/py/src/avro/datafile.py (original)
+++ avro/trunk/lang/py/src/avro/datafile.py Mon Jan 23 23:43:26 2012
@@ -118,6 +118,14 @@ class DataFileWriter(object):
   sync_marker = property(lambda self: self._sync_marker)
   meta = property(lambda self: self._meta)
 
+  def __enter__(self):
+    return self
+
+  def __exit__(self, type, value, traceback):
+    # Perform a close if there's no exception
+    if type is None:
+      self.close()
+
   # read/write properties
   def set_block_count(self, new_val):
     self._block_count = new_val
@@ -229,7 +237,15 @@ class DataFileReader(object):
     # get ready to read
     self._block_count = 0
     self.datum_reader.writers_schema = schema.parse(self.get_meta(SCHEMA_KEY))
-  
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, type, value, traceback):
+    # Perform a close if there's no exception
+    if type is None:
+      self.close()
+
   def __iter__(self):
     return self
 

Modified: avro/trunk/lang/py/test/test_datafile.py
URL: http://svn.apache.org/viewvc/avro/trunk/lang/py/test/test_datafile.py?rev=1235069&r1=1235068&r2=1235069&view=diff
==============================================================================
--- avro/trunk/lang/py/test/test_datafile.py (original)
+++ avro/trunk/lang/py/test/test_datafile.py Mon Jan 23 23:43:26 2012
@@ -159,5 +159,30 @@ class TestDataFile(unittest.TestCase):
     os.remove(FILENAME)
     self.assertEquals(correct, len(CODECS_TO_VALIDATE)*len(SCHEMAS_TO_VALIDATE))
 
+  def test_context_manager(self):
+    # Context manager was introduced as a first class
+    # member only in Python 2.6 and above.
+    import sys
+    if sys.version_info < (2,6):
+      print 'Skipping context manager tests on this Python version.'
+      return
+    # Test the writer with a 'with' statement.
+    writer = open(FILENAME, 'wb')
+    datum_writer = io.DatumWriter()
+    sample_schema, sample_datum = SCHEMAS_TO_VALIDATE[1]
+    schema_object = schema.parse(sample_schema)
+    with datafile.DataFileWriter(writer, datum_writer, schema_object) as dfw:
+      dfw.append(sample_datum)
+    self.assertTrue(writer.closed)
+
+    # Test the reader with a 'with' statement.
+    datums = []
+    reader = open(FILENAME, 'rb')
+    datum_reader = io.DatumReader()
+    with datafile.DataFileReader(reader, datum_reader) as dfr:
+      for datum in dfr:
+        datums.append(datum)
+    self.assertTrue(reader.closed)
+
 if __name__ == '__main__':
   unittest.main()