Add some type hints

Author: paulromano

openmc_plotter/docks.py

@@ -670,7 +670,7 @@ def updateNuclides(self):
         self.model.appliedNuclides = tuple(applied_nuclides)
 
         if 'total' in applied_nuclides:
-            self.model.appliedNuclides = ['total',]
+            self.model.appliedNuclides = ('total',)
             for nuclide, nuclide_box in self.nuclide_map.items():
                 if nuclide != 'total':
                     nuclide_box.setFlags(QtCore.Qt.ItemIsUserCheckable)

openmc_plotter/plotmodel.py

@@ -1,12 +1,12 @@
+from __future__ import annotations
 from ast import literal_eval
 from collections import defaultdict
 import copy
 import hashlib
 import itertools
-import os
-from pathlib import Path
 import pickle
 import threading
+from typing import Literal, Tuple
 
 from PySide6.QtWidgets import QItemDelegate, QColorDialog, QLineEdit, QMessageBox
 from PySide6.QtCore import QAbstractTableModel, QModelIndex, Qt, QSize, QEvent
@@ -60,6 +60,8 @@
                  'Std. Dev.': 'std_dev',
                  'Rel. Error': 'rel_err'}
 
+TallyValueType = Literal['mean', 'std_dev', 'rel_err']
+
 
 def hash_file(path):
     # return the md5 hash of a file
@@ -386,7 +388,7 @@ def create_tally_image(self, view=None):
         """
         Parameters
         ----------
-        view :
+        view : PlotView
             View used to set bounds of the tally data
 
         Returns
@@ -635,7 +637,10 @@ def _create_distribcell_image(self, tally, tally_value, scores, nuclides, cellin
 
         return image_data, None, data_min, data_max
 
-    def _create_tally_mesh_image(self, tally, tally_value, scores, nuclides, view=None):
+    def _create_tally_mesh_image(
+            self, tally: openmc.Tally, tally_value: TallyValueType,
+            scores: Tuple[str], nuclides: Tuple[str], view: PlotView = None
+        ):
         # some variables used throughout
         if view is None:
             view = self.currentView