From 26f30cbd9d5f4c7c686b6fd609963a2aa4955583 Mon Sep 17 00:00:00 2001
From: Kyle Fuller <kyle@fuller.li>
Date: Wed, 7 Dec 2016 21:46:04 +0000
Subject: [PATCH] feat: Allow subclassing templates (#79)

---
 CHANGELOG.md                             |  3 +++
 Sources/Environment.swift                |  6 ++++--
 Sources/Loader.swift                     |  6 ++++--
 Sources/Template.swift                   |  6 +++---
 Tests/StencilTests/EnvironmentSpec.swift | 15 ++++++++++++++-
 5 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fe1d608..7d2e693 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -56,6 +56,9 @@
 - You can now use `{{ block.super }}` to render a super block from another `{%
   block %}`.
 
+- `Environment` allows you to provide a custom `Template` subclass, allowing
+  new template to use a specific subclass.
+
 ### Deprecations
 
 - `Template` initialisers have been deprecated in favour of using a template
diff --git a/Sources/Environment.swift b/Sources/Environment.swift
index 5b68a1c..7473fc3 100644
--- a/Sources/Environment.swift
+++ b/Sources/Environment.swift
@@ -1,9 +1,11 @@
 public struct Environment {
+  public let templateClass: Template.Type
   public let extensions: [Extension]
 
   public var loader: Loader?
 
-  public init(loader: Loader? = nil, extensions: [Extension]? = nil) {
+  public init(loader: Loader? = nil, extensions: [Extension]? = nil, templateClass: Template.Type = Template.self) {
+    self.templateClass = templateClass
     self.loader = loader
     self.extensions = [DefaultExtension()] + (extensions ?? [])
   }
@@ -30,7 +32,7 @@ public struct Environment {
   }
 
   public func renderTemplate(string: String, context: [String: Any]? = nil) throws -> String {
-    let template = Template(templateString: string, environment: self)
+    let template = templateClass.init(templateString: string, environment: self)
     return try template.render(context)
   }
 }
diff --git a/Sources/Loader.swift b/Sources/Loader.swift
index fd49ece..a5457c9 100644
--- a/Sources/Loader.swift
+++ b/Sources/Loader.swift
@@ -51,7 +51,8 @@ public class FileSystemLoader: Loader, CustomStringConvertible {
         continue
       }
 
-      return try Template(path: templatePath, environment: environment, name: name)
+      let content: String = try templatePath.read()
+      return try environment.templateClass.init(templateString: content, environment: environment, name: name)
     }
 
     throw TemplateDoesNotExist(templateNames: [name], loader: self)
@@ -63,7 +64,8 @@ public class FileSystemLoader: Loader, CustomStringConvertible {
         let templatePath = try path.safeJoin(path: Path(templateName))
 
         if templatePath.exists {
-          return try Template(path: templatePath, environment: environment, name: templateName)
+          let content: String = try templatePath.read()
+          return try environment.templateClass.init(templateString: content, environment: environment, name: templateName)
         }
       }
     }
diff --git a/Sources/Template.swift b/Sources/Template.swift
index 825bc48..d9b1893 100644
--- a/Sources/Template.swift
+++ b/Sources/Template.swift
@@ -6,7 +6,7 @@ let NSFileNoSuchFileError = 4
 #endif
 
 /// A class representing a template
-public class Template: ExpressibleByStringLiteral {
+open class Template: ExpressibleByStringLiteral {
   let environment: Environment
   let tokens: [Token]
 
@@ -14,7 +14,7 @@ public class Template: ExpressibleByStringLiteral {
   public let name: String?
 
   /// Create a template with a template string
-  public init(templateString: String, environment: Environment? = nil, name: String? = nil) {
+  public required init(templateString: String, environment: Environment? = nil, name: String? = nil) {
     self.environment = environment ?? Environment()
     self.name = name
 
@@ -71,7 +71,7 @@ public class Template: ExpressibleByStringLiteral {
   }
 
   /// Render the given template
-  public func render(_ dictionary: [String: Any]? = nil) throws -> String {
+  open func render(_ dictionary: [String: Any]? = nil) throws -> String {
     return try render(Context(dictionary: dictionary, environment: environment))
   }
 }
diff --git a/Tests/StencilTests/EnvironmentSpec.swift b/Tests/StencilTests/EnvironmentSpec.swift
index 1c1e942..8bb5bfc 100644
--- a/Tests/StencilTests/EnvironmentSpec.swift
+++ b/Tests/StencilTests/EnvironmentSpec.swift
@@ -25,11 +25,17 @@ func testEnvironment() {
       let result = try environment.renderTemplate(name: "example.html")
       try expect(result) == "Hello World!"
     }
+
+    $0.it("allows you to provide a custom template class") {
+      let environment = Environment(loader: ExampleLoader(), templateClass: CustomTemplate.self)
+      let result = try environment.renderTemplate(string: "Hello World")
+
+      try expect(result) == "here"
+    }
   }
 }
 
 
-
 fileprivate class ExampleLoader: Loader {
   func loadTemplate(name: String, environment: Environment) throws -> Template {
     if name == "example.html" {
@@ -39,3 +45,10 @@ fileprivate class ExampleLoader: Loader {
     throw TemplateDoesNotExist(templateNames: [name], loader: self)
   }
 }
+
+
+class CustomTemplate: Template {
+  override func render(_ dictionary: [String: Any]? = nil) throws -> String {
+    return "here"
+  }
+}