[Context] Ensure pop happens when an error is thrown

Sources/Context.swift

@@ -1,9 +1,9 @@
 /// A container for template variables.
 public class Context {
-  var dictionaries:[[String: Any]]
+  var dictionaries: [[String: Any]]
 
   /// Initialise a Context with a dictionary
-  public init(dictionary:[String: Any]) {
+  public init(dictionary: [String: Any]) {
     dictionaries = [dictionary]
   }
 
@@ -35,20 +35,19 @@ public class Context {
   }
 
   /// Push a new level into the Context
-  public func push(dictionary: [String: Any]? = nil) {
+  private func push(dictionary: [String: Any]? = nil) {
     dictionaries.append(dictionary ?? [:])
   }
 
   /// Pop the last level off of the Context
-  public func pop() -> [String: Any]? {
+  private func pop() -> [String: Any]? {
     return dictionaries.popLast()
   }
 
   /// Push a new level onto the context for the duration of the execution of the given closure
   public func push<Result>(dictionary: [String: Any]? = nil, @noescape closure: (() throws -> Result)) rethrows -> Result {
     push(dictionary)
-    let result = try closure()
-    pop()
-    return result
+    defer { pop() }
+    return try closure()
   }
 }

Sources/IfTag.swift

@@ -67,15 +67,12 @@ public class IfNode : NodeType {
       truthy = true
     }
 
-    context.push()
-    let output:String
+    return try context.push {
       if truthy {
-      output = try renderNodes(trueNodes, context)
+        return try renderNodes(trueNodes, context)
       } else {
-      output = try renderNodes(falseNodes, context)
+        return try renderNodes(falseNodes, context)
+      }
     }
-    context.pop()
-
-    return output
   }
 }

Sources/Inheritence.swift

@@ -74,10 +74,9 @@ class ExtendsNode : NodeType {
     }
 
     let blockContext = BlockContext(blocks: blocks)
-    context.push([BlockContext.contextKey: blockContext])
-    let result = try template.render(context)
-    context.pop()
-    return result
+    return try context.push([BlockContext.contextKey: blockContext]) {
+      return try template.render(context)
+    }
   }
 }
 

Tests/ContextSpec.swift

@@ -27,29 +27,24 @@ func testContext() {
     }
 
     $0.it("allows you to retrieve a value from a parent") {
-      context.push()
-
+      try context.push {
         try expect(context["name"] as? String) == "Kyle"
       }
+    }
 
     $0.it("allows you to override a parent's value") {
-      context.push()
+      try context.push {
         context["name"] = "Katie"
-
         try expect(context["name"] as? String) == "Katie"
       }
+    }
 
     $0.it("allows you to pop to restore previous state") {
-      context.push()
+      context.push {
         context["name"] = "Katie"
-      context.pop()
-
-      try expect(context["name"] as? String) == "Kyle"
       }
 
-    $0.it("allows you to push a dictionary onto the stack") {
-      context.push(["name": "Katie"])
-      try expect(context["name"] as? String) == "Katie"
+      try expect(context["name"] as? String) == "Kyle"
     }
 
     $0.it("allows you to push a dictionary and run a closure then restoring previous state") {

Tests/VariableSpec.swift

@@ -21,7 +21,7 @@ func testVariable() {
     ])
 
 #if os(OSX)
-    context.push(["object": Object()])
+    context["object"] = Object()
 #endif
 
     $0.it("can resolve a string literal with double quotes") {