Merge pull request #99 from LSSTDESC/issue/96/gzipped_files_includeobj

handle gzipped instance catalogs and includeobj entries

Closes #96

Author: jchiang87

python/desc/imsim/__init__.py

@@ -8,3 +8,4 @@
 from .camera_readout import *
 from .focalplane_info import *
 from .skyModel import *
+from .fopen import *

python/desc/imsim/fopen.py

@@ -0,0 +1,54 @@
+"""
+Utilities to read in gzipped instance catalogs with includeobj directives.
+"""
+import os
+import contextlib
+import gzip
+
+__all__ = ["fopen"]
+
+
+@contextlib.contextmanager
+def fopen(filename, **kwds):
+    """
+    Return a file descriptor-like object that closes the underlying
+    file descriptor when used with the with-statement.
+
+    Parameters
+    ----------
+    filename: str
+        Filename of the instance catalog.
+    **kwds: dict
+        Keyword arguments to pass to the gzip.open or open functions.
+
+    Returns
+    -------
+    generator: file descriptor-like generator object that can be iterated
+        over to return the lines in a file.
+    """
+    abspath = os.path.split(os.path.abspath(filename))[0]
+    try:
+        if filename.endswith('.gz'):
+            fd = gzip.open(filename, **kwds)
+        else:
+            fd = open(filename, **kwds)
+        yield fopen_generator(fd, abspath, **kwds)
+    finally:
+        fd.close()
+
+
+def fopen_generator(fd, abspath, **kwds):
+    """
+    Return a generator for the provided file descriptor that knows how
+    to recursively read in instance catalogs specified by the
+    includeobj directive.
+    """
+    with fd as input_:
+        for line in input_:
+            if not line.startswith('includeobj'):
+                yield line
+            else:
+                filename = os.path.join(abspath, line.strip().split()[-1])
+                with fopen(filename, **kwds) as my_input:
+                    for line in my_input:
+                        yield line

python/desc/imsim/imSim.py

@@ -36,7 +36,8 @@
 from lsst.sims.GalSimInterface import GalSimCelestialObject
 from lsst.sims.photUtils import BandpassDict, Sed, getImsimFluxNorm
 from lsst.sims.utils import defaultSpecMap
-from desc.imsim import CosmicRays
+from .cosmic_rays import CosmicRays
+from .fopen import fopen
 
 _POINT_SOURCE = 1
 _SERSIC_2D = 2
@@ -103,7 +104,7 @@ def metadata_from_file(file_name):
     InstanceCatalog.
     """
     input_params = {}
-    with open(file_name, 'r') as in_file:
+    with fopen(file_name, mode='rt') as in_file:
         for line in in_file:
             if line[0] == '#':
                 continue
@@ -181,7 +182,7 @@ def sources_from_file(file_name, obs_md, phot_params, numRows=None):
 
     num_objects = 0
     ct_rows = 0
-    with open(file_name, 'r') as input_file:
+    with fopen(file_name, mode='rt') as input_file:
         for line in input_file:
             ct_rows += 1
             params = line.strip().split()
@@ -214,7 +215,7 @@ def sources_from_file(file_name, obs_md, phot_params, numRows=None):
     object_type = np.zeros(num_objects, dtype=int)
 
     i_obj = -1
-    with open(file_name, 'r') as input_file:
+    with fopen(file_name, mode='rt') as input_file:
         for line in input_file:
             params = line.strip().split()
             if params[0] != 'object':

tests/test_fopen.py

@@ -0,0 +1,54 @@
+"Unit tests for fopen function"
+
+import os
+import unittest
+import gzip
+import desc.imsim
+
+class FopenTestCase(unittest.TestCase):
+    "TestCase class for fopen unit tests."
+    def setUp(self):
+        self.test_dir = 'fopen_dir'
+        os.makedirs(self.test_dir, exist_ok=True)
+        self.fopen_test_file \
+            = os.path.join(self.test_dir, 'fopen_test_file.txt')
+        self.fopen_include_file1 \
+            = os.path.join(self.test_dir, 'fopen_include_file1.txt.gz')
+        self.fopen_include_file2 \
+            = os.path.join(self.test_dir, 'fopen_include_file2.txt.gz')
+        self.lines = 'line1 line2 line3 line4'.split()
+        with open(self.fopen_test_file, 'w') as output:
+            output.write('%s\n' % self.lines[0])
+            output.write('%s\n' % self.lines[1])
+            output.write('includeobj %s\n'
+                         % os.path.basename(self.fopen_include_file1))
+            output.write('includeobj %s\n'
+                         % os.path.basename(self.fopen_include_file2))
+        with gzip.open(self.fopen_include_file1, 'wt') as output:
+            output.write('%s\n' % self.lines[2])
+        with gzip.open(self.fopen_include_file2, 'wt') as output:
+            output.write('%s\n' % self.lines[3])
+
+    def tearDown(self):
+        for item in (self.fopen_test_file, self.fopen_include_file1,
+                     self.fopen_include_file2):
+            try:
+                os.remove(item)
+            except OSError:
+                pass
+        try:
+            os.rmdir(self.test_dir)
+        except OSError:
+            pass
+
+    def test_fopen(self):
+        "Test the fopen function."
+        with desc.imsim.fopen(self.fopen_test_file, mode='rt') as input_:
+            for i, line in enumerate(input_):
+                expected = self.lines[i]
+                self.assertEqual(line.strip(), expected)
+            # Make sure all of the expected lines have been processed.
+            self.assertEqual(self.lines[-1], expected)
+
+if __name__ == '__main__':
+    unittest.main()