Bugfix: Trace with 0 valid records ends iteration

1 files changed, 9 insertions, 1 deletions

diff --git a/unit_trace/trace_reader.py b/unit_trace/trace_reader.py
index 62c7ea7..a4c3c05 100644
--- a/unit_trace/trace_reader.py
+++ b/unit_trace/trace_reader.py
@@ -58,7 +58,13 @@ def trace_reader(files):
     for file in files:
         file_iter = _get_file_iter(file)
         file_iters.append(file_iter)
-        file_iter_buff.append([file_iter.next()])
+        try:
+            file_iter_buff.append([file_iter.next()])
+        # What if there isn't a single valid record in a trace file?
+        # file_iter.next() will raise a StopIteration that we need to catch
+        except:
+            # Forget that file iter
+            file_iters.pop()
 
     # We keep 200 records in each buffer and then keep the buffer sorted
     # This is because records may have been recorded slightly out of order
@@ -264,6 +270,8 @@ def _get_type(type_num):
     types = [None,StNameData,StParamData,StReleaseData,StAssignedData,
              StSwitchToData,StSwitchAwayData,StCompletionData,StBlockData,
              StResumeData,StSysReleaseData]
+    if type_num > len(types)-1 or type_num < 1:
+        raise Exception
     return types[type_num]
 
 # Return the type name, given the type_num (this is simply a convenience to