Repository: audreyr/cookiecutter
Pull request #613 (2015-12-04 22:13:16)

Strict env in prompt
Complement #598 

Related to #610 


Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index 733f444..f11ff9e 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -6,18 +6,13 @@ History
 1.4.0 (2015-12-??) ??
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-The goal of this release is changing to a strict Jinja2 environment, paving the way to more awesome in the future.
-
-New Feature:
-
-* Now raises an error if Cookiecutter tries to render a template that contains an undefined variable. Makes generation more robust and secure (#586). Work done by `@hackebrot`_ (#111, #586, #592)
+The goal of this release is to solve problems with version 1.3.0.
 
 Other Changes:
 
 * Removed django-cms-plugin on account of 404 error, thanks to `@mativs`_ and `@pydanny`_ (#593)
 * Fixed docs/usage.rst, thanks to `@macrotim`_ (#604)
 * Update .gitignore to latest Python.gitignore and ignore PyCharm files, thanks to `@audreyr`_
-* Use open context manager to read context_file in generate() function, thanks to `@hackebrot`_ (#607, #608)
 
 Bug Fixes:
 
diff --git a/cookiecutter/prompt.py b/cookiecutter/prompt.py
index 0809f9e..7b931c8 100755
--- a/cookiecutter/prompt.py
+++ b/cookiecutter/prompt.py
@@ -14,7 +14,12 @@ import click
 from past.builtins import basestring
 
 from future.utils import iteritems
+
+from jinja2 import StrictUndefined
 from jinja2.environment import Environment
+from jinja2.exceptions import UndefinedError
+
+from .exceptions import UndefinedVariableInTemplate
 
 
 def read_user_variable(var_name, default_value):
@@ -86,6 +91,7 @@ def render_variable(env, raw, cookiecutter_dict):
     if not isinstance(raw, basestring):
         raw = str(raw)
     template = env.from_string(raw)
+
     rendered_template = template.render(cookiecutter=cookiecutter_dict)
     return rendered_template
 
@@ -111,24 +117,28 @@ def prompt_for_config(context, no_input=False):
     :param no_input: Prompt the user at command line for manual configuration?
     """
     cookiecutter_dict = {}
-    env = Environment()
+    env = Environment(undefined=StrictUndefined)
 
     for key, raw in iteritems(context[u'cookiecutter']):
         if key.startswith(u'_'):
             cookiecutter_dict[key] = raw
             continue
 
-        if isinstance(raw, list):
-            # We are dealing with a choice variable
-            val = prompt_choice_for_config(
-                cookiecutter_dict, env, key, raw, no_input
-            )
-        else:
-            # We are dealing with a regular variable
-            val = render_variable(env, raw, cookiecutter_dict)
-
-            if not no_input:
-                val = read_user_variable(key, val)
+        try:
+            if isinstance(raw, list):
+                # We are dealing with a choice variable
+                val = prompt_choice_for_config(
+                    cookiecutter_dict, env, key, raw, no_input
+                )
+            else:
+                # We are dealing with a regular variable
+                val = render_variable(env, raw, cookiecutter_dict)
+
+                if not no_input:
+                    val = read_user_variable(key, val)
+        except UndefinedError as err:
+            msg = "Unable to render variable '{}'".format(key)
+            raise UndefinedVariableInTemplate(msg, err, context)
 
         cookiecutter_dict[key] = val
     return cookiecutter_dict


Test patch:
diff --git a/tests/test_prompt.py b/tests/test_prompt.py
index 583da1d..42aad0b 100644
--- a/tests/test_prompt.py
+++ b/tests/test_prompt.py
@@ -14,7 +14,7 @@ import platform
 import pytest
 from past.builtins import basestring
 
-from cookiecutter import prompt
+from cookiecutter import prompt, exceptions
 from jinja2.environment import Environment
 
 
@@ -251,3 +251,33 @@ class TestPromptChoiceForConfig(object):
         )
         read_choice.assert_called_once_with('orientation', choices)
         assert expected_choice == actual_choice
+
+
+def test_undefined_variable_in_cookiecutter_dict():
+    context = {
+        'cookiecutter': {
+            'hello': 'world',
+            'foo': '{{cookiecutter.nope}}'
+        }
+    }
+    with pytest.raises(exceptions.UndefinedVariableInTemplate) as err:
+        prompt.prompt_for_config(context, no_input=True)
+
+    error = err.value
+    assert error.message == "Unable to render variable 'foo'"
+    assert error.context == context
+
+
+def test_undefined_variable_in_cookiecutter_dict_with_choices():
+    context = {
+        'cookiecutter': {
+            'hello': 'world',
+            'foo': ['123', '{{cookiecutter.nope}}', '456']
+        }
+    }
+    with pytest.raises(exceptions.UndefinedVariableInTemplate) as err:
+        prompt.prompt_for_config(context, no_input=True)
+
+    error = err.value
+    assert error.message == "Unable to render variable 'foo'"
+    assert error.context == context