Removing temporary files create while processing HTML pages from string and providing a reference to the pdf output to be deleted manually #18

src/main/java/com/github/jhonnymertz/wkhtmltopdf/wrapper/Pdf.java

@@ -1,19 +1,23 @@
 package com.github.jhonnymertz.wkhtmltopdf.wrapper;
 
-import java.io.File;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.UUID;
-import org.apache.commons.io.FileUtils;
-import org.apache.commons.io.IOUtils;
-import org.apache.commons.lang3.StringUtils;
 import com.github.jhonnymertz.wkhtmltopdf.wrapper.configurations.WrapperConfig;
 import com.github.jhonnymertz.wkhtmltopdf.wrapper.page.Page;
 import com.github.jhonnymertz.wkhtmltopdf.wrapper.page.PageType;
 import com.github.jhonnymertz.wkhtmltopdf.wrapper.params.Param;
 import com.github.jhonnymertz.wkhtmltopdf.wrapper.params.Params;
+import org.apache.commons.io.FileUtils;
+import org.apache.commons.io.IOUtils;
+import org.apache.commons.lang3.StringUtils;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
 
+/**
+ * Represents a Pdf file
+ */
 public class Pdf {
 
     private static final String STDINOUT = "-";
@@ -27,7 +31,7 @@ public class Pdf {
     private boolean hasToc = false;
 
     public Pdf() {
-      this(new WrapperConfig());
+        this(new WrapperConfig());
     }
 
     public Pdf(WrapperConfig wrapperConfig) {
@@ -36,42 +40,74 @@ public class Pdf {
         this.pages = new ArrayList<Page>();
     }
 
+    /**
+     * Add a page to the pdf
+     *
+     * @deprecated Use the specific type method to a better semantic
+     */
+    @Deprecated
     public void addPage(String source, PageType type) {
         this.pages.add(new Page(source, type));
     }
 
+    /**
+     * Add a page from an URL to the pdf
+     */
+    public void addPageFromUrl(String source) {
+        this.pages.add(new Page(source, PageType.url));
+    }
+
+    /**
+     * Add a page from a HTML-based string to the pdf
+     */
+    public void addPageFromString(String source) {
+        this.pages.add(new Page(source, PageType.htmlAsString));
+    }
+
+    /**
+     * Add a page from a file to the pdf
+     */
+    public void addPageFromFile(String source) {
+        this.pages.add(new Page(source, PageType.file));
+    }
+
     public void addToc() {
         this.hasToc = true;
     }
 
     public void addParam(Param param, Param... params) {
-        this.params.add( param, params );
+        this.params.add(param, params);
     }
 
-    public void saveAs(String path) throws IOException, InterruptedException {
-        saveAs(path, getPDF());
+    public File saveAs(String path) throws IOException, InterruptedException {
+        return saveAs(path, getPDF());
     }
 
-    private static File saveAs(String path, byte[] document) throws IOException {
+    private File saveAs(String path, byte[] document) throws IOException {
         File file = new File(path);
-        FileUtils.writeByteArrayToFile( file, document );
+        FileUtils.writeByteArrayToFile(file, document);
 
         return file;
     }
 
     public byte[] getPDF() throws IOException, InterruptedException {
-        Process process = Runtime.getRuntime().exec(getCommandAsArray());
 
-        byte[] inputBytes = IOUtils.toByteArray( process.getInputStream() );
-        byte[] errorBytes = IOUtils.toByteArray( process.getErrorStream() );
-        
-        process.waitFor();
+        try {
+            Process process = Runtime.getRuntime().exec(getCommandAsArray());
 
-        if (process.exitValue() != 0) {
-            throw new RuntimeException("Process (" + getCommand() + ") exited with status code " + process.exitValue() + ":\n" + new String(errorBytes));
-        }
+            byte[] inputBytes = IOUtils.toByteArray(process.getInputStream());
+            byte[] errorBytes = IOUtils.toByteArray(process.getErrorStream());
 
-        return inputBytes;
+            process.waitFor();
+
+            if (process.exitValue() != 0) {
+                throw new RuntimeException("Process (" + getCommand() + ") exited with status code " + process.exitValue() + ":\n" + new String(errorBytes));
+            }
+
+            return inputBytes;
+        } finally {
+            cleanTempFiles();
+        }
     }
 
     private String[] getCommandAsArray() throws IOException {
@@ -102,6 +138,14 @@ public class Pdf {
         return commandLine.toArray(new String[commandLine.size()]);
     }
 
+    private void cleanTempFiles() {
+        for (Page page : pages) {
+            if (page.getType().equals(PageType.htmlAsString)) {
+                new File(page.getSource()).delete();
+            }
+        }
+    }
+
     public String getCommand() throws IOException {
         return StringUtils.join(getCommandAsArray(), " ");
     }

src/tests/java/com/github/jhonnymertz/wkhtmltopdf/wrapper/PdfTest.java

@@ -1,7 +1,6 @@
 package com.github.jhonnymertz.wkhtmltopdf.wrapper;
 
 import com.github.jhonnymertz.wkhtmltopdf.wrapper.configurations.WrapperConfig;
-import com.github.jhonnymertz.wkhtmltopdf.wrapper.page.PageType;
 import com.github.jhonnymertz.wkhtmltopdf.wrapper.params.Param;
 import org.apache.pdfbox.pdfparser.PDFParser;
 import org.apache.pdfbox.pdmodel.PDDocument;
@@ -10,6 +9,7 @@ import org.junit.Assert;
 import org.junit.Test;
 
 import java.io.ByteArrayInputStream;
+import java.io.File;
 
 import static org.hamcrest.core.StringContains.containsString;
 
@@ -20,7 +20,7 @@ public class PdfTest {
         Pdf pdf = new Pdf();
         pdf.addToc();
         pdf.addParam(new Param("--enable-javascript"), new Param("--html-header", "file:///example.html"));
-        pdf.addPage("http://www.google.com", PageType.url);
+        pdf.addPageFromUrl("http://www.google.com");
         Assert.assertThat("command params should contain the --enable-javascript and --html-header", pdf.getCommand(), containsString("--enable-javascript --html-header file:///example.html"));
     }
 
@@ -35,7 +35,7 @@ public class PdfTest {
 
         // GIVEN a html template containing special characters that java stores in utf-16 internally
         Pdf pdf = new Pdf();
-        pdf.addPage("<html><head><meta charset=\"utf-8\"></head><h1>Müller</h1></html>", PageType.htmlAsString);
+        pdf.addPageFromString("<html><head><meta charset=\"utf-8\"></head><h1>Müller</h1></html>");
 
         // WHEN
         byte[] pdfBytes = pdf.getPDF();
@@ -53,10 +53,10 @@ public class PdfTest {
     @Test
     public void testMultiplePages() throws Exception {
         Pdf pdf = new Pdf();
-        pdf.addPage("<html><head><meta charset=\"utf-8\"></head><h1>Page 1</h1></html>", PageType.htmlAsString);
-        pdf.addPage("<html><head><meta charset=\"utf-8\"></head><h1>Page 2</h1></html>", PageType.htmlAsString);
-        pdf.addPage("http://www.google.com", PageType.url);
-        pdf.addPage("<html><head><meta charset=\"utf-8\"></head><h1>Page 4</h1></html>", PageType.htmlAsString);
+        pdf.addPageFromString("<html><head><meta charset=\"utf-8\"></head><h1>Page 1</h1></html>");
+        pdf.addPageFromString("<html><head><meta charset=\"utf-8\"></head><h1>Page 2</h1></html>");
+        pdf.addPageFromUrl("http://www.google.com");
+        pdf.addPageFromString("<html><head><meta charset=\"utf-8\"></head><h1>Page 4</h1></html>");
 
         // WHEN
         byte[] pdfBytes = pdf.getPDF();
@@ -70,4 +70,16 @@ public class PdfTest {
 
         Assert.assertThat("document should contain the fourth page name", pdfText, containsString("Page 4"));
     }
+
+    @Test
+    public void testRemovingGeneratedFile() throws Exception {
+        Pdf pdf = new Pdf();
+        pdf.addPageFromString("<html><head><meta charset=\"utf-8\"></head><h1>Page 1</h1></html>");
+        pdf.addPageFromString("<html><head><meta charset=\"utf-8\"></head><h1>Page 2</h1></html>");
+        pdf.addPageFromUrl("http://www.google.com");
+        pdf.addPageFromString("<html><head><meta charset=\"utf-8\"></head><h1>Page 4</h1></html>");
+
+        File savedPdf = pdf.saveAs("output.pdf");
+        Assert.assertTrue(savedPdf.delete());
+    }
 }