Merge branch 'HLA-809-chunking' into HLA-812-testing

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

@@ -27,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)
 
@@ -248,13 +248,15 @@ def write_pb(self, filepath: PathLike):
             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)
@@ -276,6 +278,7 @@ def pb_2_txt():
 
     pb = PBUtils(pb_file)
     pb.write_to_txt(txt_file)
+    print("Write completed!")
 
 
 def print_header():