feat(Table): define row_attrs property

5j9 · 5j9 · commit 5fcb4a1f13ff · 2022-08-29T07:34:54.000+04:30
It's both a getter and setter! closes #76
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -1,5 +1,6 @@
 Unreleased
 ----------
+- Feat: ``Table`` objects now have ``row_attrs`` property.
 - Fixed: Infinite loop on parsing tables containing ``\r``. (this is just to prevent infinite loop, CRLF line endings are not supported)
 
 v0.49.4
diff --git a/tests/test_table.py b/tests/test_table.py
@@ -572,3 +572,94 @@ def test_partially_invalid_table():  # 107
 def test_data_with_carriage_return():  # 107
     # this used to hang
     assert Table('{|\n}\n\r').data() == [[]]
+
+
+def test_row_attrs_getter():
+    # simple case
+    assert Table(
+        '{|\n'
+        '|- style="background: black;"\n'
+        '| cell1\n'
+        '|- style="background: white;"\n'
+        '| cell2\n'
+        '|}'
+    ).row_attrs == [
+        {'style': 'background: black;'},
+        {'style': 'background: white;'},
+    ]
+
+    # multiple attrs on a row
+    assert Table(
+        '{|\n'
+        '|-style="color: red;" bgcolor=yellow"\n'
+        '| cell1\n'
+        '|}'
+    ).row_attrs == [
+        {'style': 'color: red;', 'bgcolor': 'yellow'},
+    ]
+
+    # test middle row with no attrs
+    assert Table(
+        '{|\n'
+        '|- style="background: black;"\n'
+        '| cell1\n'
+        '|-\n'
+        '| cell2\n'
+        '|- style="background: white;"\n'
+        '| cell3\n'
+        '|}'
+    ).row_attrs == [
+        {'style': 'background: black;'},
+        {},
+        {'style': 'background: white;'},
+    ]
+
+
+def test_row_attrs_containing_comment():
+    assert Table(
+        '{|\n'
+        '|- style="color: red;" <!--comment--> bgcolor=yellow\n'
+        '| cell1\n'
+        '|}'
+    ).row_attrs == [
+        {'style': 'color: red;', 'bgcolor': 'yellow'},
+    ]
+
+
+def test_row_attrs_containing_template():
+    assert Table(
+        '{|\n'
+        '|- style={{text|"color: red;"}}\n'
+        '| cell1\n'
+        '|}'
+    ).row_attrs == [
+        {'style': '{{text|"color: red;"}}'},
+    ]
+
+
+def test_row_attrs_setter():
+    table = Table(
+        '{|\n'
+        '|- style="background: black;"\n'
+        '| cell1\n'
+        '|-\n'
+        '| cell2\n'
+        '|- style="background: white;"\n'
+        '| cell3\n'
+        '|}'
+    )
+    table.row_attrs = [
+        {'style': 'color: white;'},
+        {'style': 'color: black;'},
+        {},
+    ]
+    assert table.string == (
+        '{|\n'
+        '|- style="color: white;"\n'
+        '| cell1\n'
+        '|- style="color: black;"\n'
+        '| cell2\n'
+        '|-\n'
+        '| cell3\n'
+        '|}'
+    )
diff --git a/wikitextparser/_spans.py b/wikitextparser/_spans.py
@@ -123,7 +123,8 @@
 # HTML tags
 # Tags:
 # https://infra.spec.whatwg.org/#ascii-whitespace
-SPACE_CHARS = rb' \t\n\u000C\r'  # \s - \v
+# \0 was added as a special case for wikitextparser
+SPACE_CHARS = rb' \t\n\u000C\r\0'  # \s - \v
 # http://stackoverflow.com/a/93029/2705757
 # chrs = (chr(i) for i in range(sys.maxunicode))
 # control_chars = ''.join(c for c in chrs if unicodedata.category(c) == 'Cc')
diff --git a/wikitextparser/_table.py b/wikitextparser/_table.py
@@ -37,6 +37,7 @@
     (?:\n[\|\!]|\|\|)
     """, DOTALL | VERBOSE).match
 T = TypeVar('T')
+FIND_ROWS = rc(rb'\|-(.*)').finditer
 
 
 HEAD_DIGITS = rc(rb'\s*+\d+').match
@@ -312,6 +313,39 @@ def caption_attrs(self, attrs: str) -> None:
             if end != -1:
                 self[m.end('preattrs'):end] = attrs
 
+    @property
+    def row_attrs(self) -> list[dict]:
+        """Row attributes.
+
+        Use the setter of this property to set attributes for all rows.
+        Note that it will overwrite all the existing attr values.
+        """
+        shadow = self._table_shadow
+        string = self.string
+        attrs = []
+        append = attrs.append
+        for row_match in FIND_ROWS(shadow):
+            s, e = row_match.span(1)
+            spans = ATTRS_MATCH(shadow, s, e).spans
+            append({
+                string[ns: ne]: string[vs: ve]
+                for (ns, ne), (vs, ve) in
+                zip(spans('attr_name'), spans('attr_value'))
+            })
+        return attrs
+
+    @row_attrs.setter
+    def row_attrs(self, attrs: list[dict]):
+        for row_match, attrs_dict in reversed(
+            [*zip(FIND_ROWS(self._table_shadow), attrs)]
+        ):
+            s, e = row_match.span(1)
+            del self[s: e]
+            self.insert(s, ''.join([
+                f' {name}="{value}"' if value else f' {name}'
+                for name, value in attrs_dict.items()
+            ]))
+
 
 def _apply_attr_spans(
     table_attrs: List[List[Dict[str, str]]], table_data: List[List[T]]
