G3Writer as a python context (#158)

This PR enables using G3Writer and G3MultiFileWriter as a python context object,
so that the object can be created and cleanly closed using the python `with`
statement.  For example:

    with core.G3Writer(filename="path/to/file.g3") as writer:
        writer(frame1)
        writer(frame2)

This will open the output file, write the two frames to it, then properly close
the file by appending an EndProcessing frame.

core/python/__init__.py

@@ -11,7 +11,7 @@ def fix_logging_crash():
 atexit.register(fix_logging_crash)
 del fix_logging_crash
 
-from spt3g.core.G3File import G3File
+from spt3g.core.fileio import *
 from spt3g.core.modconstruct import pipesegment, indexmod, pipesegment_nodoc
 from spt3g.core.funcconstruct import usefulfunc
 try:

core/python/fileio.py

@@ -0,0 +1,59 @@
+from spt3g import core
+
+__all__ = ["G3File"]
+
+
+class G3File(object):
+    """Iterable class for G3 files, as created by G3Writer. Loop through frames by doing something like:
+
+        with core.G3File("/path/to/file.g3") as f:
+            for frame in f:
+                print(frame)
+
+    An entire file can also be read into an indexable list by doing:
+
+        f = list(core.G3File("/path/to/file.g3"))
+    """
+
+    def __init__(self, path):
+        self.reader = core.G3Reader(path)
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        if self.reader is None:
+            raise StopIteration("Reader closed")
+        frames = self.reader.Process(None)
+        if len(frames) == 0:
+            raise StopIteration("No more frames in file")
+        if len(frames) > 1:
+            raise ValueError("Too many frames returned by reader")
+        return frames[0]
+
+    __next__ = next
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        del self.reader
+        self.reader = None
+
+
+# writer context methods
+
+
+def writer_enter(self):
+    return self
+
+
+def writer_exit(self, *args, **kwargs):
+    fr = core.G3Frame(core.G3FrameType.EndProcessing)
+    self(fr)
+
+
+core.G3Writer.__enter__ = writer_enter
+core.G3Writer.__exit__ = writer_exit
+core.G3MultiFileWriter.__enter__ = writer_enter
+core.G3MultiFileWriter.__exit__ = writer_exit

core/tests/fileio.py

@@ -50,8 +50,8 @@ def checkinfo(fr):
 assert n == 10, 'Wrong number of frames read (%d should be %d)' % (n, 10)
 
 # Skip empty files
-wr = core.G3Writer("empty.g3")
-del wr
+with core.G3Writer("empty.g3") as wr:
+    pass
 n = 0
 pipe = core.G3Pipeline()
 pipe.Add(core.G3Reader, filename=["empty.g3", "test.g3", "empty.g3"], track_filename=True)

core/tests/vecint.py

@@ -35,12 +35,11 @@ def tearDown(self):
     def test_serialize(self):
         """Confirm full ranges can be saved and loaded."""
 
-        w = core.G3Writer(test_filename)
-        for isize, val in bit_sizes:
-            f = core.G3Frame()
-            f['v'] = core.G3VectorInt([val] * 10)
-            w(f)
-        del w
+        with core.G3Writer(test_filename) as w:
+            for isize, val in bit_sizes:
+                f = core.G3Frame()
+                f['v'] = core.G3VectorInt([val] * 10)
+                w(f)
 
         r = core.G3Reader(test_filename)
         for isize, val in bit_sizes:
@@ -53,12 +52,11 @@ def test_serialize(self):
     def test_serialize_map(self):
         """Confirm full ranges can be saved and loaded."""
 
-        w = core.G3Writer(test_filename)
-        for isize, val in bit_sizes:
-            f = core.G3Frame()
-            f['m'] = core.G3MapInt({str(i): val for i in range(10)})
-            w(f)
-        del w
+        with core.G3Writer(test_filename) as w:
+            for isize, val in bit_sizes:
+                f = core.G3Frame()
+                f['m'] = core.G3MapInt({str(i): val for i in range(10)})
+                w(f)
 
         r = core.G3Reader(test_filename)
         for isize, val in bit_sizes:
@@ -71,12 +69,11 @@ def test_serialize_map(self):
     def test_serialize_map_vector(self):
         """Confirm full ranges can be saved and loaded."""
 
-        w = core.G3Writer(test_filename)
-        for isize, val in bit_sizes:
-            f = core.G3Frame()
-            f['m'] = core.G3MapVectorInt({'a': [val] * 10})
-            w(f)
-        del w
+        with core.G3Writer(test_filename) as w:
+            for isize, val in bit_sizes:
+                f = core.G3Frame()
+                f['m'] = core.G3MapVectorInt({'a': [val] * 10})
+                w(f)
 
         r = core.G3Reader(test_filename)
         for isize, val in bit_sizes:
@@ -91,11 +88,10 @@ def test_compression(self):
         count = 10000
         overhead = 200
         for isize, val in bit_sizes:
-            w = core.G3Writer(test_filename)
-            f = core.G3Frame()
-            f['v'] = core.G3VectorInt([val] * count)
-            w(f)
-            del w
+            with core.G3Writer(test_filename) as w:
+                f = core.G3Frame()
+                f['v'] = core.G3VectorInt([val] * count)
+                w(f)
             on_disk = os.path.getsize(test_filename)
             self.assertTrue(abs(on_disk - count * isize / 8) <= overhead,
                             "Storage for val %i took %.2f bytes/item, "
@@ -107,11 +103,10 @@ def test_compression_map(self):
         count = 10000
         overhead = 12
         for isize, val in bit_sizes:
-            w = core.G3Writer(test_filename)
-            f = core.G3Frame()
-            f['v'] = core.G3MapInt({str(i): val for i in range(count)})
-            w(f)
-            del w
+            with core.G3Writer(test_filename) as w:
+                f = core.G3Frame()
+                f['v'] = core.G3MapInt({str(i): val for i in range(count)})
+                w(f)
             on_disk = os.path.getsize(test_filename)
             self.assertTrue(abs(on_disk - count * isize / 8) <= (overhead * count),
                             "Storage for val %i took %.2f bytes/item, "
@@ -123,11 +118,10 @@ def test_compression_map_vector(self):
         count = 10000
         overhead = 200
         for isize, val in bit_sizes:
-            w = core.G3Writer(test_filename)
-            f = core.G3Frame()
-            f['v'] = core.G3MapVectorInt({'a': [val] * count})
-            w(f)
-            del w
+            with core.G3Writer(test_filename) as w:
+                f = core.G3Frame()
+                f['v'] = core.G3MapVectorInt({'a': [val] * count})
+                w(f)
             on_disk = os.path.getsize(test_filename)
             self.assertTrue(abs(on_disk - count * isize / 8) <= overhead,
                             "Storage for val %i took %.2f bytes/item, "

doc/fileio.rst

@@ -56,6 +56,17 @@ All frames emitted by ``dosomethings`` will be written to ``/path/to/file.g3``.
 
 Like G3Reader, G3Writer supports gzip and will write compressed output if its output file name ends in ``.gz``.
 
+G3Writer can also be used outside of a pipeline as a context to write particular frames to file.  For example:
+
+.. code-block:: python
+
+	with core.G3Writer(filename='/path/to/file.g3') as writer:
+		writer(frame1)
+		writer(frame2)
+
+This will open the file, write both ``frame1`` and ``frame2``, then properly close the file.
+
+
 G3MultiFileWriter
 =================
 
@@ -97,6 +108,8 @@ For more complex cases, you can also pass a callable as divide_on that takes a f
 	pipe.Add(core.G3MultiFileWriter, filename='/path/to/file-%02u.g3', size_limit = 1024**3, divide_on=lambda frame: frame.type in [core.G3FrameType.Observation])
 	pipe.Run()
 
+G3MultiFileWriter can also be used as a python context object, just like G3Writer.
+
 G3File
 ======