1 files changed, 9 insertions, 6 deletions
diff --git a/elf.py b/elf.py
index 4427263..2759679 100755
--- a/elf.py
+++ b/elf.py
@@ -80,12 +80,11 @@ def write_program_header(segment, c, start):
     c += wi(align, 4)  # p_align
     return c
 
-if __name__ == "__main__":
+@subv.with_parsed_lines
+def elf(iter):
     segments = []
     segment = None
-    for line in sys.stdin:
-        line = subv.parse(line)
-
+    for line in iter:
         if line['type'] == 'segment':
             (name, addr) = line['segment']
             segment = { 'name': name, 'addr': addr, 'content': [] }
@@ -98,8 +97,6 @@ if __name__ == "__main__":
         else:
             raise ValueError("elf input should contain only segments and data!")
 
-    cursor = 0
-
     c = write_elf_header(segments)
     segment_start = 0x1000
     for seg in segments:
@@ -110,3 +107,9 @@ if __name__ == "__main__":
         c = padto(seg['offset'], c)
         for part in seg['content']:
             c += wi(part[0], 1)
+
+    yield
+
+if __name__ == '__main__':
+    for line in elf(sys.stdin):
+        pass