Defer resolution of abstract nums until necessary

3 files changed, 41 insertions, 21 deletions

diff --git a/src/main/java/org/zwobble/mammoth/internal/docx/Numbering.java b/src/main/java/org/zwobble/mammoth/internal/docx/Numbering.java
index 5ee09db..2680cfa 100644
--- a/src/main/java/org/zwobble/mammoth/internal/docx/Numbering.java
+++ b/src/main/java/org/zwobble/mammoth/internal/docx/Numbering.java
@@ -9,16 +9,31 @@ import static org.zwobble.mammoth.internal.util.Maps.lookup;
 import static org.zwobble.mammoth.internal.util.Maps.map;
 
 public class Numbering {
-    public static final Numbering EMPTY = new Numbering(map());
+    public static class Num {
+        private final Optional<String> abstractNumId;
 
-    private final Map<String, Map<String, NumberingLevel>> numbering;
+        public Num(Optional<String> abstractNumId) {
+            this.abstractNumId = abstractNumId;
+        }
+    }
+
+    public static final Numbering EMPTY = new Numbering(map(), map());
+
+    private final Map<String, Map<String, NumberingLevel>> abstractNums;
+    private final Map<String, Num> nums;
 
-    public Numbering(Map<String, Map<String, NumberingLevel>> numbering) {
-        this.numbering = numbering;
+    public Numbering(
+        Map<String, Map<String, NumberingLevel>> abstractNums,
+        Map<String, Num> nums
+    ) {
+        this.abstractNums = abstractNums;
+        this.nums = nums;
     }
 
     public Optional<NumberingLevel> findLevel(String numId, String level) {
-        return lookup(numbering, numId)
+        return lookup(nums, numId)
+            .flatMap(num -> num.abstractNumId)
+            .flatMap(abstractNumId -> lookup(this.abstractNums, abstractNumId))
             .flatMap(levels -> lookup(levels, level));
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/docx/NumberingXml.java b/src/main/java/org/zwobble/mammoth/internal/docx/NumberingXml.java
index fa1941f..33e5e27 100644
--- a/src/main/java/org/zwobble/mammoth/internal/docx/NumberingXml.java
+++ b/src/main/java/org/zwobble/mammoth/internal/docx/NumberingXml.java
@@ -8,13 +8,13 @@ import java.util.Map;
 import java.util.Optional;
 
 import static org.zwobble.mammoth.internal.util.Maps.entry;
-import static org.zwobble.mammoth.internal.util.Maps.lookup;
 import static org.zwobble.mammoth.internal.util.Maps.toMap;
 
 public class NumberingXml {
     public static Numbering readNumberingXmlElement(XmlElement element) {
         Map<String, Map<String, NumberingLevel>> abstractNums = readAbstractNums(element.findChildren("w:abstractNum"));
-        return new Numbering(readNums(element.findChildren("w:num"), abstractNums));
+        Map<String, Numbering.Num> nums = readNums(element.findChildren("w:num"));
+        return new Numbering(abstractNums, nums);
     }
 
     private static Map<String, Map<String, NumberingLevel>> readAbstractNums(XmlElementList children) {
@@ -38,20 +38,14 @@ public class NumberingXml {
         return entry(levelIndex, new NumberingLevel(levelIndex, isOrdered));
     }
 
-    private static Map<String, Map<String, NumberingLevel>> readNums(
-        XmlElementList numElements,
-        Map<String, Map<String, NumberingLevel>> abstractNums
-    ) {
-        return toMap(numElements, numElement -> readNum(numElement, abstractNums));
+    private static Map<String, Numbering.Num> readNums(XmlElementList numElements) {
+        return toMap(numElements, NumberingXml::readNum);
     }
 
-    private static Map.Entry<String,Map<String,NumberingLevel>> readNum(
-        XmlElement numElement,
-        Map<String, Map<String, NumberingLevel>> abstractNums
-    ) {
+    private static Map.Entry<String, Numbering.Num> readNum(XmlElement numElement) {
         // TODO: in python-mammoth, we allow None here. Check whether that's actually possible or not
         String numId = numElement.getAttribute("w:numId");
-        String abstractNumId = numElement.findChild("w:abstractNumId").get().getAttribute("w:val");
-        return entry(numId, lookup(abstractNums, abstractNumId).get());
+        Optional<String> abstractNumId = numElement.findChildOrEmpty("w:abstractNumId").getAttributeOrNone("w:val");
+        return entry(numId, new Numbering.Num(abstractNumId));
     }
 }
diff --git a/src/test/java/org/zwobble/mammoth/tests/docx/BodyXmlTests.java b/src/test/java/org/zwobble/mammoth/tests/docx/BodyXmlTests.java
index b12a5fe..e115c1a 100644
--- a/src/test/java/org/zwobble/mammoth/tests/docx/BodyXmlTests.java
+++ b/src/test/java/org/zwobble/mammoth/tests/docx/BodyXmlTests.java
@@ -23,6 +23,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 
 import static java.util.Arrays.asList;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -197,7 +198,7 @@ public class BodyXmlTests {
                     element("w:ilvl", map("w:val", "1")),
                     element("w:numId", map("w:val", "42"))))))));
 
-        Numbering numbering = new Numbering(map("42", map("1", NumberingLevel.ordered("1"))));
+        Numbering numbering = numberingMap(map("42", map("1", NumberingLevel.ordered("1"))));
 
         assertThat(
             readSuccess(bodyReader(numbering), element),
@@ -212,7 +213,7 @@ public class BodyXmlTests {
                 element("w:numPr", map(), list(
                     element("w:numId", map("w:val", "42"))))))));
 
-        Numbering numbering = new Numbering(map("42", map("1", NumberingLevel.ordered("1"))));
+        Numbering numbering = numberingMap(map("42", map("1", NumberingLevel.ordered("1"))));
 
         assertThat(
             readSuccess(bodyReader(numbering), element),
@@ -227,7 +228,7 @@ public class BodyXmlTests {
                 element("w:numPr", map(), list(
                     element("w:ilvl", map("w:val", "1"))))))));
 
-        Numbering numbering = new Numbering(map("42", map("1", NumberingLevel.ordered("1"))));
+        Numbering numbering = numberingMap(map("42", map("1", NumberingLevel.ordered("1"))));
 
         assertThat(
             readSuccess(bodyReader(numbering), element),
@@ -1437,4 +1438,14 @@ public class BodyXmlTests {
     private static Relationship imageRelationship(String relationshipId, String target) {
         return new Relationship(relationshipId, target, "http://schemas.openxmlformats.org/officeDocument/2006/relationships/image");
     }
+
+    private static Numbering numberingMap(Map<String, Map<String, NumberingLevel>> numbering) {
+        return new Numbering(
+            numbering,
+            numbering.keySet().stream().collect(Collectors.toMap(
+                numId -> numId,
+                numId -> new Numbering.Num(Optional.of(numId))
+            ))
+        );
+    }
 }