Merge d22ed5a0a4 into ba21909655

lib/ansible/_internal/_templating/_lazy_containers.py

@@ -117,15 +117,25 @@ class _AnsibleLazyTemplateMixin:
     _lazy_options: LazyOptions
 
     def __init_subclass__(cls, **kwargs) -> None:
-        tagged_type = cls.__mro__[1]
-        native_type = tagged_type.__mro__[1]
+        # Allow explicit type declaration via _native_type class attribute
+        if hasattr(cls, '_native_type'):
+            native_type = cls._native_type
+            tagged_type = None  # no tagged variant for explicit types
+        else:
+            tagged_type = cls.__mro__[1]
+            native_type = tagged_type.__mro__[1]
 
-        for check_type in (tagged_type, native_type):
+        check_types = [native_type]
+        if tagged_type is not None:
+            check_types.append(tagged_type)
+
+        for check_type in check_types:
             if conflicting_type := cls._dispatch_types.get(check_type):
                 raise TypeError(f"Lazy mixin {cls.__name__!r} type {check_type.__name__!r} conflicts with {conflicting_type.__name__!r}.")
 
         cls._dispatch_types[native_type] = cls
-        cls._dispatch_types[tagged_type] = cls
+        if tagged_type is not None:
+            cls._dispatch_types[tagged_type] = cls
         cls._container_types.add(native_type)
         cls._empty_tags_as_native = False  # never revert to the native type when no tags remain
 
@@ -224,19 +234,40 @@ class _AnsibleLazyTemplateMixin:
         return new_value
 
 
+class _AnsibleLazyTemplateGenerator(_AnsibleLazyTemplateMixin):
+    """Wrapper for generators from lazy containers that preserves templar context."""
+    _native_type = types.GeneratorType
+    __slots__ = _AnsibleLazyTemplateMixin._SLOTS + ('_source',)
+
+    def __init__(self, source):
+        self._source = source
+        _AnsibleLazyTemplateMixin.__init__(self, source)
+
+    def __iter__(self):
+        return iter(self._source.source)
+
+    @staticmethod
+    def _lazy_values(values: t.Any, lazy_options: LazyOptions) -> _LazyValueSource:
+        return _LazyValueSource(source=values, templar=TemplateContext.current().templar, lazy_options=lazy_options)
+
+
 @t.final  # consumers of lazy collections rely heavily on the concrete types being final
 class _AnsibleLazyTemplateDict(_AnsibleTaggedDict, _AnsibleLazyTemplateMixin):
     __slots__ = _AnsibleLazyTemplateMixin._SLOTS
 
     def __init__(self, contents: t.Iterable | _LazyValueSource, /, **kwargs) -> None:
+        init_contents = contents
         if isinstance(contents, _AnsibleLazyTemplateDict):
             super().__init__(dict.items(contents), **kwargs)
         elif isinstance(contents, _LazyValueSource):
             super().__init__(contents.source, **kwargs)
+        elif isinstance(contents, _AnsibleLazyTemplateGenerator):
+            super().__init__(contents._source.source, **kwargs)
+            init_contents = contents._source
         else:
             raise UnsupportedConstructionMethodError()
 
-        _AnsibleLazyTemplateMixin.__init__(self, contents)
+        _AnsibleLazyTemplateMixin.__init__(self, init_contents)
 
     def get(self, key: t.Any, default: t.Any = None) -> t.Any:
         if (value := super().get(key, _NoKeySentinel)) is _NoKeySentinel:
@@ -267,8 +298,23 @@ class _AnsibleLazyTemplateDict(_AnsibleTaggedDict, _AnsibleLazyTemplateMixin):
         return default
 
     def items(self):
-        for key, value in super().items():
-            yield key, self._proxy_or_render_lazy_value(key, value)
+        def _item_generator(items):
+            for k, v in items:
+                yield (
+                    k,
+                    _AnsibleLazyTemplateMixin._try_create(
+                            self._proxy_or_render_lazy_value(k, v),
+                            self._lazy_options
+                    )
+                )
+
+        return _AnsibleLazyTemplateGenerator(
+            _LazyValueSource(
+                source=_item_generator(dict.items(self)),
+                templar=self._templar,
+                lazy_options=self._lazy_options
+            )
+        )
 
     def values(self):
         for _key, value in self.items():