major refactoring to improve readability

tacview-splitter.py

@@ -1,75 +1,40 @@
 #!/usr/bin/env ipython
 
+from __future__ import annotations
+from dataclasses import dataclass
 from os import listdir as os_listdir
-import zipfile
+from typing import Tuple
+from zipfile import ZipFile, ZIP_DEFLATED
+
 
 EXTENSION_TXT = '.txt.acmi'
 EXTENSION_ZIP = '.zip.acmi'
 
-is_zip = None
+NO_EXTENSION = 'no_extension'
 
-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
+def main():
+    filename_input, is_zip = find_input_file()
     print('Processing ' + str(filename_input))
+    filenames = Filenames(filename_input, is_zip)
+    tacview_lines = read_data(filenames)
 
+    # set up all the file descriptors we will need
     if is_zip:
-    filename_no_extension = filename_input.replace(EXTENSION_ZIP, '')
+        fd_blue_zip = ZipFile(filenames.output.blue.zip, 'w', ZIP_DEFLATED)
+        fd_red_zip = ZipFile(filenames.output.red.zip, 'w', ZIP_DEFLATED)
+        fd_violet_zip = ZipFile(filenames.output.violet.zip, 'w', ZIP_DEFLATED)
+
+        fd_blue_txt = fd_blue_zip.open(filenames.output.blue.txt, 'w')
+        fd_red_txt = fd_red_zip.open(filenames.output.blue.txt, 'w')
+        fd_violet_txt = fd_violet_zip.open(filenames.output.blue.txt, 'w')
     else:
-    filename_no_extension = filename_input.replace(EXTENSION_TXT, '')
+        fd_blue_txt = open(filenames.output.blue.txt, 'w')
+        fd_red_txt = open(filenames.output.red.txt, 'w')
+        fd_violet_txt = open(filenames.output.violet.txt, 'w')
 
-filename_blue_no_extension, filename_red_no_extension, filename_violet_no_extension = \
+    # replicate the header from the input file into the output files
-        (f'{filename_no_extension}_{color}' for color in ('blue', 'red', 'violet'))
+    for i, line in enumerate(tacview_lines):
-
-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:
@@ -83,8 +48,7 @@ for i, line in enumerate(tacview_raw_lines):
     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
+    tacview_lines = tacview_lines[i:]  # remove the header, we don't need it anymore
-
     blue_ids = []
     red_ids = []
     # violet = neutral faction, used for chaffs, flares, decoys and shrapnel
@@ -94,8 +58,9 @@ red_ids = []
     violet_ids = []
     undecided_ids = []
 
+    # core routine: process all the lines, put them in the correct output file
     continued = False
-for line in tacview_raw_lines:
+    for line in tacview_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
@@ -122,6 +87,7 @@ for line in tacview_raw_lines:
             line_output = line.encode()
         else:
             line_output = line
+        # code checker thinks that id_ can be unbound because, which it cannot
         # noinspection PyUnboundLocalVariable
         if id_ in blue_ids:
             fd_blue_txt.write(line_output)
@@ -138,19 +104,89 @@ for line in tacview_raw_lines:
             continued = True
         else:
             continued = False
-
+    # main work completed, close all descriptors
     fd_blue_txt.close()
     fd_red_txt.close()
     fd_violet_txt.close()
-
     if is_zip:
+        # there is no way the variable is unbound when we are in this branch
         # noinspection PyUnboundLocalVariable
         fd_blue_zip.close()
         # noinspection PyUnboundLocalVariable
         fd_red_zip.close()
         # noinspection PyUnboundLocalVariable
         fd_violet_zip.close()
-
+    # sanity check
     if len(undecided_ids) != 0:
         print('There were units that are neither BLUE, RED nor NEUTRAL. Please investigate.')
         print(undecided_ids)
+
+
+def read_data(filenames: Filenames) -> list[str]:
+    if filenames.input.is_zip:
+        with ZipFile(filenames.input.zip) as fd_zip:
+            with fd_zip.open(filenames.input.txt) as fd_tacview:
+                tacview_lines_binary = fd_tacview.readlines()
+        tacview_lines_ascii = []
+        for line in tacview_lines_binary:
+            tacview_lines_ascii.append(line.decode())
+    else:
+        with open(filenames.input.txt) as fd_tacview:
+            tacview_lines_ascii = fd_tacview.readlines()
+    return tacview_lines_ascii
+
+
+def find_input_file() -> Tuple[str, bool]:
+    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.')
+    return filename, is_zip
+
+
+@dataclass
+class Filenames:
+    @dataclass
+    class _Input:
+        def __init__(self, filename_input: str, is_zip: bool):
+            self.zip: str = ''
+            self.txt: str = ''
+            self.no_extension: str = ''
+            self.is_zip: bool = is_zip
+            if is_zip:
+                self.zip = filename_input
+                self.txt = filename_input.replace(EXTENSION_ZIP, EXTENSION_TXT)
+                self.no_extension = filename_input.replace(EXTENSION_ZIP, '')
+            else:
+                self.zip = ''
+                self.txt = filename_input
+                self.no_extension = filename_input.replace(EXTENSION_TXT, '')
+
+    @dataclass
+    class _Output:
+        @dataclass
+        class _Coalition:
+            def __init__(self, filename_input_no_extension: str, color: str):
+                self.no_extension: str = f'{filename_input_no_extension}_{color}'
+                self.zip: str = f'{self.no_extension}{EXTENSION_ZIP}'
+                self.txt: str = f'{self.no_extension}{EXTENSION_TXT}'
+
+        def __init__(self, filename_input_no_extension: str):
+            self.blue = self._Coalition(filename_input_no_extension, 'blue')
+            self.red = self._Coalition(filename_input_no_extension, 'red')
+            self.violet = self._Coalition(filename_input_no_extension, 'violet')
+
+    def __init__(self, filename_input: str, is_zip: bool):
+        self.input = self._Input(filename_input, is_zip)
+        self.output = self._Output(self.input.no_extension)
+
+
+if __name__ == '__main__':
+    main()