Small refactoring.

cal_tools/cal_tools/ana_tools.py

-import h5py
-import numpy as np
-import matplotlib.pyplot as plt
+from enum import Enum
 import datetime
-import dateutil.parser
 import glob
 
+import dateutil.parser
+import h5py
+import matplotlib.pyplot as plt
+import numpy as np
+
 
 def save_dict_to_hdf5(dic, filename):
     """
@@ -59,37 +61,36 @@ def load_data_from_hdf5(filelist):
                         # Load ndarray if they are there
                         item = f.get("/".join((cKey, mKey, tKey)))
                         if isinstance(item, h5py._hl.dataset.Dataset):
-                            val = item.value
-                            data[cKey][mKey][tKey] = val
+                            data[cKey][mKey][tKey] = item.value
                             continue
 
                         # Loop over stored data
-                        for dKay in f.get("/".join((cKey, mKey, tKey))):
+                        for dKey in f.get("/".join((cKey, mKey, tKey))):
 
                             # Loop over metadata
-                            if dKay == "mdata":
-                                mdata = {}
+                            if dKey == "mdata":
+                                mdata_d = {}
                                 for mdKey in f.get(
                                         "/".join((cKey, mKey, tKey, 'mdata'))):
-                                    mdata[mdKey] = f.get(
+                                    mdata_d[mdKey] = f.get(
                                         "/".join(
                                             (cKey, mKey, tKey, 'mdata',
                                              mdKey))).value
                                 mdata_l = data[cKey][mKey].get("mdata", [])
-                                mdata_l.append(mdata)
+                                mdata_l.append(mdata_d)
                                 data[cKey][mKey]["mdata"] = mdata_l
                                 continue
 
-                            if dKay not in data[cKey][mKey]:
-                                data[cKey][mKey][dKay] = []
+                            if dKey not in data[cKey][mKey]:
+                                data[cKey][mKey][dKey] = []
 
                             value = f.get(
-                                "/".join((cKey, mKey, tKey, dKay))).value
+                                "/".join((cKey, mKey, tKey, dKey))).value
 
-                            if dKay == "ctime":
+                            if dKey == "ctime":
                                 value = dateutil.parser.parse(value)
 
-                            data[cKey][mKey][dKay].append(value)
+                            data[cKey][mKey][dKey].append(value)
     return data
 
 
@@ -303,10 +304,21 @@ def multi_intersect(a, b):
     return from_multiset(aa & bb)
 
 
-def hm_combine(data, fname=None, type=1, **kwargs):
+class HMType(Enum):
+    INSET_1D = 1
+    INSET_AXIS = 2
+
+
+def hm_combine(data, fname=None, htype=None, **kwargs):
     """
     Plot heatmap for calibration report
 
+    Plotting is based on pyDetLib heatmapPlot with additional
+    insets depending on heatmap type.
+    HMType.INSET_1D insets additional 1D histogram for each line
+    of the heatmap
+    HMType.INSET_AXIS insets axis for memory cells
+
     :param data: 2D numpy.array of data points to plot
     :param fname: file name of output file
     :param type: type of figure
@@ -319,7 +331,7 @@ def hm_combine(data, fname=None, type=1, **kwargs):
     ax = fig.add_subplot(111)
 
     xana.heatmapPlot(data, add_panels=False, cmap='viridis',
-                     cb_pad=0.6 if type == 1 else 0.1,
+                     cb_pad=0.6 if htype == HMType.INSET_AXIS else 0.1,
                      use_axis=ax,
                      **kwargs)
     plt.setp(ax.yaxis.get_majorticklabels(), rotation=90)
@@ -338,7 +350,7 @@ def hm_combine(data, fname=None, type=1, **kwargs):
     h_frame = 1 - pad_b - pad_t
     w_frame = 1 - pad_l - pad_r
 
-    if type == 2:
+    if htype == HMType.INSET_1D:
         ax.tick_params(axis='y', which='major', pad=50)
         for y in range(data.shape[0]):
 
@@ -369,7 +381,6 @@ def hm_combine(data, fname=None, type=1, **kwargs):
                   'drawstyle': 'steps-mid',
                   'linestyle': 'dotted',
                   'linewidth': 2,
-
                   'color': 'white',
                   },
                  {'x': np.array([0, data.shape[1] - 1]),
@@ -391,7 +402,7 @@ def hm_combine(data, fname=None, type=1, **kwargs):
                                 y_label=y_lab,
                                 use_axis=ax,
                                 y_log=False)
-    if type == 1:
+    if htype == HMType.INSET_AXIS:
 
         ax2 = plt.axes([pad_l, pad_b, w_frame, h_frame / 16.], frame_on=False,
                        yticks=range(3), yticklabels=[32, 16, 0],
@@ -400,7 +411,7 @@ def hm_combine(data, fname=None, type=1, **kwargs):
         ax2.yaxis.set_label_position("right")
         ax2.set_ylabel('Memory cell ID', color='r')
         ax2.tick_params('y', colors='r', right=True, left=False,
-                        labelright=True, labelleft=False, )
+                        labelright=True, labelleft=False)
 
     if fname:
         fig.savefig(fname, bbox_inches='tight')