Add notes to document

9 files changed, 98 insertions, 29 deletions

diff --git a/src/main/java/org/zwobble/mammoth/Mammoth.java b/src/main/java/org/zwobble/mammoth/Mammoth.java
index 26db504..2201bd2 100644
--- a/src/main/java/org/zwobble/mammoth/Mammoth.java
+++ b/src/main/java/org/zwobble/mammoth/Mammoth.java
@@ -1,5 +1,6 @@
 package org.zwobble.mammoth;
 
+import org.zwobble.mammoth.documents.Notes;
 import org.zwobble.mammoth.docx.*;
 import org.zwobble.mammoth.html.Html;
 import org.zwobble.mammoth.results.Result;
@@ -7,7 +8,6 @@ import org.zwobble.mammoth.xml.XmlElement;
 
 import java.io.File;
 import java.io.IOException;
-import java.io.InputStream;
 import java.util.zip.ZipFile;
 
 public class Mammoth {
@@ -19,11 +19,9 @@ public class Mammoth {
             Numbering numbering = Numbering.EMPTY;
             Relationships relationships = Relationships.EMPTY;
             ContentTypes contentTypes = ContentTypes.DEFAULT;
-            FileReader fileReader = new FileReader() {
-                @Override
-                public InputStream getInputStream(String uri) throws IOException {
-                    throw new UnsupportedOperationException();
-                }
+            Notes notes = Notes.EMPTY;
+            FileReader fileReader = uri -> {
+                throw new UnsupportedOperationException();
             };
             DocumentXmlReader reader = new DocumentXmlReader(new BodyXmlReader(
                 styles,
@@ -31,7 +29,7 @@ public class Mammoth {
                 relationships,
                 contentTypes,
                 zipFile,
-                fileReader));
+                fileReader), notes);
             // TODO: prefix
             String idPrefix = "document";
             return reader.readElement(documentXml)
diff --git a/src/main/java/org/zwobble/mammoth/documents/Document.java b/src/main/java/org/zwobble/mammoth/documents/Document.java
index 4d91925..c3a14d3 100644
--- a/src/main/java/org/zwobble/mammoth/documents/Document.java
+++ b/src/main/java/org/zwobble/mammoth/documents/Document.java
@@ -4,12 +4,18 @@ import java.util.List;
 
 public class Document implements HasChildren {
     private final List<DocumentElement> children;
+    private final Notes notes;
 
-    public Document(List<DocumentElement> children) {
+    public Document(List<DocumentElement> children, Notes notes) {
         this.children = children;
+        this.notes = notes;
     }
 
     public List<DocumentElement> getChildren() {
         return children;
     }
+
+    public Notes getNotes() {
+        return notes;
+    }
 }
diff --git a/src/main/java/org/zwobble/mammoth/documents/Note.java b/src/main/java/org/zwobble/mammoth/documents/Note.java
index 1a0d3c0..da6f68f 100644
--- a/src/main/java/org/zwobble/mammoth/documents/Note.java
+++ b/src/main/java/org/zwobble/mammoth/documents/Note.java
@@ -3,11 +3,21 @@ package org.zwobble.mammoth.documents;
 import java.util.List;
 
 public class Note {
+    private final NoteType noteType;
     private final String id;
     private final List<DocumentElement> body;
 
-    public Note(String id, List<DocumentElement> body) {
+    public Note(NoteType noteType, String id, List<DocumentElement> body) {
+        this.noteType = noteType;
         this.id = id;
         this.body = body;
     }
+
+    public NoteType getNoteType() {
+        return noteType;
+    }
+
+    public String getId() {
+        return id;
+    }
 }
diff --git a/src/main/java/org/zwobble/mammoth/documents/Notes.java b/src/main/java/org/zwobble/mammoth/documents/Notes.java
new file mode 100644
index 0000000..2959bb0
--- /dev/null
+++ b/src/main/java/org/zwobble/mammoth/documents/Notes.java
@@ -0,0 +1,27 @@
+package org.zwobble.mammoth.documents;
+
+import com.google.common.collect.Maps;
+import com.google.common.collect.Multimaps;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.zwobble.mammoth.util.MammothLists.list;
+
+public class Notes {
+    public final static Notes EMPTY = new Notes(list());
+
+    private final Map<NoteType, Map<String, Note>> notes;
+
+    public Notes(List<Note> notes) {
+        this.notes = Maps.transformValues(
+            Multimaps.index(notes, Note::getNoteType).asMap(),
+            notesOfType -> Maps.uniqueIndex(notesOfType, Note::getId));
+    }
+
+    public Optional<Note> findNote(NoteType noteType, String noteId) {
+        return Optional.ofNullable(notes.get(noteType))
+            .flatMap(notesOfType -> Optional.ofNullable(notesOfType.get(noteId)));
+    }
+}
diff --git a/src/main/java/org/zwobble/mammoth/docx/DocumentXmlReader.java b/src/main/java/org/zwobble/mammoth/docx/DocumentXmlReader.java
index 2d7bec0..5e94e2f 100644
--- a/src/main/java/org/zwobble/mammoth/docx/DocumentXmlReader.java
+++ b/src/main/java/org/zwobble/mammoth/docx/DocumentXmlReader.java
@@ -1,20 +1,23 @@
 package org.zwobble.mammoth.docx;
 
-import org.zwobble.mammoth.results.Result;
 import org.zwobble.mammoth.documents.Document;
+import org.zwobble.mammoth.documents.Notes;
+import org.zwobble.mammoth.results.Result;
 import org.zwobble.mammoth.xml.XmlElement;
 
 public class DocumentXmlReader {
     private final BodyXmlReader bodyReader;
+    private final Notes notes;
 
-    public DocumentXmlReader(BodyXmlReader bodyReader) {
+    public DocumentXmlReader(BodyXmlReader bodyReader, Notes notes) {
         this.bodyReader = bodyReader;
+        this.notes = notes;
     }
 
     public Result<Document> readElement(XmlElement element) {
         XmlElement body = element.findChild("w:body");
         return bodyReader.readElements(body.children())
             .toResult()
-            .map(Document::new);
+            .map(children -> new Document(children, notes));
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/docx/NotesXmlReader.java b/src/main/java/org/zwobble/mammoth/docx/NotesXmlReader.java
index 0838076..726c476 100644
--- a/src/main/java/org/zwobble/mammoth/docx/NotesXmlReader.java
+++ b/src/main/java/org/zwobble/mammoth/docx/NotesXmlReader.java
@@ -1,5 +1,6 @@
 package org.zwobble.mammoth.docx;
 
+import org.zwobble.mammoth.documents.NoteType;
 import org.zwobble.mammoth.results.Result;
 import org.zwobble.mammoth.documents.Note;
 import org.zwobble.mammoth.xml.XmlElement;
@@ -10,16 +11,22 @@ import static com.google.common.collect.Iterables.filter;
 import static com.google.common.collect.Iterables.transform;
 
 public class NotesXmlReader {
+    public static NotesXmlReader footnote(BodyXmlReader bodyReader) {
+        return new NotesXmlReader(bodyReader, "footnote", NoteType.FOOTNOTE);
+    }
+
     private final BodyXmlReader bodyReader;
-    private final String noteType;
+    private final String tagName;
+    private final NoteType noteType;
 
-    public NotesXmlReader(BodyXmlReader bodyReader, String noteType) {
+    private NotesXmlReader(BodyXmlReader bodyReader, String tagName, NoteType noteType) {
         this.bodyReader = bodyReader;
+        this.tagName = tagName;
         this.noteType = noteType;
     }
 
     public Result<List<Note>> readElement(XmlElement element) {
-        Iterable<XmlElement> elements = filter(element.findChildren("w:" + noteType), this::isNoteElement);
+        Iterable<XmlElement> elements = filter(element.findChildren("w:" + tagName), this::isNoteElement);
         return Result.concat(transform(elements, this::readNoteElement));
     }
 
@@ -37,6 +44,7 @@ public class NotesXmlReader {
         return bodyReader.readElements(element.children())
             .toResult()
             .map(children -> new Note(
+                noteType,
                 element.getAttribute("w:id"),
                 children));
     }
diff --git a/src/test/java/org/zwobble/mammoth/tests/DocumentConverterTests.java b/src/test/java/org/zwobble/mammoth/tests/DocumentConverterTests.java
index 3ae5d16..17bcdf1 100644
--- a/src/test/java/org/zwobble/mammoth/tests/DocumentConverterTests.java
+++ b/src/test/java/org/zwobble/mammoth/tests/DocumentConverterTests.java
@@ -37,7 +37,7 @@ public class DocumentConverterTests {
         assertThat(
             documentConverter().convertToHtml(new Document(list(
                 make(a(PARAGRAPH, with(CHILDREN, list(runWithText("Hello"))))),
-                make(a(PARAGRAPH, with(CHILDREN, list(runWithText("there")))))))),
+                make(a(PARAGRAPH, with(CHILDREN, list(runWithText("there")))))), Notes.EMPTY)),
 
             deepEquals(list(
                 Html.element("p", list(Html.text("Hello"))),
diff --git a/src/test/java/org/zwobble/mammoth/tests/docx/DocumentXmlTests.java b/src/test/java/org/zwobble/mammoth/tests/docx/DocumentXmlTests.java
index 5370db3..3eaa4ee 100644
--- a/src/test/java/org/zwobble/mammoth/tests/docx/DocumentXmlTests.java
+++ b/src/test/java/org/zwobble/mammoth/tests/docx/DocumentXmlTests.java
@@ -2,23 +2,25 @@ package org.zwobble.mammoth.tests.docx;
 
 import org.junit.Test;
 import org.zwobble.mammoth.documents.Document;
-import org.zwobble.mammoth.documents.Text;
-import org.zwobble.mammoth.docx.BodyXmlReader;
+import org.zwobble.mammoth.documents.Note;
+import org.zwobble.mammoth.documents.NoteType;
+import org.zwobble.mammoth.documents.Notes;
 import org.zwobble.mammoth.docx.DocumentXmlReader;
 import org.zwobble.mammoth.results.Result;
 import org.zwobble.mammoth.xml.XmlElement;
 import org.zwobble.mammoth.xml.XmlNodes;
 
-import static com.natpryce.makeiteasy.MakeItEasy.*;
+import static com.natpryce.makeiteasy.MakeItEasy.a;
+import static com.natpryce.makeiteasy.MakeItEasy.make;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.zwobble.mammoth.results.Result.success;
 import static org.zwobble.mammoth.tests.DeepReflectionMatcher.deepEquals;
-import static org.zwobble.mammoth.tests.documents.DocumentElementMakers.*;
+import static org.zwobble.mammoth.tests.documents.DocumentElementMakers.paragraphWithText;
+import static org.zwobble.mammoth.tests.docx.BodyXmlReaderMakers.bodyReader;
 import static org.zwobble.mammoth.util.MammothLists.list;
 import static org.zwobble.mammoth.xml.XmlNodes.element;
 
 public class DocumentXmlTests {
-    private static final BodyXmlReader bodyReader = make(a(BodyXmlReaderMakers.bodyReader));
 
     @Test
     public void canReadTextWithinDocument() {
@@ -29,13 +31,27 @@ public class DocumentXmlTests {
                         element("w:t", list(
                             XmlNodes.text("Hello!")))))))))));
 
-        Result<Document> document = new DocumentXmlReader(bodyReader).readElement(documentElement);
+        DocumentXmlReader reader = new DocumentXmlReader(make(a(bodyReader)), Notes.EMPTY);
+        Result<Document> document = reader.readElement(documentElement);
 
         assertThat(
             document,
-            deepEquals(success(new Document(list(
-                make(a(PARAGRAPH, with(CHILDREN, list(
-                    make(a(RUN, with(CHILDREN, list(
-                        new Text("Hello!"))))))))))))));
+            deepEquals(success(new Document(
+                list(paragraphWithText("Hello!")),
+                Notes.EMPTY))));
+    }
+
+    @Test
+    public void notesOfDocumentAreIncludedInDocument() {
+        Note note = new Note(NoteType.FOOTNOTE, "4", list(paragraphWithText("Hello")));
+        Notes notes = new Notes(list(note));
+        DocumentXmlReader reader = new DocumentXmlReader(make(a(bodyReader)), notes);
+
+        XmlElement documentElement = element("w:document", list(element("w:body")));
+        Result<Document> document = reader.readElement(documentElement);
+
+        assertThat(
+            document.getValue().getNotes().findNote(NoteType.FOOTNOTE, "4").get(),
+            deepEquals(note));
     }
 }
diff --git a/src/test/java/org/zwobble/mammoth/tests/docx/NotesXmlReaderTests.java b/src/test/java/org/zwobble/mammoth/tests/docx/NotesXmlReaderTests.java
index ef8489a..a5d5be7 100644
--- a/src/test/java/org/zwobble/mammoth/tests/docx/NotesXmlReaderTests.java
+++ b/src/test/java/org/zwobble/mammoth/tests/docx/NotesXmlReaderTests.java
@@ -1,6 +1,7 @@
 package org.zwobble.mammoth.tests.docx;
 
 import org.junit.Test;
+import org.zwobble.mammoth.documents.NoteType;
 import org.zwobble.mammoth.results.Result;
 import org.zwobble.mammoth.documents.Note;
 import org.zwobble.mammoth.docx.BodyXmlReader;
@@ -28,12 +29,12 @@ public class NotesXmlReaderTests {
             element("w:footnote", map("w:id", "1"), list(
                 element("w:p")))));
 
-        NotesXmlReader reader = new NotesXmlReader(bodyReader, "footnote");
+        NotesXmlReader reader = NotesXmlReader.footnote(bodyReader);
         Result<List<Note>> notes = reader.readElement(element);
 
         assertThat(
             notes,
-            deepEquals(success(list(new Note("1", list(make(a(PARAGRAPH))))))));
+            deepEquals(success(list(new Note(NoteType.FOOTNOTE, "1", list(make(a(PARAGRAPH))))))));
     }
 
     @Test
@@ -51,7 +52,7 @@ public class NotesXmlReaderTests {
             element("w:footnote", map("w:id", "1", "w:type", noteType), list(
                 element("w:p")))));
 
-        NotesXmlReader reader = new NotesXmlReader(bodyReader, "footnote");
+        NotesXmlReader reader = NotesXmlReader.footnote(bodyReader);
         Result<List<Note>> notes = reader.readElement(element);
 
         assertThat(notes, deepEquals(success(list())));