Fix Function vs TemplateFunction so that you can grab output

spullara · spullara · commit d730bb7f3ae6 · 2012-12-05T21:36:56.000-08:00
diff --git a/compiler/src/main/java/com/github/mustachejava/codes/IterableCode.java b/compiler/src/main/java/com/github/mustachejava/codes/IterableCode.java
@@ -87,17 +87,22 @@ public void run() {
   protected Writer handleFunction(Writer writer, Function function, Object[] scopes) {
     StringWriter sw = new StringWriter();
     runIdentity(sw);
-    Object newtemplate = function.apply(sw.toString());
-    if (newtemplate != null) {
-      if (function instanceof TemplateFunction) {
+    if (function instanceof TemplateFunction) {
+      Object newtemplate = function.apply(sw.toString());
+      if (newtemplate != null) {
         String templateText = newtemplate.toString();
         writer = writeTemplate(writer, templateText, scopes);
-      } else {
-        try {
-          writer.write(newtemplate.toString());
-        } catch (IOException e) {
-          throw new MustacheException("Failed to write function result", e);
+      }
+    } else {
+      try {
+        StringWriter capture = new StringWriter();
+        writeTemplate(capture, sw.toString(), scopes).close();
+        Object apply = function.apply(capture.toString());
+        if (apply != null) {
+          writer.write(apply.toString());
         }
+      } catch (IOException e) {
+        throw new MustacheException("Failed to write function result", e);
       }
     }
     return writer;
diff --git a/compiler/src/test/java/com/github/mustachejava/InterpreterTest.java b/compiler/src/test/java/com/github/mustachejava/InterpreterTest.java
@@ -4,6 +4,7 @@
 import com.github.mustachejava.util.CapturingMustacheVisitor;
 import com.github.mustachejavabenchmarks.JsonCapturer;
 import com.github.mustachejavabenchmarks.JsonInterpreterTest;
+import com.google.common.base.Function;
 import junit.framework.TestCase;
 import org.codehaus.jackson.JsonGenerator;
 import org.codehaus.jackson.JsonNode;
@@ -338,6 +339,39 @@ public String apply(@Nullable String s) {
     assertEquals("This is not interesting.", sw.toString());
   }
 
+  public void testFunctions() throws IOException {
+    MustacheFactory c = init();
+    Mustache m = c.compile(new StringReader("{{#f}}{{foo}}{{/f}}"), "test");
+    {
+      StringWriter sw = new StringWriter();
+      m.execute(sw, new Object() {
+        Function f = new Function<String, String>() {
+          @Override
+          public String apply(String s) {
+            return s.toUpperCase();
+          }
+        };
+        String foo = "bar";
+      }).flush();
+      assertEquals("BAR", sw.toString());
+    }
+    {
+      StringWriter sw = new StringWriter();
+      m.execute(sw, new Object() {
+        Function f = new TemplateFunction() {
+          @Override
+          public String apply(String s) {
+            return s.toUpperCase();
+          }
+        };
+        String foo = "bar";
+        String FOO = "baz";
+      }).flush();
+      assertEquals("baz", sw.toString());
+    }
+  }
+
+
   public void testComplex() throws MustacheException, IOException {
     StringWriter json = new StringWriter();
     MappingJsonFactory jf = new MappingJsonFactory();
diff --git a/compiler/src/test/java/com/github/mustachejava/SpecTest.java b/compiler/src/test/java/com/github/mustachejava/SpecTest.java
@@ -108,7 +108,7 @@ public String apply(String input) {
       });
       put("Section", new Object() {
         Function lambda() {
-          return new Function<String, String>() {
+          return new TemplateFunction() {
             @Override
             public String apply(String input) {
               return input.equals("{{x}}") ? "yes" : "no";
