Replace do0d with a wrapper arond dond

Author: jenshnielsen

src/qcodes/dataset/dond/do_0d.py

@@ -1,92 +1,71 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Literal, overload
 
 from opentelemetry import trace
 
-from qcodes import config
-from qcodes.parameters import ParameterBase
+from .do_nd import DondKWargs, dond
 
-from ..descriptions.detect_shapes import detect_shape_of_measurement
-from ..measurements import Measurement
-from ..threading import process_params_meas
-from .do_nd_utils import _handle_plotting, _register_parameters, _set_write_period
+if TYPE_CHECKING:
+    from typing_extensions import Unpack
+
+    from .do_nd_utils import (
+        AxesTupleListWithDataSet,
+        MultiAxesTupleListWithDataSet,
+        ParamMeasT,
+    )
 
 LOG = logging.getLogger(__name__)
 TRACER = trace.get_tracer(__name__)
 
-if TYPE_CHECKING:
-    from ..descriptions.versioning.rundescribertypes import Shapes
-    from ..experiment_container import Experiment
-    from .do_nd_utils import AxesTupleListWithDataSet, ParamMeasT
+
+@overload
+def do0d(
+    *param_meas: ParamMeasT, squeeze: Literal[False], **kwargs: Unpack[DondKWargs]
+) -> MultiAxesTupleListWithDataSet: ...
+
+
+@overload
+def do0d(
+    *param_meas: ParamMeasT, squeeze: Literal[True], **kwargs: Unpack[DondKWargs]
+) -> AxesTupleListWithDataSet | MultiAxesTupleListWithDataSet: ...
+
+
+@overload
+def do0d(
+    *param_meas: ParamMeasT, squeeze: bool = True, **kwargs: Unpack[DondKWargs]
+) -> AxesTupleListWithDataSet | MultiAxesTupleListWithDataSet: ...
 
 
 @TRACER.start_as_current_span("qcodes.dataset.do0d")
 def do0d(
-    *param_meas: ParamMeasT,
-    write_period: float | None = None,
-    measurement_name: str = "",
-    exp: Experiment | None = None,
-    do_plot: bool | None = None,
-    use_threads: bool | None = None,
-    log_info: str | None = None,
-) -> AxesTupleListWithDataSet:
+    *param_meas: ParamMeasT, squeeze: bool = True, **kwargs: Unpack[DondKWargs]
+) -> AxesTupleListWithDataSet | MultiAxesTupleListWithDataSet:
     """
     Perform a measurement of a single parameter. This is probably most
-    useful for an ArrayParameter that already returns an array of data points
+    useful for a ParameterWithSetpoints that already returns an array of data points.
 
     Args:
         *param_meas: Parameter(s) to measure at each step or functions that
           will be called at each step. The function should take no arguments.
           The parameters and functions are called in the order they are
           supplied.
-        write_period: The time after which the data is actually written to the
-            database.
-        measurement_name: Name of the measurement. This will be passed down to
-            the dataset produced by the measurement. If not given, a default
-            value of 'results' is used for the dataset.
-        exp: The experiment to use for this measurement.
-        do_plot: should png and pdf versions of the images be saved after the
-            run. If None the setting will be read from ``qcodesrc.json``
-        use_threads: If True measurements from each instrument will be done on
-            separate threads. If you are measuring from several instruments
-            this may give a significant speedup.
-        log_info: Message that is logged during the measurement. If None a default
-            message is used.
+        squeeze: If True, will return a tuple of QCoDeS DataSet, Matplotlib axis,
+            Matplotlib colorbar if only one group of measurements was performed
+            and a tuple of tuples of these if more than one group of measurements
+            was performed. If False, will always return a tuple where the first
+            member is a tuple of QCoDeS DataSet(s) and the second member is a tuple
+            of Matplotlib axis(es) and the third member is a tuple of Matplotlib
+            colorbar(s).
+        **kwargs: kwargs are the same as for dond and forwarded directly to dond.
 
     Returns:
         The QCoDeS dataset.
 
     """
-    if do_plot is None:
-        do_plot = cast("bool", config.dataset.dond_plot)
-    meas = Measurement(name=measurement_name, exp=exp)
-    if log_info is not None:
-        meas._extra_log_info = log_info
-    else:
-        meas._extra_log_info = "Using 'qcodes.dataset.do0d'"
-
-    measured_parameters = tuple(
-        param for param in param_meas if isinstance(param, ParameterBase)
+    return dond(
+        *param_meas,
+        squeeze=squeeze,
+        **kwargs,
     )
-
-    try:
-        shapes: Shapes | None = detect_shape_of_measurement(
-            measured_parameters,
-        )
-    except TypeError:
-        LOG.exception(
-            f"Could not detect shape of {measured_parameters} "
-            f"falling back to unknown shape."
-        )
-        shapes = None
-
-    _register_parameters(meas, param_meas, shapes=shapes)
-    _set_write_period(meas, write_period)
-
-    with meas.run() as datasaver:
-        datasaver.add_result(*process_params_meas(param_meas, use_threads=use_threads))
-        dataset = datasaver.dataset
-
-    return _handle_plotting(dataset, do_plot)

src/qcodes/dataset/dond/do_nd.py

@@ -6,7 +6,7 @@
 from collections.abc import Callable, Mapping, Sequence
 from contextlib import ExitStack
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Literal, cast, overload
+from typing import TYPE_CHECKING, Any, Literal, NotRequired, cast, overload
 
 import numpy as np
 from opentelemetry import trace
@@ -33,8 +33,6 @@
 
 from .sweeps import AbstractSweep, TogetherSweep
 
-LOG = logging.getLogger(__name__)
-
 if TYPE_CHECKING:
     from qcodes.dataset.descriptions.versioning.rundescribertypes import Shapes
     from qcodes.dataset.dond.do_nd_utils import (
@@ -46,6 +44,7 @@
     )
     from qcodes.dataset.experiment_container import Experiment
 
+LOG = logging.getLogger(__name__)
 SweepVarType = Any
 
 TRACER = trace.get_tracer(__name__)
@@ -567,6 +566,22 @@ def parameters(self) -> tuple[ParameterBase, ...]:
         return self._parameters
 
 
+class DondKWargs(TypedDict):
+    write_period: NotRequired[float | None]
+    measurement_name: NotRequired[str | Sequence[str]]
+    exp: NotRequired[Experiment | Sequence[Experiment] | None]
+    enter_actions: NotRequired[ActionsT]
+    exit_actions: NotRequired[ActionsT]
+    do_plot: NotRequired[bool | None]
+    show_progress: NotRequired[bool | None]
+    use_threads: NotRequired[bool | None]
+    additional_setpoints: NotRequired[Sequence[ParameterBase]]
+    log_info: NotRequired[str | None]
+    break_condition: NotRequired[BreakConditionT | None]
+    dataset_dependencies: NotRequired[Mapping[str, Sequence[ParamMeasT]]]
+    in_memory_cache: NotRequired[bool | None]
+
+
 @overload
 def dond(
     *params: AbstractSweep | TogetherSweep | ParamMeasT | Sequence[ParamMeasT],