diff --git a/pyproject.toml b/pyproject.toml index c67e382..a65d73b 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/aa_remove_data/pb_utils.py b/src/aa_remove_data/pb_utils.py index 740cf0c..68bac91 100644 --- a/src/aa_remove_data/pb_utils.py +++ b/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():