Split YAMLCommentFormatter from the help formatter.

Author: mauvilsa

jsonargparse/_formatters.py

@@ -30,13 +30,14 @@
     supports_optionals_as_positionals,
 )
 from ._completions import ShtabAction
+from ._deprecated import deprecated
 from ._link_arguments import ActionLink
 from ._namespace import Namespace, NSKeyError
 from ._optionals import import_ruyaml
 from ._type_checking import ArgumentParser, ruyamlCommentedMap
 from ._typehints import ActionTypeHint, type_to_str
 
-__all__ = ["DefaultHelpFormatter"]
+__all__ = ["DefaultHelpFormatter", "YAMLCommentFormatter"]
 
 
 empty_help: str = "_EMPTY_HELP_"
@@ -54,6 +55,115 @@ class PercentTemplate(Template):
     """  # type: ignore[assignment]
 
 
+class YAMLCommentFormatter:
+    """Formatter class for adding YAML comments to configuration files."""
+
+    def __init__(self, help_formatter: HelpFormatter):
+        self.help_formatter = help_formatter
+
+    def add_yaml_comments(self, cfg: str) -> str:
+        """Adds help text as yaml comments."""
+        ruyaml = import_ruyaml("add_yaml_comments")
+        yaml = ruyaml.YAML()
+        cfg = yaml.load(cfg)
+
+        def get_subparsers(parser, prefix=""):
+            subparsers = {}
+            if parser._subparsers is not None:
+                for key, subparser in parser._subparsers._group_actions[0].choices.items():
+                    full_key = (prefix + "." if prefix else "") + key
+                    subparsers[full_key] = subparser
+                    subparsers.update(get_subparsers(subparser, prefix=full_key))
+            return subparsers
+
+        parser = parent_parser.get()
+        parsers = get_subparsers(parser)
+        parsers[None] = parser
+
+        group_titles = {}
+        for parser_key, parser in parsers.items():
+            group_titles[parser_key] = parser.description
+            prefix = "" if parser_key is None else parser_key + "."
+            for group in parser._action_groups:
+                actions = filter_default_actions(group._group_actions)
+                actions = [
+                    a for a in actions if not isinstance(a, (_ActionConfigLoad, ActionConfigFile, _ActionSubCommands))
+                ]
+                keys = {re.sub(r"\.?[^.]+$", "", a.dest) for a in actions if "." in a.dest}
+                for key in keys:
+                    group_titles[prefix + key] = group.title
+
+        def set_comments(cfg, prefix="", depth=0):
+            for key in cfg.keys():
+                full_key = (prefix + "." if prefix else "") + key
+                action = _find_action(parser, full_key)
+                text = None
+                if full_key in group_titles and isinstance(cfg[key], dict):
+                    text = group_titles[full_key]
+                elif action is not None and action.help not in {None, SUPPRESS}:
+                    text = self.help_formatter._expand_help(action)
+                if isinstance(cfg[key], dict):
+                    if text:
+                        self.set_yaml_group_comment(text, cfg, key, depth)
+                    set_comments(cfg[key], full_key, depth + 1)
+                elif text:
+                    self.set_yaml_argument_comment(text, cfg, key, depth)
+
+        if parser.description is not None:
+            self.set_yaml_start_comment(parser.description, cfg)
+        set_comments(cfg)
+        out = StringIO()
+        yaml.dump(cfg, out)
+        return out.getvalue()
+
+    def set_yaml_start_comment(
+        self,
+        text: str,
+        cfg: "ruyamlCommentedMap",
+    ):
+        """Sets the start comment to a ruyaml object.
+
+        Args:
+            text: The content to use for the comment.
+            cfg: The ruyaml object.
+        """
+        cfg.yaml_set_start_comment(text)
+
+    def set_yaml_group_comment(
+        self,
+        text: str,
+        cfg: "ruyamlCommentedMap",
+        key: str,
+        depth: int,
+    ):
+        """Sets the comment for a group to a ruyaml object.
+
+        Args:
+            text: The content to use for the comment.
+            cfg: The parent ruyaml object.
+            key: The key of the group.
+            depth: The nested level of the group.
+        """
+        cfg.yaml_set_comment_before_after_key(key, before="\n" + text, indent=2 * depth)
+
+    def set_yaml_argument_comment(
+        self,
+        text: str,
+        cfg: "ruyamlCommentedMap",
+        key: str,
+        depth: int,
+    ):
+        """Sets the comment for an argument to a ruyaml object.
+
+        Args:
+            text: The content to use for the comment.
+            cfg: The parent ruyaml object.
+            key: The key of the argument.
+            depth: The nested level of the argument.
+        """
+        cfg.yaml_set_comment_before_after_key(key, before="\n" + text, indent=2 * depth)
+
+
 class DefaultHelpFormatter(HelpFormatter):
     """Help message formatter that includes types, default values and env var names.
 
@@ -64,6 +174,69 @@ class DefaultHelpFormatter(HelpFormatter):
     the respective environment variable name is included preceded by 'ENV:'.
     """
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._yaml_formatter = YAMLCommentFormatter(self)
+
+    @deprecated(
+        """
+        The add_yaml_comments method is deprecated and will be removed in a future version.
+        Use :class:`YAMLCommentFormatter` instead.
+    """
+    )
+    def add_yaml_comments(self, cfg: str) -> str:
+        """Adds help text as yaml comments."""
+        return self._yaml_formatter.add_yaml_comments(cfg)
+
+    @deprecated(
+        """
+        The set_yaml_start_comment method is deprecated and will be removed in a future version.
+        Use :class:`YAMLCommentFormatter` instead.
+    """
+    )
+    def set_yaml_start_comment(self, text: str, cfg: "ruyamlCommentedMap"):
+        """Sets the start comment to a ruyaml object.
+
+        Args:
+            text: The content to use for the comment.
+            cfg: The ruyaml object.
+        """
+        self._yaml_formatter.set_yaml_start_comment(text, cfg)
+
+    @deprecated(
+        """
+        The set_yaml_group_comment method is deprecated and will be removed in a future version.
+        Use :class:`YAMLCommentFormatter` instead.
+    """
+    )
+    def set_yaml_group_comment(self, text: str, cfg: "ruyamlCommentedMap", key: str, depth: int):
+        """Sets the comment for a group to a ruyaml object.
+
+        Args:
+            text: The content to use for the comment.
+            cfg: The parent ruyaml object.
+            key: The key of the group.
+            depth: The nested level of the group.
+        """
+        self._yaml_formatter.set_yaml_group_comment(text, cfg, key, depth)
+
+    @deprecated(
+        """
+        The set_yaml_argument_comment method is deprecated and will be removed in a future version.
+        Use :class:`YAMLCommentFormatter` instead.
+    """
+    )
+    def set_yaml_argument_comment(self, text: str, cfg: "ruyamlCommentedMap", key: str, depth: int):
+        """Sets the comment for an argument to a ruyaml object.
+
+        Args:
+            text: The content to use for the comment.
+            cfg: The parent ruyaml object.
+            key: The key of the argument.
+            depth: The nested level of the argument.
+        """
+        self._yaml_formatter.set_yaml_argument_comment(text, cfg, key, depth)
+
     def _get_help_string(self, action: Action) -> str:
         action_help = " " if action.help == empty_help else action.help
         assert isinstance(action_help, str)
@@ -184,108 +357,6 @@ def add_usage(self, usage: Optional[str], actions: Iterable[Action], *args, **kw
         actions = [a for a in actions if not isinstance(a, ActionLink)]
         super().add_usage(usage, actions, *args, **kwargs)
 
-    def add_yaml_comments(self, cfg: str) -> str:
-        """Adds help text as yaml comments."""
-        ruyaml = import_ruyaml("add_yaml_comments")
-        yaml = ruyaml.YAML()
-        cfg = yaml.load(cfg)
-
-        def get_subparsers(parser, prefix=""):
-            subparsers = {}
-            if parser._subparsers is not None:
-                for key, subparser in parser._subparsers._group_actions[0].choices.items():
-                    full_key = (prefix + "." if prefix else "") + key
-                    subparsers[full_key] = subparser
-                    subparsers.update(get_subparsers(subparser, prefix=full_key))
-            return subparsers
-
-        parser = parent_parser.get()
-        parsers = get_subparsers(parser)
-        parsers[None] = parser
-
-        group_titles = {}
-        for parser_key, parser in parsers.items():
-            group_titles[parser_key] = parser.description
-            prefix = "" if parser_key is None else parser_key + "."
-            for group in parser._action_groups:
-                actions = filter_default_actions(group._group_actions)
-                actions = [
-                    a for a in actions if not isinstance(a, (_ActionConfigLoad, ActionConfigFile, _ActionSubCommands))
-                ]
-                keys = {re.sub(r"\.?[^.]+$", "", a.dest) for a in actions if "." in a.dest}
-                for key in keys:
-                    group_titles[prefix + key] = group.title
-
-        def set_comments(cfg, prefix="", depth=0):
-            for key in cfg.keys():
-                full_key = (prefix + "." if prefix else "") + key
-                action = _find_action(parser, full_key)
-                text = None
-                if full_key in group_titles and isinstance(cfg[key], dict):
-                    text = group_titles[full_key]
-                elif action is not None and action.help not in {None, SUPPRESS}:
-                    text = self._expand_help(action)
-                if isinstance(cfg[key], dict):
-                    if text:
-                        self.set_yaml_group_comment(text, cfg, key, depth)
-                    set_comments(cfg[key], full_key, depth + 1)
-                elif text:
-                    self.set_yaml_argument_comment(text, cfg, key, depth)
-
-        if parser.description is not None:
-            self.set_yaml_start_comment(parser.description, cfg)
-        set_comments(cfg)
-        out = StringIO()
-        yaml.dump(cfg, out)
-        return out.getvalue()
-
-    def set_yaml_start_comment(
-        self,
-        text: str,
-        cfg: ruyamlCommentedMap,
-    ):
-        """Sets the start comment to a ruyaml object.
-
-        Args:
-            text: The content to use for the comment.
-            cfg: The ruyaml object.
-        """
-        cfg.yaml_set_start_comment(text)
-
-    def set_yaml_group_comment(
-        self,
-        text: str,
-        cfg: ruyamlCommentedMap,
-        key: str,
-        depth: int,
-    ):
-        """Sets the comment for a group to a ruyaml object.
-
-        Args:
-            text: The content to use for the comment.
-            cfg: The parent ruyaml object.
-            key: The key of the group.
-            depth: The nested level of the group.
-        """
-        cfg.yaml_set_comment_before_after_key(key, before="\n" + text, indent=2 * depth)
-
-    def set_yaml_argument_comment(
-        self,
-        text: str,
-        cfg: ruyamlCommentedMap,
-        key: str,
-        depth: int,
-    ):
-        """Sets the comment for an argument to a ruyaml object.
-
-        Args:
-            text: The content to use for the comment.
-            cfg: The parent ruyaml object.
-            key: The key of the argument.
-            depth: The nested level of the argument.
-        """
-        cfg.yaml_set_comment_before_after_key(key, before="\n" + text, indent=2 * depth)
-
 
 def get_env_var(
     parser_or_formatter: Union[ArgumentParser, DefaultHelpFormatter],

jsonargparse/_loaders_dumpers.py

@@ -225,7 +225,8 @@ def yaml_dump(data):
 
 def yaml_comments_dump(data, parser):
     dump = dumpers["yaml"](data)
-    formatter = parser.formatter_class(parser.prog)
+    formatter_class = create_help_formatter_with_comments(parser.formatter_class)
+    formatter = formatter_class(parser.prog)
     return formatter.add_yaml_comments(dump)
 
 
@@ -333,3 +334,26 @@ def set_omegaconf_loader():
 
 
 set_loader("jsonnet", jsonnet_load, get_loader_exceptions("jsonnet"))
+
+
+def create_help_formatter_with_comments(formatter_class: Type["HelpFormatter"]) -> Type["HelpFormatter"]:
+    """Creates a dynamic class that combines a formatter with YAML comment functionality.
+
+    Args:
+        formatter_class: The base formatter class to extend.
+
+    Returns:
+        A new class that inherits from both the formatter and YAMLCommentFormatter.
+    """
+    from ._formatters import YAMLCommentFormatter
+
+    class DynamicHelpFormatter(formatter_class):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self._yaml_formatter = YAMLCommentFormatter(self)
+
+        def add_yaml_comments(self, cfg: str) -> str:
+            """Adds help text as yaml comments."""
+            return self._yaml_formatter.add_yaml_comments(cfg)
+
+    return DynamicHelpFormatter