Tidy up generators into contained objects

Neaten up the generators by returning an iterable object that
has the headers as attributes.

Change-Id: I3c47e590584bfa963e2214bf61c50e85600e1ed2

os_loganalyze/filter.py

@@ -96,11 +96,11 @@ class LogLine(object):
 
 class Filter(object):
 
-    def __init__(self, fname, generator, minsev="NONE", limit=None):
+    def __init__(self, file_generator, minsev="NONE", limit=None):
         self.minsev = minsev
-        self.gen = generator
+        self.file_generator = file_generator
-        self.supports_sev = SUPPORTS_SEV.search(fname) is not None
+        self.supports_sev = \
-        self.fname = fname
+            SUPPORTS_SEV.search(file_generator.logname) is not None
         self.limit = limit
         self.strip_control = False
 
@@ -110,9 +110,9 @@ class Filter(object):
     def __iter__(self):
         old_sev = "NONE"
         lineno = 1
-        for line in self.gen:
+        for line in self.file_generator:
             # bail early for limits
-            if self.limit and lineno >= int(self.limit):
+            if self.limit and lineno > int(self.limit):
                 raise StopIteration()
             # strip control chars in case the console is ascii colored
             if self.strip_control:

os_loganalyze/generator.py

@@ -16,6 +16,7 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+import collections
 import fileinput
 import os.path
 import re
@@ -104,73 +105,99 @@ def _get_swift_connection(swift_config):
 _get_swift_connection.con = None
 
 
-def get_swift_line_generator(logname, config):
+class SwiftIterableBuffer(collections.Iterable):
-    resp_headers = {}
+    file_headers = {}
-    if not config.has_section('swift'):
-        sys.stderr.write('Not configured to use swift..\n')
-        sys.stderr.write('logname: %s\n' % logname)
-        return resp_headers, None
 
-    try:
+    def __init__(self, logname, config):
-        swift_config = dict(config.items('swift'))
+        self.logname = logname
-        con = _get_swift_connection(swift_config)
+        self.resp_headers = {}
+        self.obj = None
+        self.file_headers['filename'] = logname
 
-        chunk_size = int(swift_config.get('chunk_size', 64))
+        if not config.has_section('swift'):
-        if chunk_size < 1:
+            sys.stderr.write('Not configured to use swift..\n')
-            chunk_size = None
+            sys.stderr.write('logname: %s\n' % logname)
+        else:
+            try:
+                swift_config = dict(config.items('swift'))
+                # NOTE(jhesketh): While _get_siwft_connection seems like it
+                # should be part of this class we actually still need it
+                # outside to maintain the connection across multiple objects.
+                # Each SwiftIterableBuffer is a new object request, not
+                # necessarily a new swift connection (hopefully we can reuse
+                # connections). I think the place to put the get connection
+                # in the future would be in the server.py (todo).
+                con = _get_swift_connection(swift_config)
 
-        resp_headers, obj = con.get_object(
+                chunk_size = int(swift_config.get('chunk_size', 64))
-            swift_config['container'], logname,
+                if chunk_size < 1:
-            resp_chunk_size=chunk_size)
+                    chunk_size = None
 
-        def line_generator():
+                self.resp_headers, self.obj = con.get_object(
-            ext = os.path.splitext(logname)[1]
+                    swift_config['container'], logname,
-            if ext == '.gz':
+                    resp_chunk_size=chunk_size)
-                # Set up a decompression object assuming the deflate
+                self.file_headers.update(self.resp_headers)
-                # compression algorithm was used
+            except Exception:
-                d = zlib.decompressobj(16 + zlib.MAX_WBITS)
+                import traceback
+                sys.stderr.write("Error fetching from swift.\n")
+                sys.stderr.write('logname: %s\n' % logname)
+                traceback.print_exc()
 
-            if isinstance(obj, types.GeneratorType):
+    def __iter__(self):
-                buf = next(obj)
+        ext = os.path.splitext(self.logname)[1]
-                partial = ''
+        if ext == '.gz':
-                while buf:
+            # Set up a decompression object assuming the deflate
-                    if ext == '.gz':
+            # compression algorithm was used
-                        string = partial + d.decompress(buf)
+            d = zlib.decompressobj(16 + zlib.MAX_WBITS)
-                    else:
+
-                        string = partial + buf
+        if isinstance(self.obj, types.GeneratorType):
-                    split = string.split('\n')
+            buf = next(self.obj)
-                    for line in split[:-1]:
+            partial = ''
-                        yield line + '\n'
+            while buf:
-                    partial = split[-1]
-                    try:
-                        buf = next(obj)
-                    except StopIteration:
-                        break
-                if partial != '':
-                    yield partial
-            else:
-                output = obj
                 if ext == '.gz':
-                    output = d.decompress(output)
+                    string = partial + d.decompress(buf)
-
+                else:
-                split = output.split('\n')
+                    string = partial + buf
+                split = string.split('\n')
                 for line in split[:-1]:
                     yield line + '\n'
                 partial = split[-1]
-                if partial != '':
+                try:
-                    yield partial
+                    buf = next(self.obj)
+                except StopIteration:
+                    break
+            if partial != '':
+                yield partial
+        else:
+            output = self.obj
+            if ext == '.gz':
+                output = d.decompress(output)
 
-        return resp_headers, line_generator()
+            split = output.split('\n')
-
+            for line in split[:-1]:
-    except Exception:
+                yield line + '\n'
-        import traceback
+            partial = split[-1]
-        sys.stderr.write("Error fetching from swift.\n")
+            if partial != '':
-        sys.stderr.write('logname: %s\n' % logname)
+                yield partial
-        traceback.print_exc()
-        return resp_headers, None
 
 
-def get(environ, root_path, config=None):
+class DiskIterableBuffer(collections.Iterable):
+    file_headers = {}
+
+    def __init__(self, logname, logpath, config):
+        self.logname = logname
+        self.logpath = logpath
+        self.resp_headers = {}
+        self.obj = fileinput.FileInput(self.logpath,
+                                       openhook=fileinput.hook_compressed)
+        self.file_headers['filename'] = logname
+        self.file_headers.update(util.get_headers_for_file(logpath))
+
+    def __iter__(self):
+        return self.obj
+
+
+def get_file_generator(environ, root_path, config=None):
     logname = log_name(environ)
     logpath = safe_path(root_path, logname)
     file_headers = {}
@@ -178,24 +205,19 @@ def get(environ, root_path, config=None):
         raise UnsafePath()
     file_headers['filename'] = os.path.basename(logpath)
 
-    flines_generator = None
+    file_generator = None
     # if we want swift only, we'll skip processing files
     use_files = (util.parse_param(environ, 'source', default='all')
                  != 'swift')
     if use_files and does_file_exist(logpath):
-        flines_generator = fileinput.FileInput(
+        file_generator = DiskIterableBuffer(logname, logpath, config)
-            logpath, openhook=fileinput.hook_compressed)
-        file_headers.update(util.get_headers_for_file(logpath))
     else:
-        resp_headers, flines_generator = get_swift_line_generator(logname,
+        file_generator = SwiftIterableBuffer(logname, config)
-                                                                  config)
+        if not file_generator.obj:
-        if not flines_generator:
             logname = os.path.join(logname, 'index.html')
-            resp_headers, flines_generator = get_swift_line_generator(logname,
+            file_generator = SwiftIterableBuffer(logname, config)
-                                                                      config)
-        file_headers.update(resp_headers)
 
-    if not flines_generator:
+    if not file_generator.obj:
         raise NoSuchFile()
 
-    return logname, flines_generator, file_headers
+    return file_generator

os_loganalyze/tests/test_views.py

@@ -28,8 +28,9 @@ class TestViews(base.TestCase):
         # wsgi application. We just need the generator to give to Views.
         root_path = base.samples_path(self.samples_directory)
         kwargs = {'PATH_INFO': '/htmlify/%s' % fname}
-        logname, gen, headers = osgen.get(self.fake_env(**kwargs), root_path)
+        file_generator = osgen.get_file_generator(self.fake_env(**kwargs),
-        flines_generator = osfilter.Filter(logname, gen)
+                                                  root_path)
+        flines_generator = osfilter.Filter(file_generator)
         return flines_generator
 
     def test_html_detection(self):

os_loganalyze/view.py

@@ -161,7 +161,8 @@ class HTMLView(collections.Iterable):
         return newline
 
     def __iter__(self):
-        first_line = next(x for x in self.gen)
+        igen = (x for x in self.gen)
+        first_line = next(igen)
         self._discover_html(first_line.line)
 
         if not self.is_html:
@@ -175,7 +176,7 @@ class HTMLView(collections.Iterable):
         if first:
             yield first
 
-        for line in self.gen:
+        for line in igen:
             newline = self._process_line(line)
             if newline:
                 yield newline
@@ -198,9 +199,9 @@ class TextView(collections.Iterable):
 class PassthroughView(collections.Iterable):
     headers = []
 
-    def __init__(self, gen, file_headers):
+    def __init__(self, gen):
         self.gen = gen
-        for hn, hv in file_headers.items():
+        for hn, hv in self.gen.file_headers.items():
             self.headers.append((hn, hv))
 
     def __iter__(self):

os_loganalyze/wsgi.py

@@ -99,8 +99,7 @@ def application(environ, start_response, root_path=None,
     status = '200 OK'
 
     try:
-        logname, flines_generator, file_headers = osgen.get(environ, root_path,
+        file_generator = osgen.get_file_generator(environ, root_path, config)
-                                                            config)
     except osgen.UnsafePath:
         status = '400 Bad Request'
         response_headers = [('Content-type', 'text/plain')]
@@ -112,23 +111,21 @@ def application(environ, start_response, root_path=None,
         start_response(status, response_headers)
         return ['File Not Found']
 
-    if use_passthrough_view(file_headers):
+    if use_passthrough_view(file_generator.file_headers):
-        generator = osview.PassthroughView(flines_generator,
+        view_generator = osview.PassthroughView(file_generator)
-                                           file_headers)
     else:
         minsev = util.parse_param(environ, 'level', default="NONE")
         limit = util.parse_param(environ, 'limit')
-        flines_generator = osfilter.Filter(
+        flines_generator = osfilter.Filter(file_generator, minsev, limit)
-            logname, flines_generator, minsev, limit)
         if environ.get('OS_LOGANALYZE_STRIP', None):
             flines_generator.strip_control = True
         if should_be_html(environ):
-            generator = osview.HTMLView(flines_generator)
+            view_generator = osview.HTMLView(flines_generator)
         else:
-            generator = osview.TextView(flines_generator)
+            view_generator = osview.TextView(flines_generator)
 
-    start_response(status, generator.headers)
+    start_response(status, view_generator.headers)
-    return generator
+    return view_generator
 
 
 # for development purposes, makes it easy to test the filter output