run: register all classes using Run.register_class()

Don't hardcode cr2's classes in the Run class, register them using the
same interface as external users use.  Use register_dynamic for those
classes that only need a unique_word.

Change-Id: I9eff3fd7544c6fbaf7e285f8dab888b7b9a40673
Signed-off-by: Javi Merino <javi.merino@arm.com>

6 files changed, 35 insertions, 48 deletions

diff --git a/cr2/__init__.py b/cr2/__init__.py
index e7edceb..a474620 100644
--- a/cr2/__init__.py
+++ b/cr2/__init__.py
@@ -19,6 +19,15 @@ from results import CR2, get_results, combine_results
 from plotter.LinePlot import LinePlot
 from dynamic import register_dynamic, register_class
 
+# Load all the modules to make sure all classes are registered with Run
+import os
+for fname in os.listdir(os.path.dirname(__file__)):
+    import_name, extension = os.path.splitext(fname)
+    if (extension == ".py") and (fname != "__init__.py"):
+        __import__("cr2.{}".format(import_name))
+
+del fname, import_name, extension
+
 def summary_plots(actor_order, map_label, **kwords):
     """A summary of plots for a given run
 
diff --git a/cr2/pid_controller.py b/cr2/pid_controller.py
index 2a7eaab..874f897 100644
--- a/cr2/pid_controller.py
+++ b/cr2/pid_controller.py
@@ -17,6 +17,7 @@
 current directory's trace.dat"""
 
 from base import Base
+from cr2.run import Run
 from plot_utils import normalize_title, pre_plot_setup, post_plot_setup
 
 class PIDController(Base):
@@ -38,3 +39,5 @@ class PIDController(Base):
 
         self.data_frame[["output", "p", "i", "d"]].plot(ax=ax)
         post_plot_setup(ax, title=title)
+
+Run.register_class(PIDController, "thermal")
diff --git a/cr2/power.py b/cr2/power.py
index d924bc3..7afca19 100644
--- a/cr2/power.py
+++ b/cr2/power.py
@@ -20,6 +20,7 @@ from matplotlib import pyplot as plt
 import pandas as pd
 
 from base import Base
+from cr2.run import Run
 from plot_utils import normalize_title, pre_plot_setup, post_plot_setup
 
 def pivot_with_labels(dfr, data_col_name, new_col_name, mapping_label):
@@ -99,6 +100,8 @@ class OutPower(Base):
 
         return pivot_with_labels(dfr, "freq", "cpus", mapping_label) / 1000
 
+Run.register_class(OutPower, "thermal")
+
 class InPower(Base):
     """Process the cpufreq cooling power actor data in a ftrace dump"""
 
@@ -135,3 +138,5 @@ class InPower(Base):
         dfr = self.data_frame
 
         return pivot_with_labels(dfr, "freq", "cpus", mapping_label) / 1000
+
+Run.register_class(InPower, "thermal")
diff --git a/cr2/run.py b/cr2/run.py
index 7b10f1d..985a79e 100644
--- a/cr2/run.py
+++ b/cr2/run.py
@@ -18,10 +18,6 @@ import os
 import re
 import pandas as pd
 
-from thermal import Thermal, ThermalGovernor
-from pid_controller import PIDController
-from power import InPower, OutPower
-from sched import *
 import plot_utils
 
 def _plot_freq_hists(power_inst, map_label, what, axis, title):
@@ -67,22 +63,9 @@ classes are parsed.
 
     """
 
-    thermal_classes = {
-                "thermal": "Thermal",
-                "thermal_governor": "ThermalGovernor",
-                "pid_controller": "PIDController",
-                "in_power": "InPower",
-                "out_power": "OutPower",
-    }
-
-    sched_classes = {
-                "sched_load_avg_sched_group": "SchedLoadAvgSchedGroup",
-                "sched_load_avg_task": "SchedLoadAvgTask",
-                "sched_load_avg_cpu": "SchedLoadAvgCpu",
-                "sched_contrib_scale_factor": "SchedContribScaleFactor",
-                "sched_cpu_capacity": "SchedCpuCapacity",
-                "sched_cpu_frequency": "SchedCpuFrequency",
-    }
+    thermal_classes = {}
+
+    sched_classes = {}
 
     dynamic_classes = {}
 
diff --git a/cr2/sched.py b/cr2/sched.py
index 965a47b..8d7078c 100644
--- a/cr2/sched.py
+++ b/cr2/sched.py
@@ -15,6 +15,8 @@
 #
 
 from base import Base
+from cr2.dynamic import register_dynamic
+from cr2.run import Run
 
 class SchedLoadAvgSchedGroup(Base):
     """Corresponds to Linux kernel trace event sched_load_avg_sched_group"""
@@ -34,6 +36,8 @@ class SchedLoadAvgSchedGroup(Base):
         if self._cpu_mask_column in self.data_frame.columns:
             self.data_frame[self._cpu_mask_column] = self.data_frame[self._cpu_mask_column].apply('{:0>8}'.format)
 
+Run.register_class(SchedLoadAvgSchedGroup, "sched")
+
 class SchedLoadAvgTask(Base):
     """Corresponds to Linux kernel trace event sched_load_avg_task"""
     unique_word="sched_load_avg_task:"
@@ -51,35 +55,11 @@ class SchedLoadAvgTask(Base):
 
         return df[df['comm'].str.contains(key)].values.tolist()
 
-class SchedLoadAvgCpu(Base):
-    """Corresponds to Linux kernel trace event sched_load_avg_cpu"""
-    unique_word="sched_load_avg_cpu:"
-    name="sched_load_avg_cpu"
-
-    def __init__(self):
-        super(SchedLoadAvgCpu, self).__init__(
-            unique_word=self.unique_word,
-        )
+Run.register_class(SchedLoadAvgTask, "sched")
 
-class SchedContribScaleFactor(Base):
-    """Corresponds to Linux kernel trace event sched_contrib_scale_factor"""
-    unique_word="sched_contrib_scale_f:"
-    name="sched_contrib_scale_factor"
-
-    def __init__(self):
-        super(SchedContribScaleFactor, self).__init__(
-            unique_word=self.unique_word,
-        )
-
-class SchedCpuCapacity(Base):
-    """Corresponds to Linux kernel trace event sched_cpu_capacity"""
-    unique_word="sched_cpu_capacity:"
-    name="sched_cpu_capacity"
-
-    def __init__(self):
-        super(SchedCpuCapacity, self).__init__(
-            unique_word=self.unique_word,
-        )
+register_dynamic("SchedLoadAvgCpu", "sched_load_avg_cpu:", "sched")
+register_dynamic("SchedContribScaleFactor", "sched_contrib_scale_f:", "sched")
+register_dynamic("SchedCpuCapacity", "sched_cpu_capacity:", "sched")
 
 class SchedCpuFrequency(Base):
     """Corresponds to Linux kernel trace event power/cpu_frequency"""
@@ -97,3 +77,5 @@ class SchedCpuFrequency(Base):
         classes
         """
         self.data_frame.rename(columns={'cpu_id':'cpu'}, inplace=True)
+
+Run.register_class(SchedCpuFrequency, "sched")
diff --git a/cr2/thermal.py b/cr2/thermal.py
index 5932172..03939f6 100644
--- a/cr2/thermal.py
+++ b/cr2/thermal.py
@@ -21,6 +21,7 @@ import re
 from matplotlib import pyplot as plt
 
 from base import Base
+from cr2.run import Run
 from plot_utils import normalize_title, pre_plot_setup, post_plot_setup, plot_hist
 
 class Thermal(Base):
@@ -74,6 +75,8 @@ class Thermal(Base):
 
         plot_hist(temps, ax, title, "C", 30, "Temperature", xlim, "default")
 
+Run.register_class(Thermal, "thermal")
+
 class ThermalGovernor(Base):
     """Process the power allocator data in a ftrace dump"""
 
@@ -165,3 +168,5 @@ class ThermalGovernor(Base):
             cols = ["P" + actor + "_in", "P" + actor + "_out"]
             this_title = normalize_title(actor, title)
             dfr[cols].plot(title=this_title)
+
+Run.register_class(ThermalGovernor, "thermal")