Write chunked data to multiple files

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)
 
@@ -215,13 +215,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)