Annotate zproject/jinja2/backends.py

zproject/jinja2/backends.py

@@ -1,6 +1,8 @@
 from __future__ import absolute_import
 
 import sys
+from typing import Any, Optional, Union
+from six import text_type
 
 import jinja2
 from django.utils import six
@@ -9,6 +11,7 @@ from django.template.backends import jinja2 as django_jinja2
 from django.template import TemplateDoesNotExist, TemplateSyntaxError, Context
 from django.utils.module_loading import import_string
 from django.template.backends.utils import csrf_input_lazy, csrf_token_lazy
+from django.http import HttpRequest
 
 
 class Jinja2(django_jinja2.Jinja2):
@@ -20,6 +23,7 @@ class Jinja2(django_jinja2.Jinja2):
     the `Template` object.
     """
     def __init__(self, params, *args, **kwargs):
+        # type: (Dict[str, Any], *Any, **Any) -> None
         # We need to remove `context_processors` from `OPTIONS` because
         # `Environment` doesn't expect it
         self.context_processors = params['OPTIONS'].pop('context_processors', [])
@@ -27,6 +31,7 @@ class Jinja2(django_jinja2.Jinja2):
         super(Jinja2, self).__init__(params, *args, **kwargs)
 
     def get_template(self, template_name):
+        # type: (str) -> Template
         try:
             return Template(self.env.get_template(template_name),
                             self.context_processors,
@@ -47,16 +52,24 @@ class Template(django_jinja2.Template):
         function.
         """
         def __init__(self, template, context_processors, debug, *args, **kwargs):
+            # type: (str, List[str], bool, *Any, **Any) -> None
             self.context_processors = context_processors
             self.debug = debug
             super(Template, self).__init__(template, *args, **kwargs)
 
         def render(self, context=None, request=None):
+            # type: (Optional[Union[Dict[str, Any], Context]], Optional[HttpRequest]) -> text_type
             if context is None:
                 context = {}
 
             if isinstance(context, Context):
-                context = context.flatten()  # Jinja2 expects a dictionary
+                # Jinja2 expects a dictionary
+                # This condition makes sure that `flatten` is called only when
+                # `context` is an instance of `Context`.
+                #
+                # Note: If we don't ignore then mypy complains about missing
+                # `flatten` attribute in some members of union.
+                context = context.flatten()  # type: ignore
 
             if request is not None:
                 context['request'] = request