Add reading of rowspan

3 files changed, 183 insertions, 11 deletions

diff --git a/mammoth/documents.py b/mammoth/documents.py
index 7f17f3e..36e3458 100644
--- a/mammoth/documents.py
+++ b/mammoth/documents.py
@@ -49,6 +49,7 @@ class TableRow(HasChildren):
 @cobble.data
 class TableCell(HasChildren):
     colspan = cobble.field()
+    rowspan = cobble.field()
 
 @cobble.data
 class LineBreak(Element):
@@ -111,10 +112,12 @@ bookmark = Bookmark
 
 table = Table
 table_row = TableRow
-def table_cell(children, colspan=None):
+def table_cell(children, colspan=None, rowspan=None):
     if colspan is None:
         colspan = 1
-    return TableCell(children=children, colspan=colspan)
+    if rowspan is None:
+        rowspan = 1
+    return TableCell(children=children, colspan=colspan, rowspan=rowspan)
 
 
 line_break = LineBreak
diff --git a/mammoth/docx/body_xml.py b/mammoth/docx/body_xml.py
index a44169a..8aacc1c 100644
--- a/mammoth/docx/body_xml.py
+++ b/mammoth/docx/body_xml.py
@@ -155,6 +155,7 @@ def _create_reader(numbering, content_types, relationships, styles, docx_file, f
     
     def table(element):
         return _read_xml_elements(element.children) \
+            .flat_map(calculate_row_spans) \
             .map(documents.table)
     
     
@@ -164,8 +165,8 @@ def _create_reader(numbering, content_types, relationships, styles, docx_file, f
     
     
     def table_cell(element):
-        gridspan = element \
-            .find_child_or_null("w:tcPr") \
+        properties = element.find_child_or_null("w:tcPr")
+        gridspan = properties \
             .find_child_or_null("w:gridSpan") \
             .attributes.get("w:val")
         
@@ -175,11 +176,61 @@ def _create_reader(numbering, content_types, relationships, styles, docx_file, f
             colspan = int(gridspan)
         
         return _read_xml_elements(element.children) \
-            .map(lambda children: documents.table_cell(
-                children=children,
-                colspan=colspan
+            .map(lambda children: _add_attrs(
+                documents.table_cell(
+                    children=children,
+                    colspan=colspan
+                ),
+                _vmerge=read_vmerge(properties),
             ))
     
+    def read_vmerge(properties):
+        vmerge_element = properties.find_child("w:vMerge")
+        if vmerge_element is None:
+            return False
+        else:
+            val = vmerge_element.attributes.get("w:val")
+            return val == "continue" or not val
+    
+    
+    def calculate_row_spans(rows):
+        unexpected_non_rows = any(
+            not isinstance(row, documents.TableRow)
+            for row in rows
+        )
+        if unexpected_non_rows:
+            return _elements_result_with_messages(rows, [results.warning(
+                "unexpected non-row element in table, cell merging may be incorrect"
+            )])
+            
+        unexpected_non_cells = any(
+            not isinstance(cell, documents.TableCell)
+            for row in rows
+            for cell in row.children
+        )
+        if unexpected_non_cells:
+            return _elements_result_with_messages(rows, [results.warning(
+                "unexpected non-cell element in table row, cell merging may be incorrect"
+            )])
+        
+        columns = {}
+        for row in rows:
+            cell_index = 0
+            for cell in row.children:
+                if cell._vmerge and cell_index in columns:
+                    columns[cell_index].rowspan += 1
+                else:
+                    columns[cell_index] = cell
+                    cell._vmerge = False
+                cell_index += cell.colspan
+        
+        for row in rows:
+            row.children = lists.filter(lambda cell: not cell._vmerge, row.children)
+            for cell in row.children:
+                del cell._vmerge
+        
+        return _success(rows)
+    
     
     def read_child_elements(element):
         return _read_xml_elements(element.children)
@@ -362,22 +413,38 @@ class _ReadResult(object):
         self.messages = messages
     
     def map(self, func):
+        elements = func(self.elements)
+        if not isinstance(elements, list):
+            elements = [elements]
         return _ReadResult(
-            [func(self.elements)],
+            elements,
             self.extra,
             self.messages)
     
+    def flat_map(self, func):
+        result = func(self.elements)
+        return _ReadResult(
+            result.elements,
+            self.extra + result.extra,
+            self.messages + result.messages)
+        
+    
     def to_extra(self):
         return _ReadResult([], _concat(self.extra, self.elements), self.messages)
     
     def append_extra(self):
         return _ReadResult(_concat(self.elements, self.extra), [], self.messages)
 
-def _success(element):
-    return _ReadResult([element], [], [])
+def _success(elements):
+    if not isinstance(elements, list):
+        elements = [elements]
+    return _ReadResult(elements, [], [])
 
 def _element_result_with_messages(element, messages):
-    return _ReadResult([element], [], messages)
+    return _elements_result_with_messages([element], messages)
+
+def _elements_result_with_messages(elements, messages):
+    return _ReadResult(elements, [], messages)
 
 _empty_result = _ReadResult([], [], [])
 
@@ -390,3 +457,10 @@ def _concat(*values):
         for element in value:
             result.append(element)
     return result
+
+
+def _add_attrs(obj, **kwargs):
+    for key, value in kwargs.items():
+        setattr(obj, key, value)
+    
+    return obj
diff --git a/tests/docx/body_xml_tests.py b/tests/docx/body_xml_tests.py
index 8fb35cf..2773398 100644
--- a/tests/docx/body_xml_tests.py
+++ b/tests/docx/body_xml_tests.py
@@ -247,6 +247,88 @@ class ReadXmlElementTests(object):
 
 
     @istest
+    def vmerge_is_read_as_rowspan_for_table_cell(self):
+        element = xml_element("w:tbl", {}, [
+            w_tr(w_tc()),
+            w_tr(w_tc(properties=[w_vmerge("restart")])),
+            w_tr(w_tc(properties=[w_vmerge("continue")])),
+            w_tr(w_tc(properties=[w_vmerge("continue")])),
+            w_tr(w_tc())
+        ])
+        result = _read_and_get_document_xml_element(element)
+        expected_result = documents.table([
+            documents.table_row([documents.table_cell([])]),
+            documents.table_row([documents.table_cell([], rowspan=3)]),
+            documents.table_row([]),
+            documents.table_row([]),
+            documents.table_row([documents.table_cell([])]),
+        ])
+        assert_equal(expected_result, result)
+
+
+    @istest
+    def vmerge_without_val_is_treated_as_continue(self):
+        element = xml_element("w:tbl", {}, [
+            w_tr(w_tc(properties=[w_vmerge("restart")])),
+            w_tr(w_tc(properties=[w_vmerge(None)])),
+        ])
+        result = _read_and_get_document_xml_element(element)
+        expected_result = documents.table([
+            documents.table_row([documents.table_cell([], rowspan=2)]),
+            documents.table_row([]),
+        ])
+        assert_equal(expected_result, result)
+
+
+    @istest
+    def vmerge_accounts_for_cells_spanning_columns(self):
+        element = xml_element("w:tbl", {}, [
+            w_tr(w_tc(), w_tc(), w_tc(properties=[w_vmerge("restart")])),
+            w_tr(w_tc(properties=[w_gridspan("2")]), w_tc(properties=[w_vmerge("continue")])),
+            w_tr(w_tc(), w_tc(), w_tc(properties=[w_vmerge("continue")])),
+            w_tr(w_tc(), w_tc(), w_tc()),
+        ])
+        result = _read_and_get_document_xml_element(element)
+        expected_result = documents.table([
+            documents.table_row([documents.table_cell([]), documents.table_cell([]), documents.table_cell([], rowspan=3)]),
+            documents.table_row([documents.table_cell([], colspan=2)]),
+            documents.table_row([documents.table_cell([]), documents.table_cell([])]),
+            documents.table_row([documents.table_cell([]), documents.table_cell([]), documents.table_cell([])]),
+        ])
+        assert_equal(expected_result, result)
+
+
+    @istest
+    def no_vertical_cell_merging_if_merged_cells_do_not_line_up(self):
+        element = xml_element("w:tbl", {}, [
+            w_tr(w_tc(properties=[w_gridspan("2")]), w_tc(properties=[w_vmerge("restart")])),
+            w_tr(w_tc(), w_tc(properties=[w_vmerge("continue")])),
+        ])
+        result = _read_and_get_document_xml_element(element)
+        expected_result = documents.table([
+            documents.table_row([documents.table_cell([], colspan=2), documents.table_cell([])]),
+            documents.table_row([documents.table_cell([]), documents.table_cell([])]),
+        ])
+        assert_equal(expected_result, result)
+
+
+    @istest
+    def warning_if_non_row_in_table(self):
+        element = xml_element("w:tbl", {}, [xml_element("w:p")])
+        result = _read_document_xml_element(element)
+        expected_warning = results.warning("unexpected non-row element in table, cell merging may be incorrect")
+        assert_equal([expected_warning], result.messages)
+
+
+    @istest
+    def warning_if_non_cell_in_table_row(self):
+        element = xml_element("w:tbl", {}, [w_tr(xml_element("w:p"))])
+        result = _read_document_xml_element(element)
+        expected_warning = results.warning("unexpected non-cell element in table row, cell merging may be incorrect")
+        assert_equal([expected_warning], result.messages)
+
+
+    @istest
     def children_of_w_ins_are_converted_normally(self):
         element = xml_element("w:p", {}, [
             xml_element("w:ins", {}, [
@@ -662,3 +744,16 @@ def _linked_blip(relationship_id):
 
 def _blip(attributes):
     return xml_element("a:blip", attributes)
+
+
+def w_tr(*children):
+    return xml_element("w:tr", {}, children)
+
+def w_tc(properties=None, *children):
+    return xml_element("w:tc", {}, [xml_element("w:tcPr", {}, properties)] + list(children))
+
+def w_gridspan(val):
+    return xml_element("w:gridSpan", {"w:val": val})
+
+def w_vmerge(val):
+    return xml_element("w:vMerge", {"w:val": val})