ensure transform_to_native_types converts keys (#85389)

* added basic key visitor support to variable visitor (off by default)
* transform_to_native_types enables new key visit/conversion
* add test

Co-authored-by: Matt Clay <matt@mystile.com>

lib/ansible/_internal/_json/__init__.py

@@ -81,6 +81,7 @@ class AnsibleVariableVisitor:
         convert_custom_scalars: bool = False,
         convert_to_native_values: bool = False,
         apply_transforms: bool = False,
+        visit_keys: bool = False,
         encrypted_string_behavior: EncryptedStringBehavior = EncryptedStringBehavior.DECRYPT,
     ):
         super().__init__()  # supports StateTrackingMixIn
@@ -92,6 +93,7 @@ class AnsibleVariableVisitor:
         self.convert_custom_scalars = convert_custom_scalars
         self.convert_to_native_values = convert_to_native_values
         self.apply_transforms = apply_transforms
+        self.visit_keys = visit_keys
         self.encrypted_string_behavior = encrypted_string_behavior
 
         if apply_transforms:
@@ -134,6 +136,13 @@ class AnsibleVariableVisitor:
 
         return result
 
+    def _visit_key(self, key: t.Any) -> t.Any:
+        """Internal implementation to recursively visit a key if visit_keys is enabled."""
+        if not self.visit_keys:
+            return key
+
+        return self._visit(None, key)  # key=None prevents state tracking from seeing the key as value
+
     def _visit(self, key: t.Any, value: _T) -> _T:
         """Internal implementation to recursively visit a data structure's contents."""
         self._current = key  # supports StateTrackingMixIn
@@ -144,15 +153,15 @@ class AnsibleVariableVisitor:
             value = self._template_engine.transform(value)
             value_type = type(value)
 
-        # DTFIX0: need to handle native copy for keys too
         if self.convert_to_native_values and isinstance(value, _datatag.AnsibleTaggedObject):
             value = value._native_copy()
             value_type = type(value)
 
         result: _T
 
-        # DTFIX0: the visitor is ignoring dict/mapping keys except for debugging and schema-aware checking, it should be doing type checks on keys
+        # DTFIX-FUTURE: Visitor generally ignores dict/mapping keys by default except for debugging and schema-aware checking.
-        #                keep in mind the allowed types for keys is a more restrictive set than for values (str and tagged str only, not EncryptedString)
+        #               It could be checking keys destined for variable storage to apply more strict rules about key shape and type.
+
         # DTFIX0: some type lists being consulted (the ones from datatag) are probably too permissive, and perhaps should not be dynamic
 
         if (result := self._early_visit(value, value_type)) is not _sentinel:
@@ -160,7 +169,7 @@ class AnsibleVariableVisitor:
         # DTFIX7: de-duplicate and optimize; extract inline generator expressions and fallback function or mapping for native type calculation?
         elif value_type in _ANSIBLE_ALLOWED_MAPPING_VAR_TYPES:  # check mappings first, because they're also collections
             with self:  # supports StateTrackingMixIn
-                result = AnsibleTagHelper.tag_copy(value, ((k, self._visit(k, v)) for k, v in value.items()), value_type=value_type)
+                result = AnsibleTagHelper.tag_copy(value, ((self._visit_key(k), self._visit(k, v)) for k, v in value.items()), value_type=value_type)
         elif value_type in _ANSIBLE_ALLOWED_NON_SCALAR_COLLECTION_VAR_TYPES:
             with self:  # supports StateTrackingMixIn
                 result = AnsibleTagHelper.tag_copy(value, (self._visit(k, v) for k, v in enumerate(t.cast(t.Iterable, value))), value_type=value_type)
@@ -174,7 +183,7 @@ class AnsibleVariableVisitor:
                     result = str(value)  # type: ignore[assignment]
         elif self.convert_mapping_to_dict and _internal.is_intermediate_mapping(value):
             with self:  # supports StateTrackingMixIn
-                result = {k: self._visit(k, v) for k, v in value.items()}  # type: ignore[assignment]
+                result = {self._visit_key(k): self._visit(k, v) for k, v in value.items()}  # type: ignore[assignment]
         elif self.convert_sequence_to_list and _internal.is_intermediate_iterable(value):
             with self:  # supports StateTrackingMixIn
                 result = [self._visit(k, v) for k, v in enumerate(t.cast(t.Iterable, value))]  # type: ignore[assignment]

lib/ansible/_internal/_json/_profiles/_legacy.py

@@ -1,6 +1,6 @@
 """
 Backwards compatibility profile for serialization other than inventory (which should use inventory_legacy for backward-compatible trust behavior).
-Behavior is equivalent to pre 2.18 `AnsibleJSONEncoder` with vault_to_text=True.
+Behavior is equivalent to pre 2.19 `AnsibleJSONEncoder` with vault_to_text=True.
 """
 
 from __future__ import annotations as _annotations

lib/ansible/utils/vars.py

@@ -306,6 +306,7 @@ def transform_to_native_types(
         convert_custom_scalars=True,
         convert_to_native_values=True,
         apply_transforms=True,
+        visit_keys=True,  # ensure that keys are also converted
         encrypted_string_behavior=_json.EncryptedStringBehavior.REDACT if redact else _json.EncryptedStringBehavior.DECRYPT,
     )
 

test/units/utils/test_vars.py

@@ -23,8 +23,10 @@ from collections import defaultdict
 from unittest import mock
 
 import unittest
+
 from ansible.errors import AnsibleError
-from ansible.utils.vars import combine_vars, merge_hash
+from ansible._internal._datatag._tags import Origin
+from ansible.utils.vars import combine_vars, merge_hash, transform_to_native_types
 
 
 class TestVariableUtils(unittest.TestCase):
@@ -274,3 +276,15 @@ class TestVariableUtils(unittest.TestCase):
             "b": high['b'] + [1, 1, 2]
         }
         self.assertEqual(merge_hash(low, high, True, 'prepend_rp'), expected)
+
+
+def test_transform_to_native_types() -> None:
+    """Verify that transform_to_native_types results in native types for both keys and values."""
+    value = {Origin(description="blah").tag("tagged_key"): Origin(description="blah").tag("value with tagged key")}
+
+    result = transform_to_native_types(value)
+
+    assert result == dict(tagged_key="value with tagged key")
+
+    assert all(type(key) is str for key in result.keys())  # pylint: disable=unidiomatic-typecheck
+    assert all(type(value) is str for value in result.values())  # pylint: disable=unidiomatic-typecheck