From cf7ab558bfd2f6b0bff2a9493e7e6f254ae9d341 Mon Sep 17 00:00:00 2001 From: Ben Kurtovic Date: Thu, 20 Mar 2014 00:58:37 -0400 Subject: [PATCH] Make Wikicode.get_sections() much faster (closes #63). Patch mostly by @spointy. --- CHANGELOG | 4 +- docs/changelog.rst | 4 +- mwparserfromhell/wikicode.py | 119 ++++++++++++++++++++++++++++--------------- 3 files changed, 83 insertions(+), 44 deletions(-) diff --git a/CHANGELOG b/CHANGELOG index 2c94ebc..a396f82 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -6,8 +6,8 @@ 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. +- Re-added 'flat' argument to Wikicode.get_sections(), fixed the order in which + it returns sections, and made it faster. - Wikicode.matches() now accepts a tuple or list of strings/Wikicode objects instead of just a single string or Wikicode. - Given the frequency of issues with the (admittedly insufficient) tag parser, diff --git a/docs/changelog.rst b/docs/changelog.rst index 6e1ce47..82f06c4 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -15,8 +15,8 @@ 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. +- Re-added *flat* argument to :py:meth:`.Wikicode.get_sections`, fixed the + order in which it returns sections, and made it faster. - :py:meth:`.Wikicode.matches` now accepts a tuple or list of strings/:py:class:`.Wikicode` objects instead of just a single string or :py:class:`.Wikicode`. diff --git a/mwparserfromhell/wikicode.py b/mwparserfromhell/wikicode.py index 312f3a0..4f5cd1a 100644 --- a/mwparserfromhell/wikicode.py +++ b/mwparserfromhell/wikicode.py @@ -68,6 +68,41 @@ class Wikicode(StringMixIn): substring = "".join(nodes).replace(old, new) code.nodes[index] = parse_anything(substring).nodes + @staticmethod + def _build_matcher(matches, flags): + """Helper for :py:meth:`_indexed_ifilter` and others. + + If *matches* is a function, return it. If it's a regex, return a + wrapper around it that can be called with a node to do a search. If + it's ``None``, return a function that always returns ``True``. + """ + if matches: + if callable(matches): + return matches + return lambda obj: re.search(matches, str(obj), flags) # r + return lambda obj: True + + def _indexed_ifilter(self, recursive=True, matches=None, flags=FLAGS, + forcetype=None): + """Iterate over nodes and their corresponding indices in the node list. + + The arguments are interpreted as for :py:meth:`ifilter`. For each tuple + ``(i, node)`` yielded by this method, ``self.index(node) == i``. Note + that if *recursive* is ``True``, ``self.nodes[i]`` might not be the + node itself, but will still contain it. + """ + match = self._build_matcher(matches, flags) + if recursive: + def getter(i, node): + for ch in self._get_children(node): + yield (i, ch) + inodes = chain(*(getter(i, n) for i, n in enumerate(self.nodes))) + else: + inodes = enumerate(self.nodes) + for i, node in inodes: + if (not forcetype or isinstance(node, forcetype)) and match(node): + yield (i, node) + def _do_strong_search(self, obj, recursive=True): """Search for the specific element *obj* within the node list. @@ -411,17 +446,8 @@ class Wikicode(StringMixIn): :py:const:`re.DOTALL`, and :py:const:`re.UNICODE`, but custom flags can be specified by passing *flags*. """ - if matches and not callable(matches): - pat, matches = matches, lambda obj: re.search(pat, str(obj), flags) - if recursive: - getter = self._get_children - nodes = chain.from_iterable(getter(n) for n in self.nodes) - else: - nodes = self.nodes - for node in nodes: - if not forcetype or isinstance(node, forcetype): - if not matches or matches(node): - yield node + return (node for i, node in + self._indexed_ifilter(recursive, matches, flags, forcetype)) def filter(self, recursive=True, matches=None, flags=FLAGS, forcetype=None): @@ -442,10 +468,10 @@ class Wikicode(StringMixIn): 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. + *matches* is given, it should be either a function or a regex; only + sections whose headings match it (without the surrounding equal signs) + will be included. *flags* can be used to override the default regex + flags (see :py:meth:`ifilter`) if a regex *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; @@ -454,35 +480,48 @@ class Wikicode(StringMixIn): :py:class:`~.Heading` object will be included; otherwise, this is skipped. """ - if matches: - matches = r"^(=+?)\s*" + matches + r"\s*\1$" - headings = self.filter_headings(recursive=False, matches=matches, - flags=flags) - if levels: - headings = [head for head in headings if head.level in levels] - - sections = [] + title_matcher = self._build_matcher(matches, flags) + matcher = lambda heading: (title_matcher(heading.title) and + (not levels or heading.level in levels)) + iheadings = self._indexed_ifilter(recursive=False, forcetype=Heading) + sections = [] # Tuples of (index_of_first_node, section) + open_headings = [] # Tuples of (index, heading), where index and + # heading.level are both monotonically increasing + + # Add the lead section if appropriate: if include_lead or not (include_lead is not None or matches or levels): - iterator = self.ifilter_headings(recursive=False) + itr = self._indexed_ifilter(recursive=False, forcetype=Heading) try: - first = self.index(next(iterator)) - sections.append(Wikicode(self.nodes[:first])) + first = next(itr)[0] + sections.append((0, Wikicode(self.nodes[:first]))) except StopIteration: # No headings in page - sections.append(Wikicode(self.nodes[:])) - - for heading in headings: - start = self.index(heading) - i = start + 1 - if not include_headings: - start += 1 - while i < len(self.nodes): - node = self.nodes[i] - if isinstance(node, Heading): - if flat or node.level <= heading.level: + sections.append((0, Wikicode(self.nodes[:]))) + + # Iterate over headings, adding sections to the list as they end: + for i, heading in iheadings: + if flat: # With flat, all sections close at the next heading + newly_closed, open_headings = open_headings, [] + else: # Otherwise, figure out which sections have closed, if any + closed_start_index = len(open_headings) + for j, (start, last_heading) in enumerate(open_headings): + if heading.level <= last_heading.level: + closed_start_index = j break - i += 1 - sections.append(Wikicode(self.nodes[start:i])) - return sections + newly_closed = open_headings[closed_start_index:] + del open_headings[closed_start_index:] + for start, closed_heading in newly_closed: + if matcher(closed_heading): + sections.append((start, Wikicode(self.nodes[start:i]))) + start = i if include_headings else (i + 1) + open_headings.append((start, heading)) + + # Add any remaining open headings to the list of sections: + for start, heading in open_headings: + if matcher(heading): + sections.append((start, Wikicode(self.nodes[start:]))) + + # Ensure that earlier sections are earlier in the returned list: + return [section for i, section in sorted(sections)] def strip_code(self, normalize=True, collapse=True): """Return a rendered string without unprintable code such as templates.