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:
+EXTENSION_TXT = '.txt.acmi'
-    tacview_raw_lines = fd.readlines()
+EXTENSION_ZIP = '.zip.acmi'
+
+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')
 
-fb = open('tacview_blue.txt.acmi', 'w')
-fr = open('tacview_red.txt.acmi', 'w')
-fv = open('tacview_violet.txt.acmi', '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)
+        fd_blue_txt.write(line_header)
-        fr.write(line)
+        fd_red_txt.write(line_header)
-        fv.write(line)
+        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
 
@@ -22,45 +89,68 @@ blue_ids = []
 red_ids = []
 # violet = neutral faction, used for chaffs, flares, decoys and shrapnel
 # we can't decide easily which faction they belong to
-# we would need to find the blue or red object with the least distance 
+# we would need to find the blue or red object with the least distance
 # to violet objects around their spawn time
-violet_ids = []  
+violet_ids = []
 undecided_ids = []
 
+continued = False
 for line in tacview_raw_lines:
-    # the first time a unit appears it has the Color in its line.
+    # tacview introduced continued lines, signified by a single backslash at EOL
-    # The first part of the line (before the first comma) is the unique ID of the unit.
+    # example: DCS briefing is copied into tacview file (begins with `0,Briefing=`)
-    if line[0] == '#':  # check if line is a time stamp
+    # if the line was not continued, we need to extract the unit ID from the line
-        id_ = 'both'
+    # otherwise we reuse the ID from the previous loop
-    elif line[0] == '-':  # negative id is used to indicate a destroyed unit
+    if not continued:
-        id_ = line[1:].strip()
+        # the first time a unit appears it has the Color in its line.
-    else:  # otherwise it is a standard unit line
+        # The first part of the line (before the first comma) is the unique ID of the unit.
-        id_, rest = line.split(',', 1) 
+        if line[0] == '#':  # check if line is a time stamp
-    if 'Color=Red' in line: 
+            id_ = 'both'
-        red_ids.append(id_)
+        elif line[0] == '-':  # negative id is used to indicate a destroyed unit
-    elif 'Color=Blue' in line:
+            id_ = line[1:].strip()
-        blue_ids.append(id_)
+        else:  # otherwise it is a standard unit line
-    elif 'Color=Violet' in line:
+            id_, rest = line.split(',', 1)
-        violet_ids.append(id_)
+        if 'Color=Red' in line:
-    elif 'Color=' in line: 
+            red_ids.append(id_)
-        undecided_ids.append(id_)
+        elif 'Color=Blue' in line:
+            blue_ids.append(id_)
+        elif 'Color=Violet' in line:
+            violet_ids.append(id_)
+        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)
+        fd_blue_txt.write(line_output)
-        fr.write(line)
+        fd_red_txt.write(line_output)
-        fv.write(line)
+        fd_violet_txt.write(line_output)
 
-fb.close()
+    if line.endswith('\\\n'):
-fr.close()
+        continued = True
-fv.close()
+    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)
-