feature: Use unit in label (#495)

* feature: Use unit in label

* feature: Use units in labels

* feature: Introduce kwargs, and fix x axis

* fix: Update docstring

* fix: Allow partial kwargs else fall back to data

qcodes/plots/base.py

@@ -79,7 +79,13 @@ def add(self, *args, updater=None, **kwargs):
 
             kwargs: after inserting info found in args and possibly in set_arrays
                 into `x`, `y`, and optionally `z`, these are passed along to
-                self.add_to_plot
+                self.add_to_plot.
+                To use custom labels and units pass for example:
+                    plot.add(x=set, y=amplitude, 
+                             xlabel="set"
+                             xunit="V",
+                             ylabel= "Amplitude",
+                             yunit ="V")
 
         Array shapes for 2D plots:
             x:(1D-length m), y:(1D-length n), z: (2D- n*m array)
@@ -171,8 +177,10 @@ def get_label(self, data_array):
 
         """
         # TODO this should really be a static method
-        return (getattr(data_array, 'label', '') or
+        name = (getattr(data_array, 'label', '') or
                 getattr(data_array, 'name', ''))
+        unit = getattr(data_array, 'unit', '')
+        return  name, unit
 
     def expand_trace(self, args, kwargs):
         """

qcodes/plots/pyqtgraph.py

@@ -151,17 +151,23 @@ def _line_data(self, x, y):
         return [self._clean_array(arg) for arg in [x, y] if arg is not None]
 
     def _draw_image(self, subplot_object, z, x=None, y=None, cmap='hot',
+                    zlabel=None,
+                    zunit=None,
                     **kwargs):
         img = self.rpg.ImageItem()
         subplot_object.addItem(img)
 
         hist = self.rpg.HistogramLUTItem()
         hist.setImageItem(img)
         hist.axis.setPen(self.theme[0])
-        if 'zlabel' in kwargs:  # used to specify a custom zlabel
-            hist.axis.setLabel(kwargs['zlabel'])
-        else:  # otherwise extracts the label from the dataarray
-            hist.axis.setLabel(self.get_label(z))
+
+        if zunit is None:
+            _, zunit = self.get_label(z)
+        if zlabel is None:
+            zlabel, _ = self.get_label(z)
+
+        hist.axis.setLabel(zlabel, zunit)
+
         # TODO - ensure this goes next to the correct subplot?
         self.win.addItem(hist)
 
@@ -342,22 +348,42 @@ def _update_labels(self, subplot_object, config):
         """
         Updates x and y labels, by default tries to extract label from
         the DataArray objects located in the trace config. Custom labels
-        can be specified the **kwargs "xlabel" and "ylabel"
+        can be specified the **kwargs "xlabel" and "ylabel". Custom units
+        can be specified using the kwargs xunit, ylabel
         """
         for axletter, side in (('x', 'bottom'), ('y', 'left')):
             ax = subplot_object.getAxis(side)
-            # pyqtgraph doesn't seem able to get labels, only set
-            # so we'll store it in the axis object and hope the user
-            # doesn't set it separately before adding all traces
+            # danger: 🍝
+            # find if any kwarg from plot.add in the base class
+            # matches xlabel or ylabel, signaling a custom label
             if axletter+'label' in config and not ax._qcodes_label:
                 label = config[axletter+'label']
-                ax._qcodes_label = label
-                ax.setLabel(label)
+            else:
+                label = None
+
+            # find if any kwarg from plot.add in the base class
+            # matches xunit or yunit, signaling a custom unit
+            if axletter+'unit' in config and not ax._qcodes_label:
+                unit = config[axletter+'unit']
+            else:
+                unit = None
+
+            #  find ( more hope to) unit and label from
+            # the data array inside the config
             if axletter in config and not ax._qcodes_label:
-                label = self.get_label(config[axletter])
-                if label:
-                    ax._qcodes_label = label
-                    ax.setLabel(label)
+                # now if we did not have any kwark gor label or unit
+                # fallback to the data_array
+                if unit is  None:
+                    _, unit = self.get_label(config[axletter])
+                if label is None:
+                    label, _ = self.get_label(config[axletter])
+
+            # pyqtgraph doesn't seem able to get labels, only set
+            # so we'll store it in the axis object and hope the user
+            # doesn't set it separately before adding all traces
+            ax._qcodes_label = label
+            ax._qcodes_unit = unit
+            ax.setLabel(label, unit)
 
     def update_plot(self):
         for trace in self.traces:

qcodes/plots/qcmatplotlib.py

@@ -109,10 +109,32 @@ def _get_axes(self, config):
         return self.subplots[config.get('subplot', 1) - 1]
 
     def _update_labels(self, ax, config):
-        if 'x' in config and not ax.get_xlabel():
-            ax.set_xlabel(self.get_label(config['x']))
-        if 'y' in config and not ax.get_ylabel():
-            ax.set_ylabel(self.get_label(config['y']))
+        for axletter in ("x", "y"):
+            if axletter+'label' in config:
+                label = config[axletter+'label']
+            else:
+                label = None
+
+            # find if any kwarg from plot.add in the base class
+            # matches xunit or yunit, signaling a custom unit
+            if axletter+'unit' in config:
+                unit = config[axletter+'unit']
+            else:
+                unit = None
+
+            #  find ( more hope to) unit and label from
+            # the data array inside the config
+            getter = getattr(ax, "get_{}label".format(axletter))
+            if axletter in config and not getter():
+                # now if we did not have any kwark gor label or unit
+                # fallback to the data_array
+                if unit is  None:
+                    _, unit = self.get_label(config[axletter])
+                if label is None:
+                    label, _ = self.get_label(config[axletter])
+
+            axsetter = getattr(ax, "set_{}label".format(axletter))
+            axsetter("{} ({})".format(label, unit))
 
     def update_plot(self):
         """
@@ -162,7 +184,14 @@ def update_plot(self):
 
         self.fig.canvas.draw()
 
-    def _draw_plot(self, ax, y, x=None, fmt=None, subplot=1, **kwargs):
+    def _draw_plot(self, ax, y, x=None, fmt=None, subplot=1,
+                   xlabel=None,
+                   ylabel=None,
+                   zlabel=None,
+                   xunit=None,
+                   yunit=None,
+                    zunit=None,
+                   **kwargs):
         # NOTE(alexj)stripping out subplot because which subplot we're in is already
         # described by ax, and it's not a kwarg to matplotlib's ax.plot. But I
         # didn't want to strip it out of kwargs earlier because it should stay
@@ -171,7 +200,14 @@ def _draw_plot(self, ax, y, x=None, fmt=None, subplot=1, **kwargs):
         line, = ax.plot(*args, **kwargs)
         return line
 
-    def _draw_pcolormesh(self, ax, z, x=None, y=None, subplot=1, **kwargs):
+    def _draw_pcolormesh(self, ax, z, x=None, y=None, subplot=1,
+                         xlabel=None,
+                         ylabel=None,
+                         zlabel=None,
+                         xunit=None,
+                         yunit=None,
+                         zunit=None,
+                         **kwargs):
         # NOTE(alexj)stripping out subplot because which subplot we're in is already
         # described by ax, and it's not a kwarg to matplotlib's ax.plot. But I
         # didn't want to strip it out of kwargs earlier because it should stay
@@ -201,7 +237,13 @@ def _draw_pcolormesh(self, ax, z, x=None, y=None, subplot=1, **kwargs):
             # I guess we could create the colorbar no matter what,
             # and just give it a dummy mappable to start, so we could
             # put this where it belongs.
-            ax.qcodes_colorbar.set_label(self.get_label(z))
+            if zunit is None:
+                _, zunit = self.get_label(z)
+            if zlabel is None:
+                zlabel, _ = self.get_label(z)
+
+            label = "{} ({})".format(zlabel, zunit)
+            ax.qcodes_colorbar.set_label(label)
 
         return pc