Ensure embedded style map is not removed by Word

7 files changed, 242 insertions, 12 deletions

diff --git a/src/main/java/org/zwobble/mammoth/internal/archives/MutableArchive.java b/src/main/java/org/zwobble/mammoth/internal/archives/MutableArchive.java
index d32da20..f8f69d2 100644
--- a/src/main/java/org/zwobble/mammoth/internal/archives/MutableArchive.java
+++ b/src/main/java/org/zwobble/mammoth/internal/archives/MutableArchive.java
@@ -1,7 +1,5 @@
 package org.zwobble.mammoth.internal.archives;
 
-import org.zwobble.mammoth.internal.archives.Archive;
-
 public interface MutableArchive extends Archive {
 
     void writeEntry(String path, String content);
diff --git a/src/main/java/org/zwobble/mammoth/internal/archives/ZippedArchive.java b/src/main/java/org/zwobble/mammoth/internal/archives/ZippedArchive.java
index 32b8894..853ced9 100644
--- a/src/main/java/org/zwobble/mammoth/internal/archives/ZippedArchive.java
+++ b/src/main/java/org/zwobble/mammoth/internal/archives/ZippedArchive.java
@@ -1,7 +1,5 @@
 package org.zwobble.mammoth.internal.archives;
 
-import org.zwobble.mammoth.internal.archives.Archive;
-
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
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 b01c0f5..4acce3f 100644
--- a/src/main/java/org/zwobble/mammoth/internal/docx/EmbeddedStyleMap.java
+++ b/src/main/java/org/zwobble/mammoth/internal/docx/EmbeddedStyleMap.java
@@ -1,20 +1,84 @@
 package org.zwobble.mammoth.internal.docx;
 
 import org.zwobble.mammoth.internal.archives.Archive;
+import org.zwobble.mammoth.internal.archives.Archives;
 import org.zwobble.mammoth.internal.archives.MutableArchive;
 import org.zwobble.mammoth.internal.util.Streams;
+import org.zwobble.mammoth.internal.xml.*;
+import org.zwobble.mammoth.internal.xml.parsing.XmlParser;
 
-import java.io.*;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 
+import static org.zwobble.mammoth.internal.util.Casts.tryCast;
+import static org.zwobble.mammoth.internal.util.Iterables.findIndex;
+import static org.zwobble.mammoth.internal.util.Maps.map;
+
 public class EmbeddedStyleMap {
-    private static final String PATH = "mammoth/style-map";
+    private static final String STYLE_MAP_PATH = "mammoth/style-map";
+    private static final String ABSOLUTE_STYLE_MAP_PATH = "/" + STYLE_MAP_PATH;
+    private static final String RELATIONSHIPS_PATH = "word/_rels/document.xml.rels";
+    private static final String CONTENT_TYPES_PATH = "[Content_Types].xml";
+
+    public static final NamespacePrefixes RELATIONSHIPS_NAMESPACES = NamespacePrefixes.builder()
+        .defaultPrefix("http://schemas.openxmlformats.org/package/2006/relationships")
+        .build();
+
+    public static final NamespacePrefixes CONTENT_TYPES_NAMESPACES = NamespacePrefixes.builder()
+        .defaultPrefix("http://schemas.openxmlformats.org/package/2006/content-types")
+        .build();
 
     public static Optional<String> readStyleMap(Archive file) throws IOException {
-        return file.tryGetInputStream(PATH).map(Streams::toString);
+        return file.tryGetInputStream(STYLE_MAP_PATH).map(Streams::toString);
+    }
+
+    public static void embedStyleMap(MutableArchive archive, String styleMap) throws IOException {
+        archive.writeEntry(STYLE_MAP_PATH, styleMap);
+        updateRelationships(archive);
+        updateContentTypes(archive);
     }
 
-    public static void embedStyleMap(MutableArchive archive, String styleMap) {
-        archive.writeEntry(PATH, styleMap);
+    private static void updateRelationships(MutableArchive archive) throws IOException {
+        XmlParser parser = new XmlParser(RELATIONSHIPS_NAMESPACES);
+        XmlElement relationships = parser.parseStream(Archives.getInputStream(archive, RELATIONSHIPS_PATH));
+        XmlElement relationship = XmlNodes.element("Relationship", map(
+            "Id", "rMammothStyleMap",
+            "Type", "http://schemas.zwobble.org/mammoth/style-map",
+            "Target", ABSOLUTE_STYLE_MAP_PATH
+        ));
+        XmlElement updatedRelationships = updateOrAddElement(relationships, relationship, "Id");
+        archive.writeEntry(RELATIONSHIPS_PATH, XmlWriter.toString(updatedRelationships, RELATIONSHIPS_NAMESPACES));
+    }
+
+    private static void updateContentTypes(MutableArchive archive) throws IOException {
+        XmlParser parser = new XmlParser(CONTENT_TYPES_NAMESPACES);
+        XmlElement contentTypes = parser.parseStream(Archives.getInputStream(archive, CONTENT_TYPES_PATH));
+        XmlElement override = XmlNodes.element("Override", map(
+            "PartName", ABSOLUTE_STYLE_MAP_PATH,
+            "ContentType", "text/prs.mammoth.style-map"
+        ));
+        XmlElement updatedRelationships = updateOrAddElement(contentTypes, override, "PartName");
+        archive.writeEntry(CONTENT_TYPES_PATH, XmlWriter.toString(updatedRelationships, CONTENT_TYPES_NAMESPACES));
+    }
+
+    private static XmlElement updateOrAddElement(XmlElement parent, XmlElement element, String identifyingAttribute) {
+        int index = findIndex(parent.getChildren(), child ->
+            tryCast(XmlElement.class, child)
+                .map(childElement ->
+                    childElement.getName().equals(element.getName()) &&
+                    childElement.getAttributeOrNone(identifyingAttribute).equals(element.getAttributeOrNone(identifyingAttribute))
+                )
+                .orElse(false)
+        );
+        List<XmlNode> children = new ArrayList<>(parent.getChildren());
+
+        if (index == -1) {
+            children.add(element);
+        } else {
+            children.set(index, element);
+        }
+        return new XmlElement(parent.getName(), parent.getAttributes(), children);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/docx/OfficeXml.java b/src/main/java/org/zwobble/mammoth/internal/docx/OfficeXml.java
index cbb1ab0..1cb7cd1 100644
--- a/src/main/java/org/zwobble/mammoth/internal/docx/OfficeXml.java
+++ b/src/main/java/org/zwobble/mammoth/internal/docx/OfficeXml.java
@@ -10,7 +10,7 @@ import static org.zwobble.mammoth.internal.util.Lists.eagerFlatMap;
 import static org.zwobble.mammoth.internal.util.Lists.list;
 
 public class OfficeXml {
-    private static NamespacePrefixes XML_NAMESPACES = NamespacePrefixes.builder()
+    private static final NamespacePrefixes XML_NAMESPACES = NamespacePrefixes.builder()
         .put("w", "http://schemas.openxmlformats.org/wordprocessingml/2006/main")
         .put("wp", "http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing")
         .put("a", "http://schemas.openxmlformats.org/drawingml/2006/main")
diff --git a/src/main/java/org/zwobble/mammoth/internal/util/Iterables.java b/src/main/java/org/zwobble/mammoth/internal/util/Iterables.java
index 04a987a..88e906e 100644
--- a/src/main/java/org/zwobble/mammoth/internal/util/Iterables.java
+++ b/src/main/java/org/zwobble/mammoth/internal/util/Iterables.java
@@ -138,4 +138,15 @@ public class Iterables {
     public static <T> Stream<T> stream(Iterable<T> iterable) {
         return StreamSupport.stream(iterable.spliterator(), false);
     }
+
+    public static <T> int findIndex(Iterable<T> iterable, Predicate<T> predicate) {
+        int index = 0;
+        for (T element : iterable) {
+            if (predicate.test(element)) {
+                return index;
+            }
+            index++;
+        }
+        return -1;
+    }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/xml/NamespacePrefixes.java b/src/main/java/org/zwobble/mammoth/internal/xml/NamespacePrefixes.java
index 44e5676..ff644f1 100644
--- a/src/main/java/org/zwobble/mammoth/internal/xml/NamespacePrefixes.java
+++ b/src/main/java/org/zwobble/mammoth/internal/xml/NamespacePrefixes.java
@@ -1,12 +1,15 @@
 package org.zwobble.mammoth.internal.xml;
 
+import org.zwobble.mammoth.internal.util.Iterables;
+
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Optional;
 
 import static org.zwobble.mammoth.internal.util.Maps.lookup;
 
-public class NamespacePrefixes {
+public class NamespacePrefixes implements Iterable<NamespacePrefix> {
     public static Builder builder() {
         return new Builder();
     }
@@ -23,7 +26,7 @@ public class NamespacePrefixes {
             uriToPrefix.put(uri, new NamespacePrefix(Optional.empty(), uri));
             return this;
         }
-        
+
         public NamespacePrefixes build() {
             return new NamespacePrefixes(uriToPrefix);
         }
@@ -38,4 +41,21 @@ public class NamespacePrefixes {
     public Optional<NamespacePrefix> lookupUri(String uri) {
         return lookup(uriToPrefix, uri);
     }
+
+    public Optional<NamespacePrefix> lookupPrefix(String prefix) {
+        return lookupPrefix(Optional.of(prefix));
+    }
+
+    public Optional<NamespacePrefix> defaultNamespace() {
+        return lookupPrefix(Optional.empty());
+    }
+
+    private Optional<NamespacePrefix> lookupPrefix(Optional<String> prefix) {
+        return Iterables.tryFind(uriToPrefix.values(), namespace -> namespace.getPrefix().equals(prefix));
+    }
+
+    @Override
+    public Iterator<NamespacePrefix> iterator() {
+        return uriToPrefix.values().iterator();
+    }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/xml/XmlWriter.java b/src/main/java/org/zwobble/mammoth/internal/xml/XmlWriter.java
new file mode 100644
index 0000000..7882258
--- /dev/null
+++ b/src/main/java/org/zwobble/mammoth/internal/xml/XmlWriter.java
@@ -0,0 +1,139 @@
+package org.zwobble.mammoth.internal.xml;
+
+import javax.xml.stream.XMLOutputFactory;
+import javax.xml.stream.XMLStreamException;
+import javax.xml.stream.XMLStreamWriter;
+import java.io.ByteArrayOutputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+public class XmlWriter {
+    public static String toString(XmlElement element, NamespacePrefixes namespaces) {
+        try {
+            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+            XmlWriter writer = new XmlWriter(createXmlWriter(outputStream), namespaces);
+            writer.writeDocument(element);
+            return new String(outputStream.toByteArray(), StandardCharsets.UTF_8);
+        } catch (XMLStreamException exception) {
+            throw new RuntimeException(exception);
+        }
+    }
+
+    private static XMLStreamWriter createXmlWriter(ByteArrayOutputStream outputStream) throws XMLStreamException {
+        XMLOutputFactory outputFactory = XMLOutputFactory.newInstance();
+        return outputFactory.createXMLStreamWriter(outputStream);
+    }
+
+    private final XMLStreamWriter writer;
+    private final NamespacePrefixes namespaces;
+
+    private XmlWriter(XMLStreamWriter writer, NamespacePrefixes namespaces) {
+        this.writer = writer;
+        this.namespaces = namespaces;
+    }
+
+    private void writeDocument(XmlElement element) throws XMLStreamException {
+        writer.writeStartDocument("UTF-8", "1.0");
+        writeStartElement(element);
+        writeNamespaces(namespaces);
+        writeAttributes(element);
+        writeNodes(element.getChildren());
+        writer.writeEndElement();
+        writer.writeEndDocument();
+    }
+
+    private void writeNodes(List<XmlNode> children) {
+        for (XmlNode child : children) {
+            writeNode(child);
+        }
+    }
+
+    private void writeNode(XmlNode child) {
+        child.accept(new XmlNodeVisitor<Object>() {
+            @Override
+            public Object visit(XmlElement element) {
+                try {
+                    writeStartElement(element);
+                    writeAttributes(element);
+                    writeNodes(element.getChildren());
+                    writer.writeEndElement();
+                } catch (XMLStreamException exception) {
+                    throw new RuntimeException(exception);
+                }
+                return null;
+            }
+
+            @Override
+            public Object visit(XmlTextNode textNode) {
+                try {
+                    writer.writeCharacters(textNode.getValue());
+                } catch (XMLStreamException exception) {
+                    throw new RuntimeException(exception);
+                }
+                return null;
+            }
+        });
+    }
+
+    private void writeNamespaces(NamespacePrefixes namespaces) throws XMLStreamException {
+        for (NamespacePrefix namespace : namespaces) {
+            Optional<String> prefix = namespace.getPrefix();
+            if (prefix.isPresent()) {
+                writer.writeNamespace(prefix.get(), namespace.getUri());
+            } else {
+                writer.writeDefaultNamespace(namespace.getUri());
+            }
+        }
+    }
+
+    private void writeStartElement(XmlElement element) throws XMLStreamException {
+        XmlName name = readName(element.getName());
+        Optional<String> prefix = name.namespace.getPrefix();
+        if (prefix.isPresent()) {
+            writer.writeStartElement(prefix.get(), name.localName, name.namespace.getUri());
+        } else {
+            writer.writeStartElement(name.localName);
+        }
+    }
+
+    private void writeAttributes(XmlElement element) throws XMLStreamException {
+        for (Map.Entry<String, String> attribute : element.getAttributes().entrySet()) {
+            writeAttribute(attribute);
+        }
+    }
+
+    private void writeAttribute(Map.Entry<String, String> attribute) throws XMLStreamException {
+        XmlName name = readName(attribute.getKey());
+        Optional<String> prefix = name.namespace.getPrefix();
+        if (prefix.isPresent()) {
+            writer.writeAttribute(prefix.get(), name.namespace.getUri(), name.localName, attribute.getValue());
+        } else {
+            writer.writeAttribute(name.localName, attribute.getValue());
+        }
+    }
+
+    private XmlName readName(String name) {
+        String[] parts = name.split(":", 2);
+        if (parts.length == 1) {
+            return new XmlName(namespaces.defaultNamespace().get(), parts[0]);
+        } else {
+            String prefix = parts[0];
+            String localName = parts[1];
+            NamespacePrefix namespace = namespaces.lookupPrefix(prefix)
+                .orElseThrow(() -> new RuntimeException("Could not find namespace for prefix: " + prefix));
+            return new XmlName(namespace, localName);
+        }
+    }
+
+    private static class XmlName {
+        private final NamespacePrefix namespace;
+        private final String localName;
+
+        private XmlName(NamespacePrefix namespace, String localName) {
+            this.namespace = namespace;
+            this.localName = localName;
+        }
+    }
+}