added type hinting to PowerDiagram

work related to [https://bugzilla.ipr.univ-rennes.fr/show_bug.cgi?id=3790]

cocluto/PowerDiagram.py

@@ -6,11 +6,16 @@
     This application takes its input from a database, currently in the form of an sql dump, but it could easily be adapted to read directly from a mysql database
 '''
 
+from typing import Dict, List
+from pathlib import Path
 import re
 import pygraphviz  # port install py-pygraphviz
-from inventory import Inventory, MachineSpecIdNotFound
+from .inventory import Inventory, MachineSpecIdNotFound
-from SimpaDbUtil import SqlFile, SqlDatabaseReader, TableAttrNotFound
+from .SimpaDbUtil import SqlFile, SqlDatabaseReader, TableAttrNotFound
-# from Lib import SimpaDbUtil
+
+MachineName = str
+PlugName = str
+RackId = str
 
 
 def add_capacity_constraints(capacity1, capacity2):
@@ -34,7 +39,15 @@ class Machine(object):
     """
     represents a device with input and output power plugs. It could represent a power consumer such as a server (in which case, it has no output plug), or a power distrubtion unit (in which case, it has one input plug and more than one output plugs), or even something else...
     """
-    def __init__(self, name, power_config):
+    name: MachineName
+    input_plugs: Dict[PlugName, 'Plug']
+    output_plugs: Dict[PlugName, 'Plug']
+    current_capacity_constraint: float
+    power_config: 'PowerConfig'
+    power_consumption: float
+    rack_id: RackId
+
+    def __init__(self, name: MachineName, power_config: 'PowerConfig'):
         self.name = name
         self.input_plugs = {}
         self.output_plugs = {}
@@ -43,27 +56,27 @@ class Machine(object):
         self.power_consumption = 0.0
         self.rack_id = None
 
-    def get_plug(self, plug_name):
+    def get_plug(self, plug_name: PlugName) -> 'Plug':
         plugs = {'i': self.input_plugs, 'o': self.output_plugs}[plug_name[0]]
         if plug_name not in plugs:
             plugs[plug_name] = Plug(plug_name, self, self.power_config)
         return plugs[plug_name]
 
-    def get_max_amperes(self):
+    def get_max_amperes(self) -> float:
         capacity = None
         if len(self.input_plugs) > 0:
             capacity = list(self.input_plugs.values())[0].get_max_amperes()
             capacity = add_capacity_constraints(capacity, self.current_capacity_constraint)
         return capacity
 
-    def get_outgoing_connections(self):
+    def get_outgoing_connections(self) -> List['Connection']:
         outgoing_connections = []
         for conn in self.power_config.connections:
             if conn.from_plug.machine == self:
                 outgoing_connections.append(conn)
         return outgoing_connections
 
-    def get_powered_machines(self):
+    def get_powered_machines(self) -> Dict[MachineName, 'Machine']:
         powered_machines = {}
         powered_machines[self.name] = self
         for conn in self.get_outgoing_connections():
@@ -73,7 +86,7 @@ class Machine(object):
                     powered_machines[machine.name] = machine
         return powered_machines
 
-    def get_power_consumption(self, worst_case_scenario=False):
+    def get_power_consumption(self, worst_case_scenario: bool = False) -> float:
         """
         Returns the number of watts going through this 'machine' (which could just be a powerstrip)
 
@@ -99,7 +112,12 @@ class Plug(object):
     """
     represents a power plug (input or output) of a device
     """
-    def __init__(self, name, machine, power_config):
+    name: PlugName
+    machine: Machine
+    current_capacity_constraint: float
+    power_config: 'PowerConfig'
+
+    def __init__(self, name: PlugName, machine: Machine, power_config: 'PowerConfig'):
         self.name = name
         self.machine = machine
         self.current_capacity_constraint = None  # the maximum amperes in this connection
@@ -108,19 +126,19 @@ class Plug(object):
     def __str__(self):
         return self.machine.name + '.' + self.name
 
-    def get_incoming_connection(self):
+    def get_incoming_connection(self) -> 'Connection':
         return self.power_config.get_connection_to(self)
 
 #     def get_outgoing_connections(self):
 #         return self.power_config.get_connections_from(self)
 
-    def is_input_plug(self):
+    def is_input_plug(self) -> bool:
         return self.name[0] == 'i'
 
-    def set_current_capacity_constraint(self, max_amps):
+    def set_current_capacity_constraint(self, max_amps: float):
         self.current_capacity_constraint = max_amps
 
-    def get_max_amperes(self):
+    def get_max_amperes(self) -> float:
         capacity = None
 
         debug = False
@@ -149,16 +167,19 @@ class Connection(object):
     """
     a power cable connecting an input power plug to an output power plug
     """
+    from_plug: Plug
+    to_plug: Plug
+    current_capacity_constraint: float  # the maximum amperes in this connection
 
-    def __init__(self, from_plug, to_plug):
+    def __init__(self, from_plug: Plug, to_plug: Plug):
         self.from_plug = from_plug
         self.to_plug = to_plug
-        self.current_capacity_constraint = None  # the maximum amperes in this connection
+        self.current_capacity_constraint = None
 
     def __str__(self):
         return str(self.from_plug) + ' -> ' + str(self.to_plug) + ' (' + str(self.from_plug.get_max_amperes()) + ' A, ' + str(self.get_max_amperes()) + ' A)'
 
-    def get_max_amperes(self):
+    def get_max_amperes(self) -> float:
         # gLogger.debug('%s (%s A) -> %s (%s A): ' % (str(self.from_plug), str(self.from_plug.get_max_amperes()), str(self.to_plug), str(self.to_plug.current_capacity_constraint)))
         capacity = self.from_plug.get_max_amperes()
         capacity = add_capacity_constraints(capacity, self.to_plug.current_capacity_constraint)
@@ -166,7 +187,7 @@ class Connection(object):
             capacity = min(capacity, self.current_capacity_constraint)
         return capacity
 
-    def is_redundancy_cable(self):
+    def is_redundancy_cable(self) -> bool:
         to_machine = self.to_plug.machine
         my_power_provider = self.get_power_provider()
         # find the first sibling cable that has the same provider as self
@@ -185,7 +206,7 @@ class Connection(object):
             # for each provider, all cable but the 1st one are considered as redundant
             return True
 
-    def get_power_provider(self):
+    def get_power_provider(self) -> Machine:
         from_machine = self.from_plug.machine
         if from_machine.is_power_provider():
             return from_machine
@@ -194,7 +215,7 @@ class Connection(object):
         input_plug = from_machine.input_plugs[input_plug_names[0]]
         return input_plug.get_incoming_connection().get_power_provider()
 
-    def get_power_consumption(self, worst_case_scenario=False):
+    def get_power_consumption(self, worst_case_scenario: bool = False) -> float:
         """
         Returns the number of watts going through this cable
 
@@ -222,8 +243,10 @@ class PowerConfig(object):
     """
     the description of how machines are connected together (in terms of electric power)
     """
+    machines: Dict[MachineName, Machine]
+    connections: List[Connection]
 
-    def __init__(self, simpa_db_sql_file_path):
+    def __init__(self, simpa_db_sql_file_path: Path):
         self.machines = {}
         self.connections = []
 
@@ -232,7 +255,7 @@ class PowerConfig(object):
         inventory = Inventory(sql_reader)
         self._parse_from_inventory(inventory)
 
-    def _parse_from_inventory(self, inventory):
+    def _parse_from_inventory(self, inventory: Inventory):
         """
         :param Inventory inventory:
         """
@@ -240,7 +263,7 @@ class PowerConfig(object):
         rows = inventory.query("SELECT * FROM machine_to_power")
 
         for row in rows:
-            (to_plug_as_str, from_plug_as_str, powercordid) = row
+            (to_plug_as_str, from_plug_as_str, _powercordid) = row
             if to_plug_as_str != '':
                 conn = self._add_connection(from_plug_as_str, to_plug_as_str)
                 for plug in (conn.from_plug, conn.to_plug):
@@ -302,7 +325,7 @@ class PowerConfig(object):
                     try:
                         cpu_power_consumption = inventory.get_num_cpus(machine_name) * inventory.get_cpu_tdp(machine_name)
                         estimated_max_power_consumption += cpu_power_consumption
-                    except TableAttrNotFound as e:  # noqa: F841
+                    except TableAttrNotFound as _e:  # noqa: F841
                         pass  # could happen for machines that have no cpu (eg simpa-switch-cisco-2)
                 if measured_max_power_consumption is not None:
                     machine.power_consumption = num_machines * measured_max_power_consumption
@@ -312,25 +335,25 @@ class PowerConfig(object):
                 elif estimated_max_power_consumption is not None:
                     machine.power_consumption = num_machines * estimated_max_power_consumption
 
-    def get_connection_to(self, to_plug):
+    def get_connection_to(self, to_plug: Plug) -> Connection:
         for connection in self.connections:
             if connection.to_plug == to_plug:
                 return connection
         return None
 
-    def _get_machine(self, machine_name):
+    def _get_machine(self, machine_name: MachineName) -> Machine:
         if machine_name not in self.machines:
             self.machines[machine_name] = Machine(machine_name, self)
         return self.machines[machine_name]
 
-    def _get_plug(self, plug_as_str):
+    def _get_plug(self, plug_as_str: str) -> Plug:
         elements = plug_as_str.split('_')
         plug_name = elements[-1]
         machine_name = plug_as_str[0:-(len(plug_name) + 1)]
         machine = self._get_machine(machine_name)
         return machine.get_plug(plug_name)
 
-    def _add_connection(self, from_plug_as_str, to_plug_as_str):
+    def _add_connection(self, from_plug_as_str: str, to_plug_as_str: str):
         from_plug = self._get_plug(from_plug_as_str)
         to_plug = self._get_plug(to_plug_as_str)
         conn = Connection(from_plug, to_plug)
@@ -346,7 +369,7 @@ class PowerConfig(object):
 
 class CableColorer(object):
 
-    def get_cable_color(self, cable, worst_case_scenario):
+    def get_cable_color(self, cable: Connection, worst_case_scenario: bool):
         """
         :param Connection cable:
         """
@@ -355,7 +378,7 @@ class CableColorer(object):
 
 class SimpleColorer(CableColorer):
 
-    def get_cable_color(self, cable, worst_case_scenario):
+    def get_cable_color(self, cable: Connection, worst_case_scenario: bool):
         """
         :param Connection cable:
         """
@@ -377,14 +400,14 @@ class SimpleColorer(CableColorer):
 class RampColorer(CableColorer):
 
     @staticmethod
-    def hotness_to_hsv_color(hotness):
+    def hotness_to_hsv_color(hotness: float):
         """
         :param float hotness: temperature of the wire ratio (0.0 : cold -> 1.0 : hot)
         """
         clamped_hotness = max(min(hotness, 1.0), 0.0)
         return "%f, 1.0, 0.8" % ((clamped_hotness) * 0.1 + 0.23)
 
-    def get_cable_color(self, cable, worst_case_scenario):
+    def get_cable_color(self, cable: Connection, worst_case_scenario: bool):
         """
         :param Connection cable:
         """
@@ -410,7 +433,7 @@ class RampColorer(CableColorer):
         return color
 
 
-def power_config_to_svg(power_config, svg_file_path, worst_case_scenario=True):
+def power_config_to_svg(power_config: PowerConfig, svg_file_path: Path, worst_case_scenario: bool = True):
     """
     creates a svg diagram representing the input power configuration
 
@@ -506,7 +529,7 @@ def power_config_to_svg(power_config, svg_file_path, worst_case_scenario=True):
         # sub.graph_attr['rank']='same'
     # assert False
     # graph.layout(prog='twopi')
-    with open('./toto.dot', 'w') as f:
+    with open(svg_file_path.with_suffix('.dot'), 'w', encoding='utf8') as f:
         f.write(graph.string())
 
     graph.layout(prog='dot')