Merge branch 'HLA-809-chunking' into HLA-810-remove-data

pyproject.toml

@@ -43,7 +43,7 @@ aa-reduce-data-freq = "aa_remove_data.remove_data:aa_reduce_freq"
 aa-remove-data = "aa_remove_data.__main__:main"
 aa-remove-data-every-nth = "aa_remove_data.remove_data:aa_remove_every_nth"
 aa-print-header = "aa_remove_data.pb_utils:print_header"
-PB_2_TXT = "aa_remove_data.pb_utils:pb_2_txt"
+pb-2-txt = "aa_remove_data.pb_utils:pb_2_txt"
 
 [project.urls]
 GitHub = "https://github.com/DiamondLightSource/aa-remove-data"

src/aa_remove_data/pb_utils.py

@@ -19,8 +19,6 @@ def __init__(self, filepath: PathLike | None = None, chunk_size=10000000):
             read. Defaults to None.
             chunk_size (Optional[int], optional): Number of lines to read/write
             at one time.
-            chunk_size (Optional[int], optional): Number of lines to read/write
-            at one time.
         """
         self.header = EPICSEvent_pb2.PayloadInfo()  # type: ignore
         self.samples = []
@@ -29,7 +27,7 @@ def __init__(self, filepath: PathLike | None = None, chunk_size=10000000):
         self.read_done = False
         self._chunk_size = chunk_size
         self._start_line = 0
-        self._write_started = False
+        self._write_started = []
         if filepath:
             self.read_pb(filepath)
 
@@ -214,17 +212,20 @@ def write_pb(self, filepath: PathLike):
         Args:
             filepath (PathLike): Path to file to be written.
         """
+        header_b = self._replace_newline_chars(self.header.SerializeToString()) + b"\n"
         samples_b = [
             self._replace_newline_chars(sample.SerializeToString()) + b"\n"
             for sample in self.samples
         ]
-        if self._write_started is False:  # Write header, start new file
+        if (
+            filepath not in self._write_started or self.chunked is False
+        ):  # Write header, start new file
             header_b = (
                 self._replace_newline_chars(self.header.SerializeToString()) + b"\n"
             )
             with open(filepath, "wb") as f:
                 f.writelines([header_b] + samples_b)
-            self._write_started = True
+            self._write_started.append(filepath)
         else:  # Add to existing file
             with open(filepath, "ab") as f:
                 f.writelines(samples_b)
@@ -244,6 +245,7 @@ def pb_2_txt():
         raise ValueError(f"Invalid file extension: '{pb_file.suffix}'. Expected '.pb'.")
     pb = PBUtils(pb_file)
     pb.write_to_txt(txt_file)
+    print("Write completed!")
 
 
 def print_header():