refactor: reduce extra code in RulesYAMLTransformer for simplicity

Signed-off-by: Jennifer Power <[email protected]>
RedHatProductSecurity · Oct 4, 2023 · 693620d · 693620d
1 parent 8cd41d5
commit 693620d
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 39 deletions.
diff --git a/trestlebot/transformers/csv_to_yaml.py b/trestlebot/transformers/csv_to_yaml.py
@@ -55,7 +55,10 @@ def write_to_yaml(self, filepath: pathlib.Path) -> None:
         try:
             with open(filepath, "w") as yaml_file:
                 yaml.dump(
-                    [self._yaml_transformer._write_yaml(rule) for rule in self._rules],
+                    [
+                        self._yaml_transformer.transform_from_rule(rule)
+                        for rule in self._rules
+                    ],
                     yaml_file,
                 )
         except Exception as e:
@@ -78,7 +81,7 @@ def write_default_trestle_rule_keys(self, filepath: pathlib.Path) -> None:
                     include_controls=[Control(id="example")],
                 ),
             )
-            yaml_blob = self._yaml_transformer._write_yaml(test_rule)
+            yaml_blob = self._yaml_transformer.transform_from_rule(test_rule)
             with open(filepath, "w") as yaml_file:
                 yaml_file.write(yaml_blob)
         except Exception as e:

diff --git a/trestlebot/transformers/yaml_transformer.py b/trestlebot/transformers/yaml_transformer.py
@@ -47,43 +47,6 @@ def __init__(self) -> None:
 
     def transform_to_rule(self, blob: str) -> TrestleRule:
         """Transform YAML data into a TrestleRule object."""
-        trestle_rule: TrestleRule = self._ingest_yaml(blob)
-        return trestle_rule
-
-    def transform_from_rule(self, trestle: TrestleRule) -> str:
-        """Transform TrestleRule object into YAML data."""
-        return self._write_yaml(trestle)
-
-    @staticmethod
-    def _write_yaml(rule: TrestleRule) -> str:
-        """Write the YAML to a string."""
-        yaml_obj = YAML()
-        yaml_obj.default_flow_style = False
-
-        rule_info: Dict[str, Any] = {
-            const.RULE_INFO_TAG: {
-                const.NAME: rule.name,
-                const.DESCRIPTION: rule.description,
-                const.PROFILE: rule.profile.dict(by_alias=True, exclude_unset=True),
-            },
-            const.COMPONENT_INFO_TAG: rule.component.dict(
-                by_alias=True, exclude_unset=True
-            ),
-        }
-
-        if rule.parameter is not None:
-            rule_info[const.RULE_INFO_TAG][const.PARAMETER] = rule.parameter.dict(
-                by_alias=True, exclude_unset=True
-            )
-
-        yaml_stream = StringIO()
-        yaml_obj.dump(rule_info, yaml_stream)
-
-        return yaml_stream.getvalue()
-
-    @staticmethod
-    def _ingest_yaml(blob: str) -> TrestleRule:
-        """Ingest the YAML blob into a TrestleData object."""
         try:
             yaml = YAML(typ="safe")
             yaml_data: Dict[str, Any] = yaml.load(blob)
@@ -119,3 +82,29 @@ def _ingest_yaml(blob: str) -> TrestleRule:
             raise RuntimeError(e)
 
         return rule_info_instance
+
+    def transform_from_rule(self, rule: TrestleRule) -> str:
+        """Transform TrestleRule object into YAML data."""
+        yaml_obj = YAML()
+        yaml_obj.default_flow_style = False
+
+        rule_info: Dict[str, Any] = {
+            const.RULE_INFO_TAG: {
+                const.NAME: rule.name,
+                const.DESCRIPTION: rule.description,
+                const.PROFILE: rule.profile.dict(by_alias=True, exclude_unset=True),
+            },
+            const.COMPONENT_INFO_TAG: rule.component.dict(
+                by_alias=True, exclude_unset=True
+            ),
+        }
+
+        if rule.parameter is not None:
+            rule_info[const.RULE_INFO_TAG][const.PARAMETER] = rule.parameter.dict(
+                by_alias=True, exclude_unset=True
+            )
+
+        yaml_stream = StringIO()
+        yaml_obj.dump(rule_info, yaml_stream)
+
+        return yaml_stream.getvalue()