even more refactoring to improve code quality

tacview-splitter.py

@@ -16,49 +16,29 @@ def main():
     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:
+    descriptors = Descriptors(filenames)
-        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:
-        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')
-
     # replicate the header from the input file into the output files
-    for i, line in enumerate(tacview_lines):
+    tacview_lines_no_header = move_header_to_output_files(tacview_lines, descriptors)
-        if is_zip:
+    undecided_ids = move_content_to_output_files(tacview_lines_no_header, descriptors)
-            line_header = line.encode()
+    descriptors.close()
-        else:
+    # sanity check
-            line_header = line
+    if len(undecided_ids) != 0:
-        if not line[0] == '#':  # header is everything before the first '#'
+        print('There were units that are neither BLUE, RED nor NEUTRAL. Please investigate.')
-            fd_blue_txt.write(line_header)
+        print(undecided_ids)
-            fd_red_txt.write(line_header)
-            fd_violet_txt.write(line_header)
-        else:
-            break
-    else:
-        raise IOError('Tacview file seems to be empty')
 
-    tacview_lines = tacview_lines[i:]  # remove the header, we don't need it anymore
+
+def move_content_to_output_files(tacview_lines_no_header: list[str], descriptors: Descriptors) -> list[str]:
+    # core routine: process all the lines, put them in the correct output file
     blue_ids = []
     red_ids = []
-    # violet = neutral faction, used for chaffs, flares, decoys and shrapnel
+    # violet = neutral faction, used for chaffs, flares, decoys and shrapnel. we can't decide easily which faction
-    # we can't decide easily which faction they belong to
+    # they belong to. we would need to find the blue or red object with the least distance to violet objects
-    # we would need to find the blue or red object with the least distance
+    # around their spawn time
-    # to violet objects around their spawn time
     violet_ids = []
     undecided_ids = []
-
-    # core routine: process all the lines, put them in the correct output file
     continued = False
-    for line in tacview_lines:
+    for line in tacview_lines_no_header:
         # 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
@@ -81,46 +61,56 @@ def main():
             elif 'Color=' in line:
                 undecided_ids.append(id_)
 
-        if is_zip:
+        if descriptors.filenames.input.is_zip:
             line_output = line.encode()
         else:
             line_output = line
-        # code checker thinks that id_ can be unbound because, which it cannot
+        # code checker thinks that id_ can be unbound, which it cannot
         # noinspection PyUnboundLocalVariable
-        if id_ in blue_ids:
+        if id_ in blue_ids or id_ == 'both':  # 'both@ refers to timestamps
-            fd_blue_txt.write(line_output)
+            descriptors.blue_txt.write(line_output)
-        elif id_ in red_ids:
+        if id_ in red_ids or id_ == 'both':
-            fd_red_txt.write(line_output)
+            descriptors.red_txt.write(line_output)
-        elif id_ in violet_ids:
+        if id_ in violet_ids or id_ == 'both':
-            fd_violet_txt.write(line_output)
+            descriptors.violet_txt.write(line_output)
-        else:  # id_ == 'both', has timestamps
-            fd_blue_txt.write(line_output)
-            fd_red_txt.write(line_output)
-            fd_violet_txt.write(line_output)
 
         if line.endswith('\\\n'):
             continued = True
         else:
             continued = False
-    # main work completed, close all descriptors
+    return undecided_ids
-    fd_blue_txt.close()
+
-    fd_red_txt.close()
+
-    fd_violet_txt.close()
+def move_header_to_output_files(tacview_lines: list[str], descriptors: Descriptors) -> list[str]:
-    if is_zip:
+    """
-        # there is no way the variable is unbound when we are in this branch
+    finds the tacview header in tacview_lines and writes it to the three output files
-        # noinspection PyUnboundLocalVariable
+    afterwards removes the header from the input data and returns the remaining content (actual telemetry)
-        fd_blue_zip.close()
+    :param tacview_lines: content of tacview file with header
-        # noinspection PyUnboundLocalVariable
+    :param descriptors: object of Descriptor class
-        fd_red_zip.close()
+    :return: content of tacview file without header
-        # noinspection PyUnboundLocalVariable
+    """
-        fd_violet_zip.close()
+    for i, line in enumerate(tacview_lines):
-    # sanity check
+        if descriptors.filenames.input.is_zip:
-    if len(undecided_ids) != 0:
+            line_header = line.encode()
-        print('There were units that are neither BLUE, RED nor NEUTRAL. Please investigate.')
+        else:
-        print(undecided_ids)
+            line_header = line
+        if not line[0] == '#':  # header is everything before the first '#'
+            descriptors.blue_txt.write(line_header)
+            descriptors.red_txt.write(line_header)
+            descriptors.violet_txt.write(line_header)
+        else:
+            break
+    else:
+        raise IOError('Tacview file seems to be empty')
+    return tacview_lines[i:]  # remove the header, we don't need it anymore
 
 
 def read_data(filenames: Filenames) -> list[str]:
+    """
+    get the tacview data out of the file
+    :param filenames: object of class Filenames
+    :return: file content as list of strings, one line per item
+    """
     if filenames.input.is_zip:
         with ZipFile(filenames.input.zip) as fd_zip:
             with fd_zip.open(filenames.input.txt) as fd_tacview:
@@ -186,5 +176,37 @@ class Filenames:
         self.output = self._Output(self.input.no_extension)
 
 
+@dataclass
+class Descriptors:
+    def __init__(self, filenames: Filenames):
+        self.filenames = filenames
+        if filenames.input.is_zip:
+            self._blue_zip = ZipFile(filenames.output.blue.zip, 'w', ZIP_DEFLATED)
+            self._red_zip = ZipFile(filenames.output.red.zip, 'w', ZIP_DEFLATED)
+            self._violet_zip = ZipFile(filenames.output.violet.zip, 'w', ZIP_DEFLATED)
+
+            self.blue_txt = self._blue_zip.open(filenames.output.blue.txt, 'w')
+            self.red_txt = self._red_zip.open(filenames.output.blue.txt, 'w')
+            self.violet_txt = self._violet_zip.open(filenames.output.blue.txt, 'w')
+        else:
+            self._blue_zip = None
+            self._red_zip = None
+            self._violet_zip = None
+
+            self.blue_txt = open(filenames.output.blue.txt, 'w')
+            self.red_txt = open(filenames.output.red.txt, 'w')
+            self.violet_txt = open(filenames.output.violet.txt, 'w')
+
+    def close(self):
+        self.blue_txt.close()
+        self.red_txt.close()
+        self.violet_txt.close()
+
+        if self.filenames.input.is_zip:
+            self._blue_zip.close()
+            self._red_zip.close()
+            self._violet_zip.close()
+
+
 if __name__ == '__main__':
     main()