PTA: handle None in parameter value normalization

2 files changed, 98 insertions, 6 deletions

diff --git a/lib/automata.py b/lib/automata.py
index ca78723..1715841 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -346,17 +346,21 @@ class PTA:
                 if 'formula' in normalization_spec:
                     normalization_spec['formula'] = NormalizationFunction(normalization_spec['formula'])
 
+    def normalize_parameter(self, parameter_name, parameter_value):
+        if parameter_value is not None and self.parameter_normalization is not None and parameter_name in self.parameter_normalization:
+            if 'enum' in self.parameter_normalization[parameter_name] and parameter_value in self.parameter_normalization[parameter_name]['enum']:
+                return self.parameter_normalization[parameter_name]['enum'][parameter_value]
+            if 'formula' in self.parameter_normalization[parameter_name]:
+                normalization_formula = self.parameter_normalization[parameter_name]['formula']
+                return normalization_formula.eval(parameter_value)
+        return parameter_value
+
     def normalize_parameters(self, param_dict):
         if self.parameter_normalization is None:
             return param_dict.copy()
         normalized_param = param_dict.copy()
         for parameter, value in param_dict.items():
-            if parameter in self.parameter_normalization:
-                if 'enum' in self.parameter_normalization[parameter] and value in self.parameter_normalization[parameter]['enum']:
-                    normalized_param[parameter] = self.parameter_normalization[parameter]['enum'][value]
-                if 'formula' in self.parameter_normalization[parameter]:
-                    normalization_formula = self.parameter_normalization[parameter]['formula']
-                    normalized_param[parameter] = normalization_formula.eval(value)
+            normalized_param[parameter] = self.normalize_parameter(parameter, value)
         return normalized_param
 
     @classmethod
diff --git a/test/test_pta.py b/test/test_pta.py
index 9bd6362..4c727cc 100755
--- a/test/test_pta.py
+++ b/test/test_pta.py
@@ -126,6 +126,51 @@ transition:
     is_interrupt: true
 """)
 
+example_yaml_2 = yaml.safe_load("""
+codegen:
+  instance: cc1200
+
+states:
+  - IDLE
+  - TX
+
+parameters:
+  - datarate
+  - txbytes
+  - txpower
+
+parameter_normalization:
+  txbytes:
+    enum:
+      FOO: 3
+      LONGER: 6
+  txpower:
+    formula: 'param - 16'
+
+transition:
+  init:
+    src: [UNINITIALIZED, IDLE]
+    dst: IDLE
+  setTxPower:
+    src: [IDLE]
+    dst: IDLE
+    arguments:
+      - name: txpower
+        values: [10, 20, 30]
+        parameter: txpower
+  send:
+    src: [IDLE]
+    dst: TX
+    arguments:
+      - name: data
+        values: ['FOO', 'LONGER']
+        parameter: txbytes
+  txComplete:
+    src: [TX]
+    dst: IDLE
+    is_interrupt: true
+""")
+
 def dfs_tran_to_name(runs: list, with_args: bool = False, with_param: bool = False) -> list:
     if with_param:
         return list(map(lambda run: list(map(lambda x: (x[0].name, x[1], x[2]), run)), runs))
@@ -338,6 +383,49 @@ class TestPTA(unittest.TestCase):
             ]
         )
 
+    def test_normalization(self):
+        pta = PTA.from_yaml(example_yaml_2)
+        no_param = {
+            'datarate' : None,
+            'txbytes' : None,
+            'txpower' : None,
+        }
+        param_tx3 = {
+            'datarate' : None,
+            'txbytes' : 3,
+            'txpower' : None,
+        }
+        param_tx6 = {
+            'datarate' : None,
+            'txbytes' : 6,
+            'txpower' : None,
+        }
+        param_txp10 = {
+            'datarate' : None,
+            'txbytes' : None,
+            'txpower' : -6,
+        }
+        param_txp20 = {
+            'datarate' : None,
+            'txbytes' : None,
+            'txpower' : 4,
+        }
+        param_txp30 = {
+            'datarate' : None,
+            'txbytes' : None,
+            'txpower' : 14,
+        }
+        self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1, with_arguments = True, with_parameters = True), True, True)),
+            [
+                [('init', (), no_param), ('init', (), no_param)],
+                [('init', (), no_param), ('send', ('FOO',), param_tx3)],
+                [('init', (), no_param), ('send', ('LONGER',), param_tx6)],
+                [('init', (), no_param), ('setTxPower', (10,), param_txp10)],
+                [('init', (), no_param), ('setTxPower', (20,), param_txp20)],
+                [('init', (), no_param), ('setTxPower', (30,), param_txp30)],
+            ]
+        )
+
     def test_simulation(self):
         pta = PTA()
         pta.add_state('IDLE', power = 5)