diff --git a/qualang_tools/results/__init__.py b/qualang_tools/results/__init__.py index 0b145c32..372d3461 100644 --- a/qualang_tools/results/__init__.py +++ b/qualang_tools/results/__init__.py @@ -1,7 +1,6 @@ from qualang_tools.results.results import fetching_tool from qualang_tools.results.results import progress_counter from qualang_tools.results.results import wait_until_job_is_paused - from qualang_tools.results.data_handler import DataHandler, data_processors __all__ = ["fetching_tool", "progress_counter", "wait_until_job_is_paused", "DataHandler", "data_processors"] diff --git a/qualang_tools/results/data_handler/__init__.py b/qualang_tools/results/data_handler/__init__.py index cb496969..7f10ffb7 100644 --- a/qualang_tools/results/data_handler/__init__.py +++ b/qualang_tools/results/data_handler/__init__.py @@ -1,6 +1,37 @@ -from .data_folder_tools import * from . import data_processors -from .data_processors import DEFAULT_DATA_PROCESSORS -from .data_handler import * -__all__ = [*data_folder_tools.__all__, data_processors, DEFAULT_DATA_PROCESSORS, *data_handler.__all__] +DEFAULT_DATA_PROCESSORS = [ + data_processors.MatplotlibPlotSaver, + data_processors.NumpyArraySaver, +] + +try: + import xarray # noqa: F401 + + DEFAULT_DATA_PROCESSORS.append(data_processors.XarraySaver) +except ImportError: + pass + +from .data_folder_tools import ( # noqa: E402 + DEFAULT_FOLDER_PATTERN, + extract_data_folder_properties, + get_latest_data_folder, + create_data_folder, +) +from .data_handler import save_data, DataHandler # noqa: E402 +from .data_processors import DataProcessor, MatplotlibPlotSaver, NumpyArraySaver, XarraySaver # noqa: E402 + + +__all__ = [ + "DEFAULT_FOLDER_PATTERN", + "extract_data_folder_properties", + "get_latest_data_folder", + "create_data_folder", + "DataProcessor", + "MatplotlibPlotSaver", + "NumpyArraySaver", + "XarraySaver", + "DEFAULT_DATA_PROCESSORS", + "save_data", + "DataHandler", +] diff --git a/qualang_tools/results/data_handler/data_folder_tools.py b/qualang_tools/results/data_handler/data_folder_tools.py index 725d2d21..4441306a 100644 --- a/qualang_tools/results/data_handler/data_folder_tools.py +++ b/qualang_tools/results/data_handler/data_folder_tools.py @@ -10,9 +10,6 @@ from datetime import datetime -__all__ = ["DEFAULT_FOLDER_PATTERN", "extract_data_folder_properties", "get_latest_data_folder", "create_data_folder"] - - DEFAULT_FOLDER_PATTERN = "%Y-%m-%d/#{idx}_{name}_%H%M%S" diff --git a/qualang_tools/results/data_handler/data_handler.py b/qualang_tools/results/data_handler/data_handler.py index c518a88b..661d8138 100644 --- a/qualang_tools/results/data_handler/data_handler.py +++ b/qualang_tools/results/data_handler/data_handler.py @@ -5,7 +5,7 @@ from typing import Any, Dict, Optional, Sequence, Union import warnings -from .data_processors import DEFAULT_DATA_PROCESSORS, DataProcessor +from .data_processors import DataProcessor from .data_folder_tools import ( DEFAULT_FOLDER_PATTERN, create_data_folder, @@ -14,8 +14,6 @@ ) -__all__ = ["save_data", "DataHandler"] - NODE_FILENAME = "node.json" @@ -106,7 +104,9 @@ class DataHandler: data_handler.save_data(data) """ - default_data_processors = DEFAULT_DATA_PROCESSORS + from . import DEFAULT_DATA_PROCESSORS as _DEFAULT_DATA_PROCESSORS + + default_data_processors = _DEFAULT_DATA_PROCESSORS root_data_folder: Path = None folder_pattern: str = DEFAULT_FOLDER_PATTERN data_filename: str = "data.json" diff --git a/qualang_tools/results/data_handler/data_processors.py b/qualang_tools/results/data_handler/data_processors.py index 5807c5a3..b5cbc42e 100644 --- a/qualang_tools/results/data_handler/data_processors.py +++ b/qualang_tools/results/data_handler/data_processors.py @@ -4,11 +4,6 @@ from matplotlib import pyplot as plt import numpy as np -__all__ = ["DEFAULT_DATA_PROCESSORS", "DataProcessor", "MatplotlibPlotSaver", "NumpyArraySaver", "XarraySaver"] - - -DEFAULT_DATA_PROCESSORS = [] - def iterate_nested_dict( d: Dict[str, Any], parent_keys: Optional[List[str]] = None @@ -78,9 +73,6 @@ def post_process(self, data_folder: Path): fig.savefig(data_folder / path) -DEFAULT_DATA_PROCESSORS.append(MatplotlibPlotSaver) - - class NumpyArraySaver(DataProcessor): merge_arrays: bool = True merged_array_name: str = "arrays.npz" @@ -118,9 +110,6 @@ def post_process(self, data_folder: Path): np.save(data_folder / path.with_suffix(".npy"), arr) -DEFAULT_DATA_PROCESSORS.append(NumpyArraySaver) - - class XarraySaver(DataProcessor): merge_arrays: bool = False merged_array_name: str = "xarrays" @@ -184,11 +173,3 @@ def post_process(self, data_folder: Path): else: for path, array in self.data_arrays.items(): array.to_netcdf(data_folder / path.with_suffix(self.file_suffix)) - - -try: - import xarray # noqa: F401 - - DEFAULT_DATA_PROCESSORS.append(XarraySaver) -except ImportError: - pass diff --git a/tests/data_handler/test_numpy_array_saver.py b/tests/data_handler/test_numpy_array_saver.py index c6f38910..0a888d77 100644 --- a/tests/data_handler/test_numpy_array_saver.py +++ b/tests/data_handler/test_numpy_array_saver.py @@ -1,6 +1,6 @@ import numpy as np -from qualang_tools.results.data_handler.data_processors import DEFAULT_DATA_PROCESSORS, NumpyArraySaver +from qualang_tools.results.data_handler import NumpyArraySaver def test_numpy_array_saver_process_merged():