ashb commented on a change in pull request #5461: [AIRFLOW-4835] Refactor 
render_template
URL: https://github.com/apache/airflow/pull/5461#discussion_r309624776
 
 

 ##########
 File path: airflow/models/baseoperator.py
 ##########
 @@ -636,51 +636,71 @@ def __setstate__(self, state):
         self.__dict__ = state
         self._log = logging.getLogger("airflow.task.operators")
 
-    def render_template_from_field(self, attr, content, context, jinja_env):
+    def render_template_fields(self, context: Dict) -> None:
         """
-        Renders a template from a field. If the field is a string, it will
-        simply render the string and return the result. If it is a collection 
or
-        nested set of collections, it will traverse the structure and render
-        all elements in it. If the field has another type, it will return it 
as it is.
+        Template all attributes listed in template_fields. Note this operation 
is irreversible.
+
+        :param context: Dict with values to apply on content
+        :type context: dict
+        """
+        if self.template_fields:
+            if self.dag.user_defined_macros:
+                context.update(self.dag.user_defined_macros)
+
+            for attr_name in self.template_fields:
+                content = getattr(self, attr_name)
+                if content:
+                    rendered_content = self.render_template(content, context)
+                    setattr(self, attr_name, rendered_content)
+
+    def render_template(
+        self, content: Any, context: Dict, jinja_env: 
Optional[jinja2.Environment] = None
+    ) -> Any:
         """
-        rt = self.render_template
+        Render a templated string. The content can be a collection holding 
multiple templated strings and will
+        be templated recursively.
+
+        :param content: Content to template. Only strings can be templated 
(may be inside collection).
+        :type content: Any
+        :param context: Dict with values to apply on templated content
+        :type context: dict
+        :param jinja_env: Jinja environment. Can be provided to avoid 
re-creating Jinja environments during
+            recursion.
+        :type jinja_env: jinja2.Environment
+        :return: Templated content
+        """
+
+        if not jinja_env:
+            jinja_env = self.get_template_env()
+
         if isinstance(content, str):
-            result = jinja_env.from_string(content).render(**context)
-        elif isinstance(content, tuple):
+            if any([content.endswith(ext) for ext in self.template_ext]):
+                # Content contains a filepath
+                return jinja_env.get_template(content).render(**context)
+            else:
+                return jinja_env.from_string(content).render(**context)
+
+        if isinstance(content, tuple):
             if type(content) is not tuple:
                 # Special case for named tuples
-                result = content.__class__(*(rt(attr, e, context) for e in 
content))
+                return content.__class__(
+                    *(self.render_template(element, context, jinja_env) for 
element in content)
+                )
             else:
-                result = tuple(rt(attr, e, context) for e in content)
+                return tuple(self.render_template(element, context, jinja_env) 
for element in content)
+
         elif isinstance(content, list):
-            result = [rt(attr, e, context) for e in content]
-        elif isinstance(content, dict):
-            result = {
-                k: rt("{}[{}]".format(attr, k), v, context)
-                for k, v in list(content.items())}
-        else:
-            result = content
-        return result
+            return [self.render_template(element, context, jinja_env) for 
element in content]
 
-    def render_template(self, attr, content, context):
-        """
-        Renders a template either from a file or directly in a field, and 
returns
-        the rendered result.
-        """
-        jinja_env = self.get_template_env()
+        elif isinstance(content, dict):
+            return {key: self.render_template(value, context, jinja_env) for 
key, value in content.items()}
 
-        exts = self.__class__.template_ext
-        if (
-                isinstance(content, str) and
-                any([content.endswith(ext) for ext in exts])):
-            return jinja_env.get_template(content).render(**context)
         else:
-            return self.render_template_from_field(attr, content, context, 
jinja_env)
+            return content
 
-    def get_template_env(self):
-        return self.dag.get_template_env() \
-            if hasattr(self, 'dag') \
-            else jinja2.Environment(cache_size=0)
+    def get_template_env(self) -> jinja2.Environment:
+        """Fetch a Jinja template environment from the DAG or instantiate 
empty environment if no DAG."""
+        return self.dag.get_template_env() if hasattr(self, 'dag') else 
jinja2.Environment(cache_size=0)
 
 Review comment:
   ```suggestion
           return self.dag.get_template_env() if self.has_dag() else 
jinja2.Environment(cache_size=0)
   ```
   Especially given the attribute is `_dag` :)
   

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


With regards,
Apache Git Services

Reply via email to