Re-added 'flat' argument to Wikicode.get_sections().

Requested by Σ.

CHANGELOG

@@ -5,6 +5,7 @@ v0.4 (unreleased):
   is a breaking change if you rely on the default behavior.
 - The 'matches' argument of Wikicode's filter methods now accepts a function
   (taking one argument, a Node, and returning a bool) in addition to a regex.
+- Re-added 'flat' argument to Wikicode.get_sections().
 - Wikicode.get_sections() now returns sections in the correct order.
 - Wikicode.matches() now accepts a tuple or list of strings/Wikicode objects
   instead of just a single string or Wikicode.

docs/changelog.rst

@@ -14,6 +14,7 @@ Unreleased
 - The *matches* argument of :py:class:`Wikicode's <.Wikicode>`
   :py:meth:`.filter` methods now accepts a function (taking one argument, a
   :py:class:`.Node`, and returning a bool) in addition to a regex.
+- Re-added *flat* argument to :py:meth:`.Wikicode.get_sections`.
 - :py:meth:`.Wikicode.get_sections` now returns sections in the correct order.
 - :py:meth:`.Wikicode.matches` now accepts a tuple or list of
   strings/:py:class:`.Wikicode` objects instead of just a single string or

mwparserfromhell/wikicode.py

@@ -409,7 +409,7 @@ class Wikicode(StringMixIn):
         """
         return list(self.ifilter(recursive, matches, flags, forcetype))
 
-    def get_sections(self, levels=None, matches=None, flags=FLAGS,
+    def get_sections(self, levels=None, matches=None, flags=FLAGS, flat=False,
                      include_lead=None, include_headings=True):
         """Return a list of sections within the page.
 
@@ -417,13 +417,13 @@ class Wikicode(StringMixIn):
         node list (implemented using :py:class:`~.SmartList`) so that changes
         to sections are reflected in the parent Wikicode object.
 
-        Each section contains all of its subsections. If *levels* is given, it
-        should be a iterable of integers; only sections whose heading levels
-        are within it will be returned. If *matches* is given, it should be a
-        regex to be matched against the titles of section headings; only
-        sections whose headings match the regex will be included. *flags* can
-        be used to override the default regex flags (see :py:meth:`ifilter`) if
-        *matches* is used.
+        Each section contains all of its subsections, unless *flat* is
+        ``True``. If *levels* is given, it should be a iterable of integers;
+        only sections whose heading levels are within it will be returned. If
+        *matches* is given, it should be a regex to be matched against the
+        titles of section headings; only sections whose headings match the
+        regex will be included. *flags* can be used to override the default
+        regex flags (see :py:meth:`ifilter`) if *matches* is used.
 
         If *include_lead* is ``True``, the first, lead section (without a
         heading) will be included in the list; ``False`` will not include it;
@@ -455,8 +455,9 @@ class Wikicode(StringMixIn):
                 start += 1
             while i < len(self.nodes):
                 node = self.nodes[i]
-                if isinstance(node, Heading) and node.level <= heading.level:
-                    break
+                if isinstance(node, Heading):
+                    if flat or node.level <= heading.level:
+                        break
                 i += 1
             sections.append(Wikicode(self.nodes[start:i]))
         return sections

tests/test_wikicode.py

@@ -362,14 +362,22 @@ class TestWikicode(TreeEqualityTestCase):
         self.assertEqual(["====Gnidaeh====\n"], page3.get_sections(levels=[4]))
         self.assertEqual(["===Heading===\nFoo bar baz\n====Gnidaeh====\n"],
                          page3.get_sections(levels=(2, 3)))
+        self.assertEqual(["===Heading===\nFoo bar baz\n"],
+                         page3.get_sections(levels=(2, 3), flat=True))
         self.assertEqual([], page3.get_sections(levels=[0]))
         self.assertEqual(["", "====Gnidaeh====\n"],
                          page3.get_sections(levels=[4], include_lead=True))
         self.assertEqual(["===Heading===\nFoo bar baz\n====Gnidaeh====\n",
                           "====Gnidaeh====\n"],
                          page3.get_sections(include_lead=False))
+        self.assertEqual(["===Heading===\nFoo bar baz\n", "====Gnidaeh====\n"],
+                         page3.get_sections(flat=True, include_lead=False))
 
         self.assertEqual([p4_IB1, p4_IIIA2], page4.get_sections(levels=[4]))
+        self.assertEqual([p4_IA, p4_IB, p4_IIIA], page4.get_sections(levels=[3]))
+        self.assertEqual([p4_IA, "=== Section I.B ===\n",
+                          "=== Section III.A ===\nText.\n"],
+                         page4.get_sections(levels=[3], flat=True))
         self.assertEqual(["", ""], page2.get_sections(include_headings=False))
         self.assertEqual(["\nSection I.B.1 body.\n\n•Some content.\n\n",
                           "\nEven more text.\n" + p4_IIIA2ai1],