diff --git a/src/pythonequipmentdrivers/utility/data_management.py b/src/pythonequipmentdrivers/utility/data_management.py index 1316713..609c964 100644 --- a/src/pythonequipmentdrivers/utility/data_management.py +++ b/src/pythonequipmentdrivers/utility/data_management.py @@ -5,6 +5,8 @@ from pathlib import Path from time import asctime, strftime from typing import Any, Dict, Iterable, List, Optional +import shutil +import importlib __all__ = ("log_to_csv", "dump_data", "dump_array_data", "create_test_log", "Logger") @@ -120,7 +122,10 @@ def dump_data(file_path: Path, data: Iterable[Iterable[Any]]) -> None: with open(file_path_ext, "w", newline="" if dict_based_data else None) as f: if dict_based_data: - writer = csv.DictWriter(f, fieldnames=data[0].keys()) + field_names = {} + for row in data: + field_names.update(row) + writer = csv.DictWriter(f, fieldnames=field_names.keys()) writer.writeheader() writer.writerows(data) else: @@ -216,7 +221,7 @@ def dump_array_data( print(*rows, sep=",", end="\n", file=f) -def create_test_log(base_dir, images=False, raw_data=False, **test_info): +def create_test_log(base_dir, images=False, raw_data=False, script=False, **test_info): """ create_test_log(base_dir, images=False, raw_data=False, **test_info) @@ -244,6 +249,10 @@ def create_test_log(base_dir, images=False, raw_data=False, **test_info): that is to be run. In addition to any keyword arguements passed an addition key-value pair, ('run_time': TIMESTAMP) will be added to the json. + Additionally if boolean arguement "script" is set then the top level python script + will be copied into the test directory. This is helpful to create backups of + scripts run for future reference. + Example: # The following code ... @@ -273,6 +282,8 @@ def create_test_log(base_dir, images=False, raw_data=False, **test_info): called "images". Defaults to False. raw_data (bool, optional): Whether or not to create a sub-directory called "raw_data". Defaults to False. + script (bool, optional): Whether or not to copy the top level python script to + the root directory. Defaults to False. Kwargs: test_info: configuration information for a test or experiment. Can include an alternate name for the directory structure with the key @@ -295,6 +306,10 @@ def create_test_log(base_dir, images=False, raw_data=False, **test_info): (test_dir / "images").mkdir() if raw_data: (test_dir / "raw_data").mkdir() + if script: + __main__ = importlib.import_module("__main__") + top_level_script = Path(__main__.__file__) + shutil.copy(top_level_script, test_dir / top_level_script.name) # dump config dictionary to file if kwargs were passed if test_info: