Improve types

Author: ssbarnea

src/ansiblelint/yaml_utils.py

@@ -17,6 +17,7 @@
 from ruamel.yaml.comments import CommentedMap, CommentedSeq, Format
 from ruamel.yaml.composer import ComposerError
 from ruamel.yaml.constructor import RoundTripConstructor
+from ruamel.yaml.docinfo import Version
 from ruamel.yaml.emitter import Emitter, ScalarAnalysis
 
 # Module 'ruamel.yaml' does not explicitly export attribute 'YAML'; implicit reexport disabled
@@ -1028,15 +1029,23 @@ def version(self) -> tuple[int, int] | None:
         return None
 
     @version.setter
-    def version(self, val: tuple[int, int] | None) -> None:
+    def version(self, val: str | tuple[int, int] | list[int] | Version | None) -> None:
         """Ensure that yaml version uses our default value.
 
         The yaml Reader updates this value based on the ``%YAML`` directive in files.
         So, if a file does not include the directive, it sets this to None.
         But, None effectively resets the parsing version to YAML 1.2 (ruamel's default).
         """
         if val is not None:
-            self._yaml_version = val
+            if isinstance(val, tuple):
+                self._yaml_version = val
+            elif isinstance(val, list):
+                self._yaml_version = (val[0], val[1])
+            elif isinstance(val, Version):
+                self._yaml_version = (val.major, val.minor)
+            else:
+                msg = f"Unsupported argument {val}"
+                raise TypeError(msg)
         elif hasattr(self, "_yaml_version_default"):
             self._yaml_version = self._yaml_version_default
         # We do nothing if the object did not have a previous default version defined