From 67d4c525350a8637e79d8ab5af945e7891fd88d1 Mon Sep 17 00:00:00 2001
From: Kyle Fuller <kyle@fuller.li>
Date: Sat, 5 Mar 2016 00:23:08 +0000
Subject: [PATCH] [Context] Ensure pop happens when an error is thrown

---
 Sources/Context.swift     | 13 ++++++-------
 Sources/IfTag.swift       | 15 ++++++---------
 Sources/Inheritence.swift |  7 +++----
 Tests/ContextSpec.swift   | 25 ++++++++++---------------
 Tests/VariableSpec.swift  |  2 +-
 5 files changed, 26 insertions(+), 36 deletions(-)

diff --git a/Sources/Context.swift b/Sources/Context.swift
index 645f37d..131b619 100644
--- a/Sources/Context.swift
+++ b/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()
   }
 }
diff --git a/Sources/IfTag.swift b/Sources/IfTag.swift
index 9d0dcb1..9b80c2f 100644
--- a/Sources/IfTag.swift
+++ b/Sources/IfTag.swift
@@ -67,15 +67,12 @@ public class IfNode : NodeType {
       truthy = true
     }
 
-    context.push()
-    let output:String
-    if truthy {
-      output = try renderNodes(trueNodes, context)
-    } else {
-      output = try renderNodes(falseNodes, context)
+    return try context.push {
+      if truthy {
+        return try renderNodes(trueNodes, context)
+      } else {
+        return try renderNodes(falseNodes, context)
+      }
     }
-    context.pop()
-
-    return output
   }
 }
diff --git a/Sources/Inheritence.swift b/Sources/Inheritence.swift
index f1efa28..ab07a8e 100644
--- a/Sources/Inheritence.swift
+++ b/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)
+    }
   }
 }
 
diff --git a/Tests/ContextSpec.swift b/Tests/ContextSpec.swift
index f9ff19d..f0973a8 100644
--- a/Tests/ContextSpec.swift
+++ b/Tests/ContextSpec.swift
@@ -27,31 +27,26 @@ func testContext() {
     }
 
     $0.it("allows you to retrieve a value from a parent") {
-      context.push()
-
-      try expect(context["name"] as? String) == "Kyle"
+      try context.push {
+        try expect(context["name"] as? String) == "Kyle"
+      }
     }
 
     $0.it("allows you to override a parent's value") {
-      context.push()
-      context["name"] = "Katie"
-
-      try expect(context["name"] as? String) == "Katie"
+      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["name"] = "Katie"
-      context.pop()
+      context.push {
+        context["name"] = "Katie"
+      }
 
       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"
-    }
-
     $0.it("allows you to push a dictionary and run a closure then restoring previous state") {
       var didRun = false
 
diff --git a/Tests/VariableSpec.swift b/Tests/VariableSpec.swift
index 0b4b7f3..a42f910 100644
--- a/Tests/VariableSpec.swift
+++ b/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") {