Convert modifier vars to use WorkloadVar class

rfbgo · Apr 19, 2024 · 8f969d0 · 8f969d0
1 parent ae63e74
commit 8f969d0
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 19 deletions.
diff --git a/lib/ramble/ramble/language/modifier_language.py b/lib/ramble/ramble/language/modifier_language.py
@@ -269,14 +269,10 @@ def _define_modifier_variable(mod):
             if mode_name not in mod.modifier_variables:
                 mod.modifier_variables[mode_name] = {}
 
-            mod.modifier_variables[mode_name][name] = {
-                'default': default,
-                'description': description,
-                'expandable': expandable
-            }
-
-            if values:
-                mod.modifier_variables[mode_name][name]['values'] = values
+            mod.modifier_variables[mode_name][name] = ramble.workload.WorkloadVariable(
+                name, default=default, description=description,
+                values=values, expandable=expandable
+            )
 
     return _define_modifier_variable
 

diff --git a/lib/ramble/ramble/modifier.py b/lib/ramble/ramble/modifier.py
@@ -127,10 +127,10 @@ def _long_print(self):
                     indent = '\t\t'
                     for var, conf in self.modifier_variables[mode_name].items():
                         out_str.append(rucolor.nested_2(f'{indent}{var}:\n'))
-                        out_str.append(f'{indent}\tDescription: {conf["description"]}\n')
-                        out_str.append(f'{indent}\tDefault: {conf["default"]}\n')
-                        if 'values' in conf:
-                            out_str.append(f'{indent}\tSuggested Values: {conf["values"]}\n')
+                        out_str.append(f'{indent}\tDescription: {conf.description}\n')
+                        out_str.append(f'{indent}\tDefault: {conf.default}\n')
+                        if conf.values:
+                            out_str.append(f'{indent}\tSuggested Values: {conf.values}\n')
 
                 if mode_name in self.variable_modifications:
                     out_str.append(rucolor.nested_1('\tVariable Modifications:\n'))
@@ -287,17 +287,16 @@ def no_expand_vars(self):
 
         if self._usage_mode in self.modifier_variables:
             for var, var_conf in self.modifier_variables[self._usage_mode].items():
-                if not var_conf['expandable']:
+                if not var_conf.expandable:
                     yield var
 
     def mode_variables(self):
         """Return a dict of variables that should be defined for the current mode"""
 
-        var_dict = {}
         if self._usage_mode in self.modifier_variables:
-            for var, conf in self.modifier_variables[self._usage_mode].items():
-                var_dict[var] = conf['default']
-        return var_dict
+            return self.modifier_variables[self._usage_mode]
+        else:
+            return {}
 
     def run_phase_hook(self, workspace, pipeline, hook_name):
         """Run a modifier hook.

diff --git a/lib/ramble/ramble/test/modifier_language.py b/lib/ramble/ramble/test/modifier_language.py
@@ -547,5 +547,5 @@ def test_modifier_variable_directive(mod_class, func_type):
 
         assert mode in mod_inst.modifier_variables
         assert test_def['name'] in mod_inst.modifier_variables[mode]
-        for attr in ['description', 'default']:
-            assert test_def[attr] == mod_inst.modifier_variables[mode][var_name][attr]
+        assert test_def['description'] == mod_inst.modifier_variables[mode][var_name].description
+        assert test_def['default'] == mod_inst.modifier_variables[mode][var_name].default