Switch to ruamel to keep comments in yaml files

PyYaml is a good library, but had the unfortunate side-effect of
stripping comments from yaml files on read/write. Switching to ruamel
now let's us retain the comments, and arguably improves usage in some
ways.
There may be some issues with awxkit objects due to the way the
representers are added, but I can't say for certain that is the case
yet.

broker/config_manager.py

@@ -3,7 +3,6 @@
 import importlib
 from importlib.metadata import version
 import json
-import os
 from pathlib import Path
 import pkgutil
 import sys
@@ -12,10 +11,14 @@
 import click
 from logzero import logger
 from packaging.version import Version
-import yaml
+from ruamel.yaml import YAML, YAMLError
 
 from broker import exceptions
 
+yaml = YAML()
+yaml.default_flow_style = False
+yaml.sort_keys = False
+
 C_SEP = "."  # chunk separator
 GH_CFG = "https://raw.githubusercontent.com/SatelliteQE/broker/master/broker_settings.yaml.example"
 
@@ -25,14 +28,6 @@ def file_name_to_ver(file_name):
     return Version(file_name[1:].replace("_", "."))
 
 
-def yaml_format(data):
-    """Format the data as yaml.
-
-    Duplicating here to avoid circular imports.
-    """
-    return yaml.dump(data, default_flow_style=False, sort_keys=False)
-
-
 class ConfigManager:
     """Class to interact with Broker's configuration file.
 
@@ -45,44 +40,31 @@ class ConfigManager:
     e.g. broker config view AnsibleTower.instances.my_instance
     """
 
+    interactive_mode = sys.stdin.isatty()
     version = version("broker")
 
     def __init__(self, settings_path=None):
-        self._interactive_mode = None
         self._settings_path = settings_path
         if settings_path:
             if settings_path.exists():
-                self._cfg = yaml.safe_load(self._settings_path.read_text())
+                self._cfg = yaml.load(self._settings_path)
             else:
                 click.secho(
                     f"Broker settings file not found at {settings_path.absolute()}.", fg="red"
                 )
                 self.init_config_file()
 
-    @property
-    def interactive_mode(self):
-        """Determine if Broker is running in interactive mode."""
-        if self._interactive_mode is None:
-            self._interactive_mode = False
-            # GitHub action context
-            if "GITHUB_WORKFLOW" not in os.environ:
-                # determine if we're being ran from a CLI
-                if sys.stdin.isatty():
-                    self._interactive_mode = True
-        return self._interactive_mode
-
     def _interactive_edit(self, chunk):
         """Write the chunk data to a temporary file and open it in an editor."""
         with NamedTemporaryFile(mode="w+", suffix=".yaml") as tmp:
-            tmp.write(yaml_format(chunk))
-            tmp.seek(0)
+            yaml.dump(chunk, tmp)
             click.edit(filename=tmp.name)
             tmp.seek(0)
             new_data = tmp.read()
         # first try to load it as yaml
         try:
-            return yaml.safe_load(new_data)
-        except yaml.YAMLError:  # then try json
+            return yaml.load(new_data)
+        except YAMLError:  # then try json
             try:
                 return json.loads(new_data)
             except json.JSONDecodeError:  # finally, just return the raw data
@@ -187,9 +169,7 @@ def update(self, chunk, new_val, curr_chunk=None):
             # update the config file if it exists
             if self._settings_path.exists():
                 self.backup()
-            self._settings_path.write_text(
-                yaml.dump(self._cfg, default_flow_style=False, sort_keys=False)
-            )
+            yaml.dump(self._cfg, self._settings_path)
         else:  # we're not at the top level, so keep going down
             if C_SEP in chunk:
                 curr, chunk = chunk.split(C_SEP, 1)
@@ -237,7 +217,7 @@ def init_config_file(self, chunk=None, _from=None):
             raise exceptions.ConfigurationError(
                 f"Broker settings file not found at {self._settings_path.absolute()}."
             )
-        chunk_data = self.get(chunk, yaml.safe_load(raw_data))
+        chunk_data = self.get(chunk, yaml.load(raw_data))
         if self.interactive_mode:
             chunk_data = self._interactive_edit(chunk_data)
         self.update(chunk, chunk_data)
@@ -253,9 +233,7 @@ def migrate(self, force_version=None):
         for migration in sorted(migrations, key=lambda m: m.TO_VERSION):
             working_config = migration.run_migrations(working_config)
         self.backup()
-        self._settings_path.write_text(
-            yaml.dump(working_config, default_flow_style=False, sort_keys=False)
-        )
+        yaml.dump(working_config, self._settings_path)
         logger.info("Config migration complete.")
 
     def validate(self, chunk, providers=None):
@@ -265,7 +243,7 @@ def validate(self, chunk, providers=None):
             for item in all_settings:
                 self.validate(item, providers)
             return
-        chunk = chunk.split(".")[0] if "." in chunk else chunk
+        chunk = chunk.split(C_SEP)[0] if C_SEP in chunk else chunk
         if chunk.lower() == "base":
             return
         if chunk.lower() == "ssh":

broker/helpers.py

@@ -7,6 +7,7 @@
 from copy import deepcopy
 import getpass
 import inspect
+from io import BytesIO
 import json
 import os
 from pathlib import Path
@@ -18,13 +19,17 @@
 
 import click
 from logzero import logger
-import yaml
+from ruamel.yaml import YAML
 
 from broker import exceptions, logger as b_log, settings
 
 FilterTest = namedtuple("FilterTest", "haystack needle test")
 INVENTORY_LOCK = threading.Lock()
 
+yaml = YAML()
+yaml.default_flow_style = False
+yaml.sort_keys = False
+
 
 def clean_dict(in_dict):
     """Remove entries from a dict where value is None."""
@@ -167,15 +172,10 @@ def load_file(file, warn=True):
         if warn:
             logger.warning(f"File {file.absolute()} is invalid or does not exist.")
         return []
-    loader_args = {}
     if file.suffix == ".json":
-        loader = json
+        return json.loads(file.read_text())
     elif file.suffix in (".yaml", ".yml"):
-        loader = yaml
-        loader_args = {"Loader": yaml.FullLoader}
-    with file.open() as f:
-        data = loader.load(f, **loader_args) or []
-    return data
+        return yaml.load(file)
 
 
 def resolve_file_args(broker_args):
@@ -251,8 +251,7 @@ def update_inventory(add=None, remove=None):
             inv_data.extend(add)
 
         settings.inventory_path.touch()
-        with settings.inventory_path.open("w") as inv_file:
-            yaml.dump(inv_data, inv_file)
+        yaml.dump(inv_data, settings.inventory_path)
 
 
 def yaml_format(in_struct):
@@ -263,8 +262,10 @@ def yaml_format(in_struct):
     :return: yaml-formatted string
     """
     if isinstance(in_struct, str):
-        in_struct = yaml.load(in_struct, Loader=yaml.FullLoader)
-    return yaml.dump(in_struct, default_flow_style=False, sort_keys=False)
+        in_struct = yaml.load(in_struct)
+    output = BytesIO()  # ruamel doesn't natively allow for string output
+    yaml.dump(in_struct, output)
+    return output.getvalue().decode("utf-8")
 
 
 def flip_provider_actions(provider_actions):

broker/providers/ansible_tower.py

@@ -1,4 +1,5 @@
 """Ansible Tower provider implementation."""
+
 from functools import cache, cached_property
 import inspect
 import json
@@ -7,7 +8,6 @@
 import click
 from dynaconf import Validator
 from logzero import logger
-import yaml
 
 from broker import exceptions
 from broker.helpers import eval_filter, find_origin
@@ -21,6 +21,8 @@
 from broker import helpers
 from broker.providers import Provider
 
+helpers.yaml.register_class(awxkit.utils.PseudoNamespace)
+
 
 class JobExecutionError(exceptions.ProviderError):
     """Raised when a job execution fails."""
@@ -776,11 +778,3 @@ def release(self, name, broker_args=None):
             source_vm=name,
             **broker_args,
         )
-
-
-def awxkit_representer(dumper, data):
-    """In order to resolve awxkit objects, a custom representer is needed."""
-    return dumper.represent_dict(dict(data))
-
-
-yaml.add_representer(awxkit.utils.PseudoNamespace, awxkit_representer)

broker/session.py

@@ -8,6 +8,7 @@
 Note: You typically want to use a Host object instance to create sessions,
       not these classes directly.
 """
+
 from contextlib import contextmanager
 from pathlib import Path
 import tempfile
@@ -19,7 +20,7 @@
 from broker.settings import settings
 
 SSH_BACKENDS = ("ssh2-python", "ssh2-python312", "ansible-pylibssh", "hussh")
-SSH_BACKEND = settings.BACKEND
+SSH_BACKEND = settings.SSH.BACKEND
 
 logger.debug(f"{SSH_BACKEND=}")
 

broker/settings.py

@@ -10,9 +10,9 @@
     settings_path: The path to the settings file.
     inventory_path: The path to the inventory file.
 """
+
 import os
 from pathlib import Path
-import sys
 
 import click
 from dynaconf import Dynaconf, Validator
@@ -21,9 +21,9 @@
 from broker.config_manager import ConfigManager
 from broker.exceptions import ConfigurationError
 
-INTERACTIVE_MODE = ConfigManager().interactive_mode
+INTERACTIVE_MODE = ConfigManager.interactive_mode
 BROKER_DIRECTORY = Path.home().joinpath(".broker")
-TEST_MODE = "pytest" in sys.modules
+TEST_MODE = os.environ.get("BROKER_TEST_MODE", False)
 
 if TEST_MODE:  # when in test mode, don't use the real broker directory
     BROKER_DIRECTORY = Path("tests/data/")