rewrite of core functionality

Can handle continued lines (introduced recently by Tacview)
Processes the first file of matching type it finds in its current dir
Handles both zip and txt files

tacview-splitter.py

@@ -1,20 +1,87 @@
 #!/usr/bin/env ipython
 
+from os import listdir as os_listdir
+import zipfile
 
-with open('Tacview-20200419-200244-DCS-OPUF_MSN18.txt.acmi') as fd:
-    tacview_raw_lines = fd.readlines()
+EXTENSION_TXT = '.txt.acmi'
+EXTENSION_ZIP = '.zip.acmi'
 
-fb = open('tacview_blue.txt.acmi', 'w')
-fr = open('tacview_red.txt.acmi', 'w')
-fv = open('tacview_violet.txt.acmi', 'w')
+is_zip = None
+
+all_files = os_listdir('.')
+for filename in all_files:
+    filename_lower = filename.lower()
+    if filename_lower.endswith(EXTENSION_ZIP):
+        is_zip = True
+        break
+    elif filename_lower.endswith(EXTENSION_TXT):
+        is_zip = False
+        break
+else:
+    raise FileNotFoundError('Could not find a tacview file in this directory.')
+
+filename_input = filename
+print('Processing ' + str(filename_input))
+
+if is_zip:
+    filename_no_extension = filename_input.replace(EXTENSION_ZIP, '')
+else:
+    filename_no_extension = filename_input.replace(EXTENSION_TXT, '')
+
+filename_blue_no_extension, filename_red_no_extension, filename_violet_no_extension = \
+        (f'{filename_no_extension}_{color}' for color in ('blue', 'red', 'violet'))
+
+if is_zip:
+    filename_blue_zip, filename_red_zip, filename_violet_zip = \
+        (f'{arg}{EXTENSION_ZIP}' for arg in
+         (filename_blue_no_extension, filename_red_no_extension, filename_violet_no_extension)
+         )
+
+filename_blue_txt, filename_red_txt, filename_violet_txt = \
+    (f'{arg}{EXTENSION_TXT}' for arg in
+     (filename_blue_no_extension, filename_red_no_extension, filename_violet_no_extension)
+     )
+
+if is_zip:
+    with zipfile.ZipFile(filename_input) as fd_zip:
+        with fd_zip.open(filename_no_extension + EXTENSION_TXT) as fd_tacview:
+            tacview_binary_lines = fd_tacview.readlines()
+    tacview_raw_lines = []
+    for line in tacview_binary_lines:
+        tacview_raw_lines.append(line.decode())
+else:
+    with open(filename_input) as fd_tacview:
+        tacview_raw_lines = fd_tacview.readlines()
+
+if is_zip:
+    # noinspection PyUnboundLocalVariable
+    fd_blue_zip = zipfile.ZipFile(filename_blue_zip, 'w', zipfile.ZIP_DEFLATED)
+    # noinspection PyUnboundLocalVariable
+    fd_red_zip = zipfile.ZipFile(filename_red_zip, 'w', zipfile.ZIP_DEFLATED)
+    # noinspection PyUnboundLocalVariable
+    fd_violet_zip = zipfile.ZipFile(filename_violet_zip, 'w', zipfile.ZIP_DEFLATED)
+
+    fd_blue_txt = fd_blue_zip.open(filename_blue_txt, 'w')
+    fd_red_txt = fd_red_zip.open(filename_red_txt, 'w')
+    fd_violet_txt = fd_violet_zip.open(filename_violet_txt, 'w')
+else:
+    fd_blue_txt = open(filename_blue_txt, 'w')
+    fd_red_txt = open(filename_red_txt, 'w')
+    fd_violet_txt = open(filename_violet_txt, 'w')
 
 for i, line in enumerate(tacview_raw_lines):
+    if is_zip:
+        line_header = line.encode()
+    else:
+        line_header = line
     if not line[0] == '#':  # header is everything before the first '#'
-        fb.write(line)
-        fr.write(line)
-        fv.write(line)
+        fd_blue_txt.write(line_header)
+        fd_red_txt.write(line_header)
+        fd_violet_txt.write(line_header)
     else:
         break
+else:
+    raise IOError('Tacview file seems to be empty')
 
 tacview_raw_lines = tacview_raw_lines[i:]  # remove the header, we don't need it anymore
 
@@ -27,7 +94,13 @@ red_ids = []
 violet_ids = []
 undecided_ids = []
 
+continued = False
 for line in tacview_raw_lines:
+    # tacview introduced continued lines, signified by a single backslash at EOL
+    # example: DCS briefing is copied into tacview file (begins with `0,Briefing=`)
+    # if the line was not continued, we need to extract the unit ID from the line
+    # otherwise we reuse the ID from the previous loop
+    if not continued:
         # the first time a unit appears it has the Color in its line.
         # The first part of the line (before the first comma) is the unique ID of the unit.
         if line[0] == '#':  # check if line is a time stamp
@@ -45,22 +118,39 @@ for line in tacview_raw_lines:
         elif 'Color=' in line:
             undecided_ids.append(id_)
 
+    if is_zip:
+        line_output = line.encode()
+    else:
+        line_output = line
+    # noinspection PyUnboundLocalVariable
     if id_ in blue_ids:
-        fb.write(line)
+        fd_blue_txt.write(line_output)
     elif id_ in red_ids:
-        fr.write(line)
+        fd_red_txt.write(line_output)
     elif id_ in violet_ids:
-        fv.write(line)
+        fd_violet_txt.write(line_output)
     else:  # id_ == 'both', has timestamps
-        fb.write(line)
-        fr.write(line)
-        fv.write(line)
+        fd_blue_txt.write(line_output)
+        fd_red_txt.write(line_output)
+        fd_violet_txt.write(line_output)
 
-fb.close()
-fr.close()
-fv.close()
+    if line.endswith('\\\n'):
+        continued = True
+    else:
+        continued = False
+
+fd_blue_txt.close()
+fd_red_txt.close()
+fd_violet_txt.close()
+
+if is_zip:
+    # noinspection PyUnboundLocalVariable
+    fd_blue_zip.close()
+    # noinspection PyUnboundLocalVariable
+    fd_red_zip.close()
+    # noinspection PyUnboundLocalVariable
+    fd_violet_zip.close()
 
 if len(undecided_ids) != 0:
-    print('There were units that are neither BLUEFOR, REDFOR nor NEUTRAL. Please investigate.')
+    print('There were units that are neither BLUE, RED nor NEUTRAL. Please investigate.')
     print(undecided_ids)
-