- [feature] Added support for context managers,

  i.e. "% with x as e:/ % endwith" support.
  Courtesy Ben Trofatter [ticket:147]

5 files changed, 24 insertions, 2 deletions

diff --git a/CHANGES b/CHANGES
index 2d3dabe..c84fb92 100644
--- a/CHANGES
+++ b/CHANGES
@@ -1,4 +1,8 @@
 0.6.3
+- [feature] Added support for context managers,
+  i.e. "% with x as e:/ % endwith" support.
+  Courtesy Ben Trofatter [ticket:147]
+
 - [bug] Fixed some Py3K resource warnings due
   to filehandles being implicitly closed.
   [ticket:182]
diff --git a/mako/ast.py b/mako/ast.py
index 515d7b8..f2f09d6 100644
--- a/mako/ast.py
+++ b/mako/ast.py
@@ -83,6 +83,8 @@ class PythonFragment(PythonCode):
             code = "if False:pass\n" + code + "pass"
         elif keyword == 'except':
             code = "try:pass\n" + code + "pass"
+        elif keyword == 'with':
+            code = code + "pass"
         else:
             raise exceptions.CompileException(
                                 "Unsupported control keyword: '%s'" % 
diff --git a/mako/parsetree.py b/mako/parsetree.py
index 8aa95d4..52bd156 100644
--- a/mako/parsetree.py
+++ b/mako/parsetree.py
@@ -64,7 +64,7 @@ class ControlLine(Node):
         self.text = text
         self.keyword = keyword
         self.isend = isend
-        self.is_primary = keyword in ['for','if', 'while', 'try']
+        self.is_primary = keyword in ['for','if', 'while', 'try', 'with']
         if self.isend:
             self._declared_identifiers = []
             self._undeclared_identifiers = []
diff --git a/mako/pygen.py b/mako/pygen.py
index e38b383..b50e60e 100644
--- a/mako/pygen.py
+++ b/mako/pygen.py
@@ -108,7 +108,7 @@ class PythonPrinter(object):
             # keep track of what the keyword was that indented us,
             # if it is a python compound statement keyword
             # where we might have to look for an "unindent" keyword
-            match = re.match(r"^\s*(if|try|elif|while|for)", line)
+            match = re.match(r"^\s*(if|try|elif|while|for|with)", line)
             if match:
                 # its a "compound" keyword, so we will check for "unindentors"
                 indentor = match.group(1)
diff --git a/test/test_ast.py b/test/test_ast.py
index adea08a..60ad6ec 100644
--- a/test/test_ast.py
+++ b/test/test_ast.py
@@ -189,6 +189,22 @@ def x(q):
         eq_(parsed.declared_identifiers, set(['x']))
         eq_(parsed.undeclared_identifiers, set())
 
+    def test_locate_identifiers_12(self):
+        code = """
+class ContextManager(object):
+    def __enter__(self):
+        return 1
+    def __exit__(self, exc_type, exc_value, traceback):
+        pass
+
+with ContextManager() as x, ContextManager():
+    print x
+"""
+        parsed = ast.PythonCode(code, **exception_kwargs)
+        eq_(parsed.declared_identifiers, set(['ContextManager', 'x']))
+        eq_(parsed.undeclared_identifiers, set())
+        
+
     def test_no_global_imports(self):
         code = """
 from foo import *