diff --git a/src/daq/caen/n1081b.py b/src/daq/caen/n1081b.py index 36972fe..bd25e91 100644 --- a/src/daq/caen/n1081b.py +++ b/src/daq/caen/n1081b.py @@ -76,13 +76,14 @@ def _poll_sections(self): self._logger.info(f"No counters in section {section}") continue - self._send_store_message(data) + self._send_store_message(data, section.name) - def _send_store_message(self, data: dict): + def _send_store_message(self, data: dict, section): self.message_out.put( DAQJobMessageStore( store_config=self.config.store_config, daq_job=self, + prefix=section, keys=[f"lemo_{x['lemo']}" for x in data["counters"]], data=[[x["value"] for x in data["counters"]]], ) diff --git a/src/daq/store/csv.py b/src/daq/store/csv.py index cc3a2d0..1615e28 100644 --- a/src/daq/store/csv.py +++ b/src/daq/store/csv.py @@ -10,7 +10,7 @@ from daq.models import DAQJobConfig from daq.store.base import DAQJobStore from daq.store.models import DAQJobMessageStore, DAQJobStoreConfig -from utils.file import add_date_to_file_name +from utils.file import modify_file_path DAQ_JOB_STORE_CSV_FLUSH_INTERVAL_SECONDS = 5 * 60 DAQ_JOB_STORE_CSV_WRITE_BATCH_SIZE = 1000 @@ -48,7 +48,9 @@ def __init__(self, config: Any): def handle_message(self, message: DAQJobMessageStore) -> bool: super().handle_message(message) store_config = cast(DAQJobStoreConfigCSV, message.store_config) - file_path = add_date_to_file_name(store_config.file_path, store_config.add_date) + file_path = modify_file_path( + store_config.file_path, store_config.add_date, message.prefix + ) file, new_file = self._open_csv_file(file_path) # Write headers if the file is new diff --git a/src/daq/store/models.py b/src/daq/store/models.py index 49b3155..fb3ecd4 100644 --- a/src/daq/store/models.py +++ b/src/daq/store/models.py @@ -22,6 +22,7 @@ class DAQJobMessageStore(DAQJobMessage): daq_job: DAQJob keys: list[str] data: list[list[Any]] + prefix: str | None = None @dataclass diff --git a/src/daq/store/root.py b/src/daq/store/root.py index 99ee19b..f8e5ae4 100644 --- a/src/daq/store/root.py +++ b/src/daq/store/root.py @@ -7,7 +7,7 @@ from daq.models import DAQJobConfig from daq.store.base import DAQJobStore from daq.store.models import DAQJobMessageStore, DAQJobStoreConfig -from utils.file import add_date_to_file_name +from utils.file import modify_file_path @dataclass @@ -34,7 +34,9 @@ def __init__(self, config: Any): def handle_message(self, message: DAQJobMessageStore) -> bool: super().handle_message(message) store_config = cast(DAQJobStoreConfigROOT, message.store_config) - file_path = add_date_to_file_name(store_config.file_path, store_config.add_date) + file_path = modify_file_path( + store_config.file_path, store_config.add_date, message.prefix + ) if file_path not in self._open_files: file_exists = os.path.exists(file_path) diff --git a/src/utils/file.py b/src/utils/file.py index 9c9030c..a08dff6 100644 --- a/src/utils/file.py +++ b/src/utils/file.py @@ -2,12 +2,16 @@ from datetime import datetime -def add_date_to_file_name(file_path: str, add_date: bool) -> str: - splitted_file_path = os.path.splitext(file_path) +def modify_file_path(file_path: str, add_date: bool, prefix: str | None) -> str: + split = os.path.splitext(file_path) date_text = datetime.now().strftime("%Y-%m-%d") - if len(splitted_file_path) > 1: - file_path = f"{splitted_file_path[0]}_{date_text}{splitted_file_path[1]}" - else: - file_path = f"{splitted_file_path[0]}_{date_text}" + if add_date: + if len(split) > 1: + file_path = f"{split[0]}_{date_text}{split[1]}" + else: + file_path = f"{split[0]}_{date_text}" + if prefix is not None: + head, tail = os.path.split(file_path) + file_path = os.path.join(head, f"{prefix}_{tail}") return file_path