From 693620d90d86ce88c8d4a77193e5af97714a913f Mon Sep 17 00:00:00 2001 From: Jennifer Power Date: Wed, 4 Oct 2023 07:11:03 -0400 Subject: [PATCH] refactor: reduce extra code in RulesYAMLTransformer for simplicity Signed-off-by: Jennifer Power --- trestlebot/transformers/csv_to_yaml.py | 7 ++- trestlebot/transformers/yaml_transformer.py | 63 +++++++++------------ 2 files changed, 31 insertions(+), 39 deletions(-) diff --git a/trestlebot/transformers/csv_to_yaml.py b/trestlebot/transformers/csv_to_yaml.py index 8cef9941..7c005e32 100644 --- 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 index 52351f7e..254759e3 100644 --- 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()