Add basic embedded of style maps

7 files changed, 108 insertions, 23 deletions

diff --git a/src/main/java/org/zwobble/mammoth/internal/docx/Archives.java b/src/main/java/org/zwobble/mammoth/internal/docx/Archives.java
index ae4f7e3..519520a 100644
--- a/src/main/java/org/zwobble/mammoth/internal/docx/Archives.java
+++ b/src/main/java/org/zwobble/mammoth/internal/docx/Archives.java
@@ -4,7 +4,7 @@ import java.io.IOException;
 import java.io.InputStream;
 
 public class Archives {
-    static InputStream getInputStream(Archive file, String name) throws IOException {
+    public static InputStream getInputStream(Archive file, String name) throws IOException {
         return file.tryGetInputStream(name)
             .orElseThrow(() -> new IOException("Missing entry in file: " + name));
     }
diff --git a/src/main/java/org/zwobble/mammoth/internal/docx/EmbeddedStyleMap.java b/src/main/java/org/zwobble/mammoth/internal/docx/EmbeddedStyleMap.java
index 2befa20..d268bf0 100644
--- a/src/main/java/org/zwobble/mammoth/internal/docx/EmbeddedStyleMap.java
+++ b/src/main/java/org/zwobble/mammoth/internal/docx/EmbeddedStyleMap.java
@@ -1,23 +1,18 @@
 package org.zwobble.mammoth.internal.docx;
 
-import java.io.BufferedReader;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.InputStreamReader;
+import org.zwobble.mammoth.internal.util.Streams;
+
+import java.io.*;
 import java.util.Optional;
-import java.util.stream.Collectors;
 
 public class EmbeddedStyleMap {
     private static final String PATH = "mammoth/style-map";
 
     public static Optional<String> readStyleMap(Archive file) throws IOException {
-        return file.tryGetInputStream(PATH)
-            .map(inputStream -> readInputStream(inputStream));
+        return file.tryGetInputStream(PATH).map(Streams::toString);
     }
 
-    private static String readInputStream(InputStream inputStream) {
-        return new BufferedReader(new InputStreamReader(inputStream))
-            .lines()
-            .collect(Collectors.joining("\n"));
+    public static void embedStyleMap(MutableArchive archive, String styleMap) {
+        archive.writeEntry(PATH, styleMap);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/docx/InMemoryArchive.java b/src/main/java/org/zwobble/mammoth/internal/docx/InMemoryArchive.java
index fa635c1..6d164f4 100644
--- a/src/main/java/org/zwobble/mammoth/internal/docx/InMemoryArchive.java
+++ b/src/main/java/org/zwobble/mammoth/internal/docx/InMemoryArchive.java
@@ -1,6 +1,10 @@
 package org.zwobble.mammoth.internal.docx;
 
+import org.zwobble.mammoth.internal.util.PassThroughException;
+import org.zwobble.mammoth.internal.util.Streams;
+
 import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
@@ -9,23 +13,23 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipInputStream;
+import java.util.zip.ZipOutputStream;
 
 import static org.zwobble.mammoth.internal.util.Maps.eagerMapValues;
 import static org.zwobble.mammoth.internal.util.Maps.lookup;
-import static org.zwobble.mammoth.internal.util.Streams.toByteArray;
 
-public class InMemoryArchive implements Archive {
-    public static Archive fromStream(InputStream stream) throws IOException {
+public class InMemoryArchive implements MutableArchive {
+    public static InMemoryArchive fromStream(InputStream stream) throws IOException {
         ZipInputStream zipStream = new ZipInputStream(stream);
         Map<String, byte[]> entries = new HashMap<>();
         ZipEntry entry;
         while ((entry = zipStream.getNextEntry()) != null) {
-            entries.put(entry.getName(), toByteArray(zipStream));
+            entries.put(entry.getName(), Streams.toByteArray(zipStream));
         }
         return new InMemoryArchive(entries);
     }
 
-    public static Archive fromStrings(Map<String, String> entries) {
+    public static InMemoryArchive fromStrings(Map<String, String> entries) {
         return new InMemoryArchive(eagerMapValues(entries, value -> value.getBytes(StandardCharsets.UTF_8)));
     }
 
@@ -42,6 +46,24 @@ public class InMemoryArchive implements Archive {
     }
 
     @Override
+    public void writeEntry(String path, String content) {
+        entries.put(path, content.getBytes(StandardCharsets.UTF_8));
+    }
+
+    @Override
     public void close() throws IOException {
     }
+
+    public byte[] toByteArray() {
+        return PassThroughException.wrap(() -> {
+            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+            try (ZipOutputStream zipStream = new ZipOutputStream(outputStream)) {
+                for (Map.Entry<String, byte[]> entry : entries.entrySet()) {
+                    zipStream.putNextEntry(new ZipEntry(entry.getKey()));
+                    zipStream.write(entry.getValue());
+                }
+            }
+            return outputStream.toByteArray();
+        });
+    }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/docx/MutableArchive.java b/src/main/java/org/zwobble/mammoth/internal/docx/MutableArchive.java
new file mode 100644
index 0000000..5615d58
--- /dev/null
+++ b/src/main/java/org/zwobble/mammoth/internal/docx/MutableArchive.java
@@ -0,0 +1,6 @@
+package org.zwobble.mammoth.internal.docx;
+
+public interface MutableArchive extends Archive {
+
+    void writeEntry(String path, String content);
+}
diff --git a/src/main/java/org/zwobble/mammoth/internal/util/Streams.java b/src/main/java/org/zwobble/mammoth/internal/util/Streams.java
index ac0f165..69cd796 100644
--- a/src/main/java/org/zwobble/mammoth/internal/util/Streams.java
+++ b/src/main/java/org/zwobble/mammoth/internal/util/Streams.java
@@ -1,9 +1,7 @@
 package org.zwobble.mammoth.internal.util;
 
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
+import java.io.*;
+import java.util.stream.Collectors;
 
 public class Streams {
     private Streams() {}
@@ -24,4 +22,10 @@ public class Streams {
         copy(stream, output);
         return output.toByteArray();
     }
+
+    public static String toString(InputStream stream) {
+        return new BufferedReader(new InputStreamReader(stream))
+            .lines()
+            .collect(Collectors.joining("\n"));
+    }
 }
diff --git a/src/test/java/org/zwobble/mammoth/tests/MammothTests.java b/src/test/java/org/zwobble/mammoth/tests/MammothTests.java
index 665e00a..ab67f07 100644
--- a/src/test/java/org/zwobble/mammoth/tests/MammothTests.java
+++ b/src/test/java/org/zwobble/mammoth/tests/MammothTests.java
@@ -3,10 +3,11 @@ package org.zwobble.mammoth.tests;
 import org.junit.Test;
 import org.zwobble.mammoth.DocumentConverter;
 import org.zwobble.mammoth.Result;
+import org.zwobble.mammoth.internal.docx.EmbeddedStyleMap;
+import org.zwobble.mammoth.internal.docx.InMemoryArchive;
 import org.zwobble.mammoth.internal.styles.parsing.ParseException;
 
-import java.io.File;
-import java.io.IOException;
+import java.io.*;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.function.Function;
@@ -235,6 +236,17 @@ public class MammothTests {
     }
 
     @Test
+    public void embeddedStyleMapCanBeWrittenAndThenRead() throws IOException {
+        InMemoryArchive archive = InMemoryArchive.fromStream(new FileInputStream(TestData.file("single-paragraph.docx")));
+        EmbeddedStyleMap.embedStyleMap(archive, "p => h1");
+
+        assertThat(
+            new DocumentConverter().convertToHtml(new ByteArrayInputStream(archive.toByteArray())),
+            isSuccess("<h1>Walking on imported air</h1>")
+        );
+    }
+
+    @Test
     public void canExtractRawTextFromFile() throws IOException {
         assertThat(
             new DocumentConverter().extractRawText(TestData.file("simple-list.docx")),
diff --git a/src/test/java/org/zwobble/mammoth/tests/docx/EmbeddedStyleMapTests.java b/src/test/java/org/zwobble/mammoth/tests/docx/EmbeddedStyleMapTests.java
new file mode 100644
index 0000000..1c1f9a8
--- /dev/null
+++ b/src/test/java/org/zwobble/mammoth/tests/docx/EmbeddedStyleMapTests.java
@@ -0,0 +1,46 @@
+package org.zwobble.mammoth.tests.docx;
+
+import org.junit.Test;
+import org.zwobble.mammoth.internal.docx.*;
+import org.zwobble.mammoth.internal.util.Streams;
+
+import java.io.IOException;
+import java.util.Optional;
+
+import static org.junit.Assert.assertEquals;
+import static org.zwobble.mammoth.internal.util.Maps.map;
+
+public class EmbeddedStyleMapTests {
+    private static final String ORIGINAL_RELATIONSHIPS_XML = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?>" +
+        "<Relationships xmlns=\"http://schemas.openxmlformats.org/package/2006/relationships\">" +
+        "<Relationship Id=\"rId3\" Type=\"http://schemas.openxmlformats.org/officeDocument/2006/relationships/settings\" Target=\"settings.xml\"/>" +
+        "</Relationships>";
+
+    private static final String ORIGINAL_CONTENT_TYPES_XML = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?>" +
+        "<Types xmlns=\"http://schemas.openxmlformats.org/package/2006/content-types\">" +
+        "<Default Extension=\"png\" ContentType=\"image/png\"/>" +
+        "</Types>";
+
+    @Test
+    public void embeddedStyleMapPreservesUnrelatedFiles() throws IOException {
+        MutableArchive archive = normalDocx();
+        EmbeddedStyleMap.embedStyleMap(archive, "p => h1");
+        assertEquals("placeholder text", Streams.toString(Archives.getInputStream(archive, "placeholder")));
+    }
+
+    @Test
+    public void embeddedStyleMapCanBeReadAfterBeingEmbedded() throws IOException {
+        MutableArchive archive = normalDocx();
+        EmbeddedStyleMap.embedStyleMap(archive, "p => h1");
+        assertEquals(Optional.of("p => h1"), EmbeddedStyleMap.readStyleMap(archive));
+    }
+
+    private static MutableArchive normalDocx() {
+        return InMemoryArchive.fromStrings(map(
+            "placeholder", "placeholder text",
+            "word/_rels/document.xml.rels", ORIGINAL_RELATIONSHIPS_XML,
+            "[Content_Types].xml", ORIGINAL_CONTENT_TYPES_XML
+        ));
+    }
+
+}