From 44c6219129fa259d51986e4ca8693e4be1edd851 Mon Sep 17 00:00:00 2001
From: Mark Wolfman <canismarko@gmail.com>
Date: Tue, 28 Jan 2025 13:14:55 -0600
Subject: [PATCH] Fixed tests and linting.

---
 src/firefly/run_browser/client.py             | 10 +--
 src/firefly/run_browser/display.py            | 27 ++++--
 src/firefly/run_browser/gridplot_view.py      | 87 ++++---------------
 src/firefly/run_browser/lineplot_view.py      | 14 ++-
 src/firefly/run_browser/metadata_view.py      |  3 +-
 src/firefly/run_browser/multiplot_view.py     | 16 +---
 src/firefly/run_browser/tests/test_client.py  |  1 +
 src/firefly/run_browser/tests/test_display.py | 33 -------
 .../run_browser/tests/test_gridplot_view.py   | 35 +++-----
 .../run_browser/tests/test_lineplot_view.py   |  2 -
 .../run_browser/tests/test_metadata_view.py   |  1 -
 .../run_browser/tests/test_multiplot_view.py  | 20 ++---
 src/firefly/run_browser/widgets.py            | 49 +----------
 src/firefly/run_browser/xrf_view.py           |  1 -
 src/haven/catalog.py                          |  6 +-
 src/haven/motor_position.py                   |  3 +-
 src/haven/tests/test_save_motor_positions.py  | 45 ++++------
 17 files changed, 93 insertions(+), 260 deletions(-)

diff --git a/src/firefly/run_browser/client.py b/src/firefly/run_browser/client.py
index 34261db3..f3d6a906 100644
--- a/src/firefly/run_browser/client.py
+++ b/src/firefly/run_browser/client.py
@@ -2,7 +2,7 @@
 import datetime as dt
 import logging
 import warnings
-from collections import OrderedDict, ChainMap
+from collections import ChainMap, OrderedDict
 from functools import partial
 from typing import Mapping, Sequence
 
@@ -162,7 +162,7 @@ async def load_all_runs(self, filters: Mapping = {}):
             all_runs.append(run_data)
         return all_runs
 
-    async def hints(self, stream: str="primary") -> tuple[list, list]:
+    async def hints(self, stream: str = "primary") -> tuple[list, list]:
         """Get hints for this stream, as two lists.
 
         (*independent_hints*, *dependent_hints*)
@@ -178,10 +178,7 @@ async def hints(self, stream: str="primary") -> tuple[list, list]:
         dhints = [hint for hints in dhints for hint in hints]
         return ihints, dhints
 
-        
-
     async def signal_names(self, stream: str, *, hinted_only: bool = False):
-
         """Get a list of valid signal names (data columns) for selected runs.
 
         Parameters
@@ -258,7 +255,7 @@ async def all_signals(self, stream: str, *, hinted_only=False) -> dict:
 
     async def dataset(
         self,
-            dataset_name: str,
+        dataset_name: str,
         *,
         stream: str,
         uids: Sequence[str] | None = None,
@@ -284,7 +281,6 @@ async def dataset(
             arrays[run.uid] = arr
         return arrays
 
-
     async def signals(
         self,
         x_signal,
diff --git a/src/firefly/run_browser/display.py b/src/firefly/run_browser/display.py
index 8b3c4b1b..bea9dc3b 100644
--- a/src/firefly/run_browser/display.py
+++ b/src/firefly/run_browser/display.py
@@ -1,14 +1,13 @@
 import asyncio
 import datetime as dt
 import logging
-from collections import Counter, ChainMap
+from collections import ChainMap, Counter
 from contextlib import contextmanager
 from functools import partial, wraps
 from typing import Mapping, Optional, Sequence
 
-import qtawesome as qta
-import yaml
 import numpy as np
+import qtawesome as qta
 from ophyd import Device as ThreadedDevice
 from ophyd_async.core import Device
 from pydm import PyDMChannel
@@ -95,7 +94,9 @@ async def change_catalog(self, catalog_name: str):
         )
 
     @asyncSlot(str)
-    async def retrieve_dataset(self, dataset_name: str, callback, task_name: str) -> np.ndarray:
+    async def retrieve_dataset(
+        self, dataset_name: str, callback, task_name: str
+    ) -> np.ndarray:
         """Retrieve a dataset from disk, and provide it to the slot.
 
         Parameters
@@ -106,9 +107,11 @@ async def retrieve_dataset(self, dataset_name: str, callback, task_name: str) ->
           Will be called with the retrieved dataset.
         task_name
           For handling parallel database tasks.
-       """
+        """
         # Retrieve data from the database
-        data = await self.db_task(self.db.dataset(dataset_name, stream=self.stream), task_name)
+        data = await self.db_task(
+            self.db.dataset(dataset_name, stream=self.stream), task_name
+        )
         # Pass it back to the slot
         callback(data)
 
@@ -414,7 +417,9 @@ async def update_data_keys(self, *args):
                 self.db_task(self.db.hints(stream), "update data hints"),
             )
         independent_hints, dependent_hints = hints
-        self.data_keys_changed.emit(data_keys, set(independent_hints), set(dependent_hints))
+        self.data_keys_changed.emit(
+            data_keys, set(independent_hints), set(dependent_hints)
+        )
 
     @asyncSlot()
     @cancellable
@@ -425,8 +430,12 @@ async def update_data_frames(self):
             assert False
             log.info("Not loading data frames for empty stream.")
         else:
-            with self.busy_hints(run_widgets=True, run_table=False, filter_widgets=False):
-                data_frames = await self.db_task(self.db.data_frames(stream), "update data frames")
+            with self.busy_hints(
+                run_widgets=True, run_table=False, filter_widgets=False
+            ):
+                data_frames = await self.db_task(
+                    self.db.data_frames(stream), "update data frames"
+                )
         self.data_frames_changed.emit(data_frames)
 
     @asyncSlot()
diff --git a/src/firefly/run_browser/gridplot_view.py b/src/firefly/run_browser/gridplot_view.py
index 9a5971dc..cdc7f97a 100644
--- a/src/firefly/run_browser/gridplot_view.py
+++ b/src/firefly/run_browser/gridplot_view.py
@@ -1,20 +1,15 @@
 import logging
-from itertools import count
 from pathlib import Path
 from typing import Mapping, Sequence
 
-from scipy.interpolate import griddata
 import numpy as np
 import pandas as pd
-import yaml
-from matplotlib.colors import TABLEAU_COLORS
-from pandas.api.types import is_numeric_dtype
-from pyqtgraph import GraphicsLayoutWidget, ImageView, PlotItem, PlotWidget
 import pyqtgraph
-import qtawesome as qta
-from qtpy import QtCore, QtWidgets, uic
-from qtpy.QtCore import Qt, Signal, Slot
-from qtpy.QtWidgets import QFileDialog, QWidget
+from matplotlib.colors import TABLEAU_COLORS
+from pyqtgraph import ImageView, PlotItem
+from qtpy import QtWidgets, uic
+from qtpy.QtCore import Slot
+from scipy.interpolate import griddata
 
 log = logging.getLogger(__name__)
 colors = list(TABLEAU_COLORS.values())
@@ -32,6 +27,7 @@ def __init__(self, *args, view=None, **kwargs):
 
 class GridplotView(QtWidgets.QWidget):
     """Handles the plotting of tabular data that was taken on a grid."""
+
     ui_file = Path(__file__).parent / "gridplot_view.ui"
     shape = ()
     extent = ()
@@ -46,7 +42,7 @@ def __init__(self, parent=None):
         self.ui = uic.loadUi(self.ui_file, self)
         # Prepare plotting style
         vbox = self.ui.plot_widget.ui.roiPlot.getPlotItem().getViewBox()
-        vbox.setBackgroundColor('k')
+        vbox.setBackgroundColor("k")
         # Connect internal signals/slots
         self.ui.use_hints_checkbox.stateChanged.connect(self.update_signal_widgets)
         self.ui.regrid_checkbox.stateChanged.connect(self.update_signal_widgets)
@@ -67,8 +63,8 @@ def set_image_dimensions(self, metadata: Sequence):
             return
         md = list(metadata.values())[0]
         try:
-            self.shape = md['start']['shape']
-            self.extent = md['start']['extents']
+            self.shape = md["start"]["shape"]
+            self.extent = md["start"]["extents"]
         except KeyError as exc:
             self.shape = ()
             self.extent = ()
@@ -116,7 +112,9 @@ def update_signal_widgets(
             self.ui.value_signal_combobox,
             self.ui.r_signal_combobox,
         ]
-        for combobox, new_cols in zip(comboboxes, [new_xcols, new_xcols, new_ycols, new_ycols]):
+        for combobox, new_cols in zip(
+            comboboxes, [new_xcols, new_xcols, new_ycols, new_ycols]
+        ):
             old_cols = [combobox.itemText(idx) for idx in range(combobox.count())]
             if old_cols != new_cols:
                 old_value = combobox.currentText()
@@ -173,13 +171,11 @@ def prepare_plotting_data(self, df: pd.DataFrame) -> tuple[np.ndarray, np.ndarra
         return img
 
     def regrid(self, points: np.ndarray, values: np.ndarray):
-        """Calculate a new image with a shape based on metadata.
-
-        """
+        """Calculate a new image with a shape based on metadata."""
         # Prepare new regular grid to interpolate to
         (ymin, ymax), (xmin, xmax) = self.extent
         ystep, xstep = (npts * 1j for npts in self.shape)
-        yy, xx = np.mgrid[ymin:ymax:ystep,xmin:xmax:xstep]
+        yy, xx = np.mgrid[ymin:ymax:ystep, xmin:xmax:xstep]
         xi = np.c_[yy.flatten(), xx.flatten()]
         # Interpolate
         new_values = griddata(points, values, xi, method="cubic")
@@ -218,7 +214,9 @@ def plot(self, dataframes: Mapping | None = None):
         try:
             ylabel, xlabel = self.independent_hints
         except ValueError:
-            log.warning(f"Could not determine grid labels from hints: {self.independent_hints}")
+            log.warning(
+                f"Could not determine grid labels from hints: {self.independent_hints}"
+            )
         else:
             view = self.ui.plot_widget.view
             view.setLabels(left=ylabel, bottom=xlabel)
@@ -233,54 +231,3 @@ def plot(self, dataframes: Mapping | None = None):
     def clear_plot(self):
         self.ui.plot_widget.getImageItem().clear()
         self.data_items = {}
-
-
-# class Browser2DPlotWidget(ImageView):
-#     """A plot widget for 2D maps."""
-
-#     def __init__(self, *args, view=None, **kwargs):
-#         if view is None:
-#             view = PlotItem()
-#         super().__init__(*args, view=view, **kwargs)
-
-#     def plot_runs(
-#         self, runs: Mapping, xlabel: str = "", ylabel: str = "", extents=None
-#     ):
-#         """Take loaded 2D or 3D mapping data and plot it.
-
-#         Parameters
-#         ==========
-#         runs
-#           Dictionary with pandas series for each curve. The keys
-#           should be the curve labels, the series' indexes are the x
-#           values and the series' values are the y data.
-#         xlabel
-#           The label for the horizontal axis.
-#         ylabel
-#           The label for the vertical axis.
-#         extents
-#           Spatial extents for the map as ((-y, +y), (-x, +x)).
-
-#         """
-#         images = np.asarray(list(runs.values()))
-#         # Combine the different runs into one image
-#         # To-do: make this respond to the combobox selection
-#         image = np.mean(images, axis=0)
-#         # To-do: Apply transformations
-
-#         # # Plot the image
-#         if 2 <= image.ndim <= 3:
-#             self.setImage(image.T, autoRange=False)
-#         else:
-#             log.info(f"Could not plot image of dataset with shape {image.shape}.")
-#             return
-#         # Determine the axes labels
-#         self.view.setLabel(axis="bottom", text=xlabel)
-#         self.view.setLabel(axis="left", text=ylabel)
-#         # Set axes extent
-#         yextent, xextent = extents
-#         x = xextent[0]
-#         y = yextent[0]
-#         w = xextent[1] - xextent[0]
-#         h = yextent[1] - yextent[0]
-#         self.getImageItem().setRect(x, y, w, h)
diff --git a/src/firefly/run_browser/lineplot_view.py b/src/firefly/run_browser/lineplot_view.py
index e07a6ed2..9f7c7e7d 100644
--- a/src/firefly/run_browser/lineplot_view.py
+++ b/src/firefly/run_browser/lineplot_view.py
@@ -1,18 +1,14 @@
 import logging
-from itertools import count
 from pathlib import Path
 from typing import Mapping, Sequence
 
 import numpy as np
 import pandas as pd
-import yaml
-from matplotlib.colors import TABLEAU_COLORS
-from pandas.api.types import is_numeric_dtype
-from pyqtgraph import GraphicsLayoutWidget, ImageView, PlotItem, PlotWidget
 import qtawesome as qta
-from qtpy import QtCore, QtWidgets, uic
-from qtpy.QtCore import Qt, Signal, Slot
-from qtpy.QtWidgets import QFileDialog, QWidget
+from matplotlib.colors import TABLEAU_COLORS
+from pyqtgraph import PlotItem, PlotWidget
+from qtpy import QtWidgets, uic
+from qtpy.QtCore import Signal, Slot
 
 log = logging.getLogger(__name__)
 colors = list(TABLEAU_COLORS.values())
@@ -148,7 +144,7 @@ def axis_labels(self):
         use_reference = self.ui.r_signal_checkbox.checkState()
         inverted = self.ui.invert_checkbox.checkState()
         logarithm = self.ui.logarithm_checkbox.checkState()
-        gradient =  self.ui.gradient_checkbox.checkState()
+        gradient = self.ui.gradient_checkbox.checkState()
         if use_reference and inverted:
             ylabel = f"{rlabel}/{ylabel}"
         elif use_reference:
diff --git a/src/firefly/run_browser/metadata_view.py b/src/firefly/run_browser/metadata_view.py
index 61e59b65..8cf48692 100644
--- a/src/firefly/run_browser/metadata_view.py
+++ b/src/firefly/run_browser/metadata_view.py
@@ -2,8 +2,7 @@
 from typing import Mapping
 
 import yaml
-from qtpy import QtCore, QtWidgets, uic
-from qtpy.QtCore import Signal
+from qtpy import QtWidgets, uic
 
 
 class MetadataView(QtWidgets.QWidget):
diff --git a/src/firefly/run_browser/multiplot_view.py b/src/firefly/run_browser/multiplot_view.py
index 2dd7b175..fe0463cc 100644
--- a/src/firefly/run_browser/multiplot_view.py
+++ b/src/firefly/run_browser/multiplot_view.py
@@ -1,22 +1,15 @@
 import logging
 from itertools import count
-from typing import Mapping, Sequence
 from pathlib import Path
+from typing import Mapping, Sequence
 
-import yaml
-from qtpy import QtCore, QtWidgets, uic
-from qtpy.QtCore import Signal, Slot
-import numpy as np
-from matplotlib.colors import TABLEAU_COLORS
 from pandas.api.types import is_numeric_dtype
-from pyqtgraph import GraphicsLayoutWidget, ImageView, PlotItem, PlotWidget
-from qtpy.QtCore import Qt, Signal
-from qtpy.QtWidgets import QFileDialog, QWidget
-
-
+from qtpy import QtWidgets, uic
+from qtpy.QtCore import Slot
 
 log = logging.getLogger(__name__)
 
+
 class MultiplotView(QtWidgets.QWidget):
     _multiplot_items: Mapping
     ui_file = Path(__file__).parent / "multiplot_view.ui"
@@ -155,4 +148,3 @@ def multiplot_items(self, n_cols: int = 3):
             # view.resize(int(width), int(plot_width * row))
             view.setFixedHeight(1200)
             yield new_item
-
diff --git a/src/firefly/run_browser/tests/test_client.py b/src/firefly/run_browser/tests/test_client.py
index f18ec339..30a41a02 100644
--- a/src/firefly/run_browser/tests/test_client.py
+++ b/src/firefly/run_browser/tests/test_client.py
@@ -26,6 +26,7 @@ async def test_data_frames(worker):
     # Check the results
     assert uids[0] in data_keys.keys()
 
+
 @pytest.mark.asyncio
 async def test_hints(worker):
     uids = (await worker.catalog.client).keys()
diff --git a/src/firefly/run_browser/tests/test_display.py b/src/firefly/run_browser/tests/test_display.py
index c8283889..da9f2066 100644
--- a/src/firefly/run_browser/tests/test_display.py
+++ b/src/firefly/run_browser/tests/test_display.py
@@ -3,11 +3,9 @@
 from functools import partial
 from unittest.mock import AsyncMock, MagicMock
 
-import numpy as np
 import pytest
 import time_machine
 from ophyd.sim import instantiate_fake_device
-from pyqtgraph import ImageItem, ImageView, PlotItem, PlotWidget
 from qtpy.QtWidgets import QFileDialog
 
 from firefly.run_browser.display import RunBrowserDisplay
@@ -103,37 +101,6 @@ async def test_metadata(display, qtbot):
         await display.update_selected_runs()
 
 
-@pytest.mark.xfail
-async def test_update_2d_plot(catalog, display):
-    display.plot_2d_item.setRect = MagicMock()
-    # Load test data
-    run = await catalog["85573831-f4b4-4f64-b613-a6007bf03a8d"]
-    display.db.selected_runs = [run]
-    await display.update_1d_signals()
-    # Set the controls to describe the data we want to test
-    val_combobox = display.ui.signal_value_combobox
-    val_combobox.addItem("It_net_counts")
-    val_combobox.setCurrentText("It_net_counts")
-    display.ui.logarithm_checkbox_2d.setChecked(True)
-    display.ui.invert_checkbox_2d.setChecked(True)
-    display.ui.gradient_checkbox_2d.setChecked(True)
-    # Update the plots
-    await display.update_2d_plot()
-    # Determine what the image data should look like
-    expected_data = await run.__getitem__("It_net_counts", stream="primary")
-    expected_data = expected_data.reshape((5, 21)).T
-    # Check that the data were added
-    image = display.plot_2d_item.image
-    np.testing.assert_almost_equal(image, expected_data)
-    # Check that the axes were formatted correctly
-    axes = display.plot_2d_view.view.axes
-    xaxis = axes["bottom"]["item"]
-    yaxis = axes["left"]["item"]
-    assert xaxis.labelText == "aerotech_horiz"
-    assert yaxis.labelText == "aerotech_vert"
-    display.plot_2d_item.setRect.assert_called_with(-100, -80, 200, 160)
-
-
 def test_busy_hints_run_widgets(display):
     """Check that the display widgets get disabled during DB hits."""
     with display.busy_hints(run_widgets=True, run_table=False):
diff --git a/src/firefly/run_browser/tests/test_gridplot_view.py b/src/firefly/run_browser/tests/test_gridplot_view.py
index 88dfde83..99055a09 100644
--- a/src/firefly/run_browser/tests/test_gridplot_view.py
+++ b/src/firefly/run_browser/tests/test_gridplot_view.py
@@ -1,9 +1,7 @@
 import numpy as np
 import pandas as pd
 import pytest
-import matplotlib.pyplot as plt
 from pyqtgraph import ImageView
-from qtpy.QtWidgets import QPlainTextEdit
 
 from firefly.run_browser.gridplot_view import GridplotView
 
@@ -34,13 +32,13 @@ def view(qtbot):
 
 
 # Set up fake data
-yy, xx = np.mgrid[0:16,0:32]
+yy, xx = np.mgrid[0:16, 0:32]
 dataframe = pd.DataFrame(
     {
         "slow_motor": yy.flatten(),
         "fast_motor": xx.flatten(),
-        "I0-net_current": np.linspace(1, 100, num=16*32),
-        "It-net_current": np.linspace(101, 200, num=16*32),
+        "I0-net_current": np.linspace(1, 100, num=16 * 32),
+        "It-net_current": np.linspace(101, 200, num=16 * 32),
     }
 )
 dataframes = {"7d1daf1d-60c7-4aa7-a668-d1cd97e5335f": dataframe}
@@ -158,7 +156,9 @@ def test_hinted_signal_options(view):
         view.ui.r_signal_combobox,
     ]
     view.ui.use_hints_checkbox.setChecked(True)
-    view.update_signal_widgets(data_keys, ["fast_motor", "slow_motor"], ["I0-net_current"])
+    view.update_signal_widgets(
+        data_keys, ["fast_motor", "slow_motor"], ["I0-net_current"]
+    )
     # Make sure we don't include array datasets
     for combobox in comboboxes:
         assert (
@@ -199,19 +199,19 @@ def test_plotting_data(view):
 
 def test_regrid_data(view):
     # Prepare some simulated measurement data
-    xx, yy = np.mgrid[-3.2:3.2:0.25,-3.2:3.2:0.2]
-    xy = np.sqrt(xx**2+yy**2)
+    xx, yy = np.mgrid[-3.2:3.2:0.25, -3.2:3.2:0.2]
+    xy = np.sqrt(xx**2 + yy**2)
     data = np.cos(xy)
-    xmax = np.sqrt(2*np.pi)
+    xmax = np.sqrt(2 * np.pi)
     view.shape = (61, 61)
     view.extent = ((-xmax, xmax), (-xmax, xmax))
     points = np.c_[yy.flatten(), xx.flatten()]
     new_data = view.regrid(points=points, values=data.flatten())
-    assert new_data.shape == (61*61, )
+    assert new_data.shape == (61 * 61,)
     # Simulate what the interpolated data should be
-    xstep = 2*xmax/60
-    xx, yy = np.mgrid[-xmax:xmax:xstep,-xmax:xmax:xstep]
-    new_xy = np.sqrt(xx**2+yy**2)
+    xstep = 2 * xmax / 60
+    xx, yy = np.mgrid[-xmax:xmax:xstep, -xmax:xmax:xstep]
+    new_xy = np.sqrt(xx**2 + yy**2)
     test_data = np.cos(new_xy).flatten()
     np.testing.assert_almost_equal(new_data, test_data, decimal=2)
 
@@ -224,12 +224,3 @@ def test_update_plot(view):
     plot_item = view.ui.plot_widget.getImageItem()
     print(view.ui.plot_widget.image)
     assert view.ui.plot_widget.image is not None
-
-
-# Need to implement still
-@pytest.mark.xfail
-def test_axis_labels(view):
-    xlabel, ylabel = view.axis_labels()
-    assert xlabel == "energy_energy"
-    assert ylabel == "grad(ln(I0-net_current/It-net_current))"
-
diff --git a/src/firefly/run_browser/tests/test_lineplot_view.py b/src/firefly/run_browser/tests/test_lineplot_view.py
index 5e2d6457..0284927a 100644
--- a/src/firefly/run_browser/tests/test_lineplot_view.py
+++ b/src/firefly/run_browser/tests/test_lineplot_view.py
@@ -2,7 +2,6 @@
 import pandas as pd
 import pytest
 from pyqtgraph import PlotWidget
-from qtpy.QtWidgets import QPlainTextEdit
 
 from firefly.run_browser.lineplot_view import LineplotView
 
@@ -186,4 +185,3 @@ def test_axis_labels(view):
     xlabel, ylabel = view.axis_labels()
     assert xlabel == "energy_energy"
     assert ylabel == "grad(ln(I0-net_current/It-net_current))"
-    
diff --git a/src/firefly/run_browser/tests/test_metadata_view.py b/src/firefly/run_browser/tests/test_metadata_view.py
index c168c7db..c10b8424 100644
--- a/src/firefly/run_browser/tests/test_metadata_view.py
+++ b/src/firefly/run_browser/tests/test_metadata_view.py
@@ -1,4 +1,3 @@
-import numpy as np
 import pytest
 from qtpy.QtWidgets import QPlainTextEdit
 
diff --git a/src/firefly/run_browser/tests/test_multiplot_view.py b/src/firefly/run_browser/tests/test_multiplot_view.py
index 0323e72f..3e180fc4 100644
--- a/src/firefly/run_browser/tests/test_multiplot_view.py
+++ b/src/firefly/run_browser/tests/test_multiplot_view.py
@@ -1,8 +1,6 @@
-
-import pandas as pd
 import numpy as np
+import pandas as pd
 import pytest
-from qtpy.QtWidgets import QPlainTextEdit
 from pyqtgraph import GraphicsLayoutWidget
 
 from firefly.run_browser.multiplot_view import MultiplotView
@@ -15,7 +13,7 @@ def view(qtbot):
     return mp_view
 
 
-data_keys =  {
+data_keys = {
     "sim_motor_2": {
         "dtype": "number",
     },
@@ -101,8 +99,8 @@ def test_hinted_xsignal_options(view):
     assert (
         combobox.findText("I0-net_current") == -1
     ), f"I0-net_current signal should not be in {combobox.objectName()}."
-    
-    
+
+
 def test_update_plot(view):
     view.use_hints_checkbox.setChecked(True)
     view.independent_hints = ["energy_energy"]
@@ -117,9 +115,7 @@ def test_update_plot(view):
             "I0": np.sin(np.linspace(0, 6.28, num=101)),
         }
     )
-    dataframes = {
-        "7d1daf1d-60c7-4aa7-a668-d1cd97e5335f": df
-    }
+    dataframes = {"7d1daf1d-60c7-4aa7-a668-d1cd97e5335f": df}
     # Configure signals
     view.ui.x_signal_combobox.addItem("energy_energy")
     view.ui.x_signal_combobox.setCurrentText("energy_energy")
@@ -128,7 +124,7 @@ def test_update_plot(view):
     view.plot_multiples(dataframes)
     # Check that the data were added
     assert len(view._multiplot_items) == 1
-    data_item = view._multiplot_items[(0,0)].listDataItems()[0]
+    data_item = view._multiplot_items[(0, 0)].listDataItems()[0]
     xdata, ydata = data_item.getData()
-    np.testing.assert_almost_equal(xdata, df['energy_energy'])
-    np.testing.assert_almost_equal(ydata, df['I0'])
+    np.testing.assert_almost_equal(xdata, df["energy_energy"])
+    np.testing.assert_almost_equal(ydata, df["I0"])
diff --git a/src/firefly/run_browser/widgets.py b/src/firefly/run_browser/widgets.py
index b1ed943a..40e9aef7 100644
--- a/src/firefly/run_browser/widgets.py
+++ b/src/firefly/run_browser/widgets.py
@@ -1,16 +1,10 @@
 import logging
-from itertools import count
-from typing import Mapping, Optional, Sequence
+from typing import Optional, Sequence
 
-import numpy as np
-from matplotlib.colors import TABLEAU_COLORS
-from pandas.api.types import is_numeric_dtype
-from pyqtgraph import GraphicsLayoutWidget, ImageView, PlotItem, PlotWidget
 from qtpy.QtCore import Qt, Signal
 from qtpy.QtWidgets import QFileDialog, QWidget
 
 log = logging.getLogger(__name__)
-colors = list(TABLEAU_COLORS.values())
 
 
 class FiltersWidget(QWidget):
@@ -37,44 +31,3 @@ def ask(self, mimetypes: Optional[Sequence[str]] = None):
             return self.selectedFiles()
         else:
             return None
-
-
-# class Browser2DPlotWidget(ImageView):
-#     """A plot widget for 2D maps."""
-
-#     def __init__(self, *args, view=None, **kwargs):
-#         if view is None:
-#             view = PlotItem()
-#         super().__init__(*args, view=view, **kwargs)
-
-#     def plot_runs(
-#         self, runs: Mapping, xlabel: str = "", ylabel: str = "", extents=None
-#     ):
-#         """Take loaded 2D or 3D mapping data and plot it.
-
-#         Parameters
-#         ==========
-#         runs
-#           Dictionary with pandas series for each curve. The keys
-#           should be the curve labels, the series' indexes are the x
-#           values and the series' values are the y data.
-#         xlabel
-#           The label for the horizontal axis.
-#         ylabel
-#           The label for the vertical axis.
-#         extents
-#           Spatial extents for the map as ((-y, +y), (-x, +x)).
-
-#         """
-#         images = np.asarray(list(runs.values()))
-#         # Combine the different runs into one image
-#         # To-do: make this respond to the combobox selection
-#         image = np.mean(images, axis=0)
-#         # To-do: Apply transformations
-
-#         # # Plot the image
-#         if 2 <= image.ndim <= 3:
-#             self.setImage(image.T, autoRange=False)
-#         else:
-#             log.info(f"Could not plot image of dataset with shape {image.shape}.")
-#             return
diff --git a/src/firefly/run_browser/xrf_view.py b/src/firefly/run_browser/xrf_view.py
index b03da6fd..7922b679 100644
--- a/src/firefly/run_browser/xrf_view.py
+++ b/src/firefly/run_browser/xrf_view.py
@@ -4,7 +4,6 @@
 
 import numpy as np
 from qtpy import QtCore, QtWidgets, uic
-from qtpy.QtCore import Signal
 
 axes = namedtuple("axes", ("z", "y", "x"))
 
diff --git a/src/haven/catalog.py b/src/haven/catalog.py
index b4109891..98138381 100644
--- a/src/haven/catalog.py
+++ b/src/haven/catalog.py
@@ -205,7 +205,7 @@ def loop(self):
 
     @run_in_executor
     def data_keys(self, stream: str = "primary"):
-        data_keys = self.container[stream].metadata.get('data_keys')
+        data_keys = self.container[stream].metadata.get("data_keys")
         return data_keys or {}
 
     async def hints(self, stream: str = "primary"):
@@ -228,7 +228,9 @@ async def hints(self, stream: str = "primary"):
         # Get hints for the independent (X)
         try:
             dimensions = metadata["start"]["hints"]["dimensions"]
-            independent = [sig for signals, strm in dimensions if strm == stream for sig in signals]
+            independent = [
+                sig for signals, strm in dimensions if strm == stream for sig in signals
+            ]
         except (KeyError, IndexError):
             warnings.warn("Could not get independent hints")
         # Get hints for the dependent (X)
diff --git a/src/haven/motor_position.py b/src/haven/motor_position.py
index 4236f090..555e9631 100644
--- a/src/haven/motor_position.py
+++ b/src/haven/motor_position.py
@@ -63,10 +63,11 @@ def _load(Cls, run_md, data_keys, data):
     @classmethod
     def load(Cls, run):
         """Create a new MotorPosition object from a Tiled Bluesky run."""
+        data_keys = run["primary"].metadata["data_keys"]
         return Cls._load(
             run_md=run.metadata,
             # Assumes the 0-th descriptor is for the primary stream
-            data_keys=run["primary"].metadata["descriptors"][0]["data_keys"],
+            data_keys=data_keys,
             data=run["primary/internal/events"].read(),
         )
 
diff --git a/src/haven/tests/test_save_motor_positions.py b/src/haven/tests/test_save_motor_positions.py
index b830b8a9..5872da4f 100644
--- a/src/haven/tests/test_save_motor_positions.py
+++ b/src/haven/tests/test_save_motor_positions.py
@@ -49,14 +49,10 @@
                     ),
                 },
                 metadata={
-                    "descriptors": [
-                        {
-                            "data_keys": {
-                                "motor_A": {"object_name": "motor_A"},
-                                "motor_B": {"object_name": "motor_B"},
-                            },
-                        }
-                    ],
+                    "data_keys": {
+                        "motor_A": {"object_name": "motor_A"},
+                        "motor_B": {"object_name": "motor_B"},
+                    },
                 },
             ),
         },
@@ -91,13 +87,9 @@
                     ),
                 },
                 metadata={
-                    "descriptors": [
-                        {
-                            "data_keys": {
-                                "motorC": {"object_name": "motorC"},
-                            },
-                        }
-                    ],
+                    "data_keys": {
+                        "motorC": {"object_name": "motorC"},
+                    },
                 },
             ),
         },
@@ -128,13 +120,9 @@
                     ),
                 },
                 metadata={
-                    "descriptors": [
-                        {
-                            "data_keys": {
-                                "motorC": {"object_name": "motorC"},
-                            },
-                        }
-                    ],
+                    "data_keys": {
+                        "motorC": {"object_name": "motorC"},
+                    },
                 },
             ),
         },
@@ -165,13 +153,9 @@
                     ),
                 },
                 metadata={
-                    "descriptors": [
-                        {
-                            "data_keys": {
-                                "motorC": {"object_name": "motorC"},
-                            },
-                        }
-                    ],
+                    "data_keys": {
+                        "motorC": {"object_name": "motorC"},
+                    },
                 },
             ),
         },
@@ -309,6 +293,9 @@ async def test_list_motor_positions(client, capsys):
             f"┗━motor_B: -113.25, offset: None",
         ]
     )
+    print(first_motor)
+    print("===")
+    print(expected)
     assert first_motor == expected