Aris/cflib2/loco tab

src/cfclient/ui/dialogs/anchor_position_dialog.py

@@ -25,31 +25,42 @@
 """
 Dialog box used to configure anchor positions. Used from the LPS tab.
 """
+
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 import cfclient
-from cfclient.utils.logconfigreader import FILE_REGEX_YAML
 from PySide6 import QtWidgets
 from PySide6.QtUiTools import loadUiType
-from PySide6.QtCore import QAbstractTableModel, Qt
+from PySide6.QtCore import QAbstractTableModel, QModelIndex, QObject, Qt
 from PySide6.QtGui import QBrush, QColor
 from PySide6.QtWidgets import QInputDialog, QFileDialog
 import yaml
 import os
 
-__author__ = 'Bitcraze AB'
-__all__ = ['AnchorPositionDialog']
+from cfclient.ui.pluginhelper import PluginHelper
+
+if TYPE_CHECKING:
+    from cfclient.ui.tabs.locopositioning_tab import LocoPositioningTab
+
+__author__ = "Bitcraze AB"
+__all__ = ["AnchorPositionDialog"]
 
 logger = logging.getLogger(__name__)
 
-(anchor_postiong_widget_class, connect_widget_base_class) = (
-    loadUiType(
-        cfclient.module_path + '/ui/dialogs/anchor_position_dialog.ui')
+FILE_REGEX_YAML = "Config *.yaml;;All *.*"
+
+Position = tuple[float, float, float]
+
+(anchor_postiong_widget_class, connect_widget_base_class) = loadUiType(
+    cfclient.module_path + "/ui/dialogs/anchor_position_dialog.ui"
 )
 
 
 class AnchorPositionConfigTableModel(QAbstractTableModel):
-    def __init__(self, headers, parent=None, *args):
+    def __init__(self, headers: list[str], parent: QObject | None = None) -> None:
         QAbstractTableModel.__init__(self, parent)
         self._anchor_positions = []
         self._headers = headers
@@ -58,13 +69,15 @@ def __init__(self, headers, parent=None, *args):
         self._green_brush = QBrush(QColor(200, 255, 200))
         self._red_brush = QBrush(QColor(255, 200, 200))
 
-    def rowCount(self, parent=None, *args, **kwargs):
+    def rowCount(self, parent: QModelIndex | None = None) -> int:
         return len(self._anchor_positions)
 
-    def columnCount(self, parent=None, *args, **kwargs):
+    def columnCount(self, parent: QModelIndex | None = None) -> int:
         return len(self._headers)
 
-    def data(self, index, role=None):
+    def data(
+        self, index: QModelIndex, role: int | None = None
+    ) -> int | float | str | QBrush | None:
         value = self._anchor_positions[index.row()][index.column()]
         if index.isValid():
             if index.column() == 0:
@@ -75,27 +88,34 @@ def data(self, index, role=None):
                     return value
             else:
                 if role == Qt.ItemDataRole.DisplayRole:
-                    return '%.2f' % (value)
+                    return "%.2f" % (value)
                 elif role == Qt.ItemDataRole.EditRole:
                     return value
                 elif role == Qt.ItemDataRole.BackgroundRole:
                     return self._get_background(index.row(), index.column())
 
-        return 
+        return
 
-    def setData(self, index, value, role=Qt.ItemDataRole.EditRole):
+    def setData(
+        self, index: QModelIndex, value: float, role: int = Qt.ItemDataRole.EditRole
+    ) -> bool:
         if not index.isValid():
             return False
 
         self._anchor_positions[index.row()][index.column()] = value
         return True
 
-    def headerData(self, col, orientation, role=None):
-        if orientation == Qt.Orientation.Horizontal and role == Qt.ItemDataRole.DisplayRole:
+    def headerData(
+        self, col: int, orientation: Qt.Orientation, role: int | None = None
+    ) -> str | None:
+        if (
+            orientation == Qt.Orientation.Horizontal
+            and role == Qt.ItemDataRole.DisplayRole
+        ):
             return self._headers[col]
-        return 
+        return
 
-    def flags(self, index):
+    def flags(self, index: QModelIndex) -> Qt.ItemFlag | None:
         if not index.isValid():
             return None
 
@@ -106,47 +126,50 @@ def flags(self, index):
         else:
             return Qt.ItemFlag.ItemIsEnabled | Qt.ItemFlag.ItemIsEditable
 
-    def add_anchor(self, anchor_id, x=0.0, y=0.0, z=0.0):
+    def add_anchor(
+        self, anchor_id: int, x: float = 0.0, y: float = 0.0, z: float = 0.0
+    ) -> None:
         if not self._id_exist(anchor_id):
             self.layoutAboutToBeChanged.emit()
             self._anchor_positions.append([0, anchor_id, x, y, z])
             self._anchor_positions.sort(key=lambda row: row[1])
             self.layoutChanged.emit()
 
-    def replace_anchors_from_latest_known_positions(self):
+    def replace_anchors_from_latest_known_positions(self) -> None:
         self.replace_anchor_positions(self._latest_known_anchor_positions)
 
-    def replace_anchor_positions(self, anchor_positions):
+    def replace_anchor_positions(self, anchor_positions: dict[int, Position]) -> None:
         self.layoutAboutToBeChanged.emit()
         self._anchor_positions = []
         for id, position in anchor_positions.items():
             self.add_anchor(id, x=position[0], y=position[1], z=position[2])
         self.layoutChanged.emit()
 
-    def get_anchor_postions(self):
+    def get_anchor_postions(self) -> dict[int, Position]:
         result = {}
         for row in self._anchor_positions:
             result[row[1]] = (row[2], row[3], row[4])
         return result
 
-    def anchor_postions_updated(self, anchor_positions):
+    def anchor_postions_updated(self, anchor_positions: dict[int, Position]) -> None:
         self.layoutAboutToBeChanged.emit()
         self._latest_known_anchor_positions = anchor_positions
         self.layoutChanged.emit()
 
-    def remove_selected_anchors(self):
+    def remove_selected_anchors(self) -> None:
         self.layoutAboutToBeChanged.emit()
-        self._anchor_positions = list(filter(
-            lambda row: row[0] == 0, self._anchor_positions))
+        self._anchor_positions = list(
+            filter(lambda row: row[0] == 0, self._anchor_positions)
+        )
         self.layoutChanged.emit()
 
-    def _id_exist(self, anchor_id):
+    def _id_exist(self, anchor_id: int) -> bool:
         for anchor in self._anchor_positions:
             if anchor[1] == anchor_id:
                 return True
         return False
 
-    def _get_background(self, row, col):
+    def _get_background(self, row: int, col: int) -> QBrush | None:
         id = self._anchor_positions[row][1]
         if id in self._latest_known_anchor_positions:
             current_value = self._anchor_positions[row][col]
@@ -157,87 +180,99 @@ def _get_background(self, row, col):
             else:
                 return self._red_brush
 
-        return 
+        return
 
 
 class AnchorPositionDialog(QtWidgets.QWidget, anchor_postiong_widget_class):
-
-    def __init__(self, lps_tab, helper):
+    def __init__(self, lps_tab: LocoPositioningTab, helper: PluginHelper) -> None:
         super(AnchorPositionDialog, self).__init__()
         self.setupUi(self)
 
         self._lps_tab = lps_tab
         self._helper = helper
 
-        self._headers = ['', 'id', 'x', 'y', 'z']
+        self._headers = ["", "id", "x", "y", "z"]
         self._data_model = AnchorPositionConfigTableModel(self._headers, self)
         self._table_view.setModel(self._data_model)
 
         self._table_view.verticalHeader().setVisible(False)
 
         header = self._table_view.horizontalHeader()
-        header.setSectionResizeMode(0, QtWidgets.QHeaderView.ResizeMode.ResizeToContents)
-        header.setSectionResizeMode(1, QtWidgets.QHeaderView.ResizeMode.ResizeToContents)
+        header.setSectionResizeMode(
+            0, QtWidgets.QHeaderView.ResizeMode.ResizeToContents
+        )
+        header.setSectionResizeMode(
+            1, QtWidgets.QHeaderView.ResizeMode.ResizeToContents
+        )
         header.setSectionResizeMode(2, QtWidgets.QHeaderView.ResizeMode.Stretch)
         header.setSectionResizeMode(3, QtWidgets.QHeaderView.ResizeMode.Stretch)
         header.setSectionResizeMode(4, QtWidgets.QHeaderView.ResizeMode.Stretch)
 
-        self._add_anchor_button.clicked.connect(
-            self._add_anchor_button_clicked)
+        self._add_anchor_button.clicked.connect(self._add_anchor_button_clicked)
         self._remove_anchors_button.clicked.connect(
-            self._data_model.remove_selected_anchors)
+            self._data_model.remove_selected_anchors
+        )
         self._get_from_anchors_button.clicked.connect(
-            self._get_from_anchors_button_clicked)
+            self._get_from_anchors_button_clicked
+        )
         self._write_to_anchors_button.clicked.connect(
-            self._write_to_anchors_button_clicked)
+            self._write_to_anchors_button_clicked
+        )
         self._close_button.clicked.connect(self.close)
-        self._load_button.clicked.connect(
-            self._load_button_clicked)
-        self._save_button.clicked.connect(
-            self._save_button_clicked)
+        self._load_button.clicked.connect(self._load_button_clicked)
+        self._save_button.clicked.connect(self._save_button_clicked)
 
-    def _add_anchor_button_clicked(self):
+    def _add_anchor_button_clicked(self) -> None:
         anchor_id, ok = QInputDialog.getInt(
-            self, "New anchor", "Enter id", min=0, max=255)
+            self, "New anchor", "Enter id", min=0, max=255
+        )
         if ok:
             self._data_model.add_anchor(anchor_id)
 
-    def _get_from_anchors_button_clicked(self):
+    def _get_from_anchors_button_clicked(self) -> None:
         self._data_model.replace_anchors_from_latest_known_positions()
 
-    def _write_to_anchors_button_clicked(self):
+    def _write_to_anchors_button_clicked(self) -> None:
         anchor_positions = self._data_model.get_anchor_postions()
         self._lps_tab.write_positions_to_anchors(anchor_positions)
 
-    def anchor_postions_updated(self, anchor_positions):
+    def anchor_postions_updated(self, anchor_positions: dict[int, Position]) -> None:
         self._data_model.anchor_postions_updated(anchor_positions)
 
-    def _load_button_clicked(self):
-        names = QFileDialog.getOpenFileName(self, 'Open file', self._helper.current_folder, FILE_REGEX_YAML)
+    def set_writing(self, writing: bool) -> None:
+        """Disable the write button while positions are being written"""
+        self._write_to_anchors_button.setEnabled(not writing)
+
+    def _load_button_clicked(self) -> None:
+        names = QFileDialog.getOpenFileName(
+            self, "Open file", self._helper.current_folder, FILE_REGEX_YAML
+        )
 
-        if names[0] == '':
+        if names[0] == "":
             return
 
         self._helper.current_folder = os.path.dirname(names[0])
 
-        f = open(names[0], 'r')
+        f = open(names[0], "r")
         with f:
             data = yaml.safe_load(f)
 
             anchor_positions = {}
             for id, pos in data.items():
-                anchor_positions[id] = (pos['x'], pos['y'], pos['z'])
+                anchor_positions[id] = (pos["x"], pos["y"], pos["z"])
             self._data_model.replace_anchor_positions(anchor_positions)
 
-    def _save_button_clicked(self):
+    def _save_button_clicked(self) -> None:
         anchor_positions = self._data_model.get_anchor_postions()
         data = {}
         for id, pos in anchor_positions.items():
-            data[id] = {'x': pos[0], 'y': pos[1], 'z': pos[2]}
+            data[id] = {"x": pos[0], "y": pos[1], "z": pos[2]}
 
-        names = QFileDialog.getSaveFileName(self, 'Save file', self._helper.current_folder, FILE_REGEX_YAML)
+        names = QFileDialog.getSaveFileName(
+            self, "Save file", self._helper.current_folder, FILE_REGEX_YAML
+        )
 
-        if names[0] == '':
+        if names[0] == "":
             return
 
         self._helper.current_folder = os.path.dirname(names[0])
@@ -247,6 +282,6 @@ def _save_button_clicked(self):
         else:
             filename = names[0]
 
-        f = open(filename, 'w')
+        f = open(filename, "w")
         with f:
             yaml.dump(data, f)

src/cfclient/ui/tabs/__init__.py

@@ -44,7 +44,8 @@
 from .ParamTab import ParamTab
 
 # from .PlotTab import PlotTab
-# from .locopositioning_tab import LocoPositioningTab
+from .locopositioning_tab import LocoPositioningTab
+
 # from .LogClientTab import LogClientTab
 # from .lighthouse_tab import LighthouseTab
 # from .TuningTab import TuningTab
@@ -64,7 +65,7 @@
     # LogTab,
     ParamTab,
     # PlotTab,
-    # LocoPositioningTab,
+    LocoPositioningTab,
     # LighthouseTab,
     # LogClientTab,
     # TuningTab,