A Python parser for MediaWiki wikicode https://mwparserfromhell.readthedocs.io/
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

wikicode.py 30 KiB

12 years ago
5 years ago
5 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679
  1. # Copyright (C) 2012-2020 Ben Kurtovic <ben.kurtovic@gmail.com>
  2. #
  3. # Permission is hereby granted, free of charge, to any person obtaining a copy
  4. # of this software and associated documentation files (the "Software"), to deal
  5. # in the Software without restriction, including without limitation the rights
  6. # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  7. # copies of the Software, and to permit persons to whom the Software is
  8. # furnished to do so, subject to the following conditions:
  9. #
  10. # The above copyright notice and this permission notice shall be included in
  11. # all copies or substantial portions of the Software.
  12. #
  13. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  14. # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  15. # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  16. # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  17. # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  18. # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  19. # SOFTWARE.
  20. import re
  21. from itertools import chain
  22. from .nodes import (Argument, Comment, ExternalLink, Heading, HTMLEntity,
  23. Node, Tag, Template, Text, Wikilink)
  24. from .smart_list.list_proxy import ListProxy
  25. from .string_mixin import StringMixIn
  26. from .utils import parse_anything
  27. __all__ = ["Wikicode"]
  28. FLAGS = re.IGNORECASE | re.DOTALL | re.UNICODE
  29. class Wikicode(StringMixIn):
  30. """A ``Wikicode`` is a container for nodes that operates like a string.
  31. Additionally, it contains methods that can be used to extract data from or
  32. modify the nodes, implemented in an interface similar to a list. For
  33. example, :meth:`index` can get the index of a node in the list, and
  34. :meth:`insert` can add a new node at that index. The :meth:`filter()
  35. <ifilter>` series of functions is very useful for extracting and iterating
  36. over, for example, all of the templates in the object.
  37. """
  38. RECURSE_OTHERS = 2
  39. def __init__(self, nodes):
  40. super().__init__()
  41. self._nodes = nodes
  42. def __str__(self):
  43. return "".join([str(node) for node in self.nodes])
  44. @staticmethod
  45. def _get_children(node, contexts=False, restrict=None, parent=None):
  46. """Iterate over all child :class:`.Node`\\ s of a given *node*."""
  47. yield (parent, node) if contexts else node
  48. if restrict and isinstance(node, restrict):
  49. return
  50. for code in node.__children__():
  51. for child in code.nodes:
  52. sub = Wikicode._get_children(child, contexts, restrict, code)
  53. yield from sub
  54. @staticmethod
  55. def _slice_replace(code, index, old, new):
  56. """Replace the string *old* with *new* across *index* in *code*."""
  57. nodes = [str(node) for node in code.get(index)]
  58. substring = "".join(nodes).replace(old, new)
  59. code.nodes[index] = parse_anything(substring).nodes
  60. @staticmethod
  61. def _build_matcher(matches, flags):
  62. """Helper for :meth:`_indexed_ifilter` and others.
  63. If *matches* is a function, return it. If it's a regex, return a
  64. wrapper around it that can be called with a node to do a search. If
  65. it's ``None``, return a function that always returns ``True``.
  66. """
  67. if matches:
  68. if callable(matches):
  69. return matches
  70. return lambda obj: re.search(matches, str(obj), flags)
  71. return lambda obj: True
  72. def _indexed_ifilter(self, recursive=True, matches=None, flags=FLAGS,
  73. forcetype=None):
  74. """Iterate over nodes and their corresponding indices in the node list.
  75. The arguments are interpreted as for :meth:`ifilter`. For each tuple
  76. ``(i, node)`` yielded by this method, ``self.index(node) == i``. Note
  77. that if *recursive* is ``True``, ``self.nodes[i]`` might not be the
  78. node itself, but will still contain it.
  79. """
  80. match = self._build_matcher(matches, flags)
  81. if recursive:
  82. restrict = forcetype if recursive == self.RECURSE_OTHERS else None
  83. def getter(i, node):
  84. for ch in self._get_children(node, restrict=restrict):
  85. yield (i, ch)
  86. inodes = chain(*(getter(i, n) for i, n in enumerate(self.nodes)))
  87. else:
  88. inodes = enumerate(self.nodes)
  89. for i, node in inodes:
  90. if (not forcetype or isinstance(node, forcetype)) and match(node):
  91. yield (i, node)
  92. def _is_child_wikicode(self, obj, recursive=True):
  93. """Return whether the given :class:`.Wikicode` is a descendant."""
  94. def deref(nodes):
  95. if isinstance(nodes, ListProxy):
  96. return nodes._parent # pylint: disable=protected-access
  97. return nodes
  98. target = deref(obj.nodes)
  99. if target is deref(self.nodes):
  100. return True
  101. if recursive:
  102. todo = [self]
  103. while todo:
  104. code = todo.pop()
  105. if target is deref(code.nodes):
  106. return True
  107. for node in code.nodes:
  108. todo += list(node.__children__())
  109. return False
  110. def _do_strong_search(self, obj, recursive=True):
  111. """Search for the specific element *obj* within the node list.
  112. *obj* can be either a :class:`.Node` or a :class:`.Wikicode` object. If
  113. found, we return a tuple (*context*, *index*) where *context* is the
  114. :class:`.Wikicode` that contains *obj* and *index* is its index there,
  115. as a :class:`slice`. Note that if *recursive* is ``False``, *context*
  116. will always be ``self`` (since we only look for *obj* among immediate
  117. descendants), but if *recursive* is ``True``, then it could be any
  118. :class:`.Wikicode` contained by a node within ``self``. If *obj* is not
  119. found, :exc:`ValueError` is raised.
  120. """
  121. if isinstance(obj, Wikicode):
  122. if not self._is_child_wikicode(obj, recursive):
  123. raise ValueError(obj)
  124. return obj, slice(0, len(obj.nodes))
  125. if isinstance(obj, Node):
  126. mkslice = lambda i: slice(i, i + 1)
  127. if not recursive:
  128. return self, mkslice(self.index(obj))
  129. for node in self.nodes:
  130. for context, child in self._get_children(node, contexts=True):
  131. if obj is child:
  132. if not context:
  133. context = self
  134. return context, mkslice(context.index(child))
  135. raise ValueError(obj)
  136. raise TypeError(obj)
  137. def _do_weak_search(self, obj, recursive):
  138. """Search for an element that looks like *obj* within the node list.
  139. This follows the same rules as :meth:`_do_strong_search` with some
  140. differences. *obj* is treated as a string that might represent any
  141. :class:`.Node`, :class:`.Wikicode`, or combination of the two present
  142. in the node list. Thus, matching is weak (using string comparisons)
  143. rather than strong (using ``is``). Because multiple nodes can match
  144. *obj*, the result is a list of tuples instead of just one (however,
  145. :exc:`ValueError` is still raised if nothing is found). Individual
  146. matches will never overlap.
  147. The tuples contain a new first element, *exact*, which is ``True`` if
  148. we were able to match *obj* exactly to one or more adjacent nodes, or
  149. ``False`` if we found *obj* inside a node or incompletely spanning
  150. multiple nodes.
  151. """
  152. obj = parse_anything(obj)
  153. if not obj or obj not in self:
  154. raise ValueError(obj)
  155. results = []
  156. contexts = [self]
  157. while contexts:
  158. context = contexts.pop()
  159. i = len(context.nodes) - 1
  160. while i >= 0:
  161. node = context.get(i)
  162. if obj.get(-1) == node:
  163. for j in range(-len(obj.nodes), -1):
  164. if obj.get(j) != context.get(i + j + 1):
  165. break
  166. else:
  167. i -= len(obj.nodes) - 1
  168. index = slice(i, i + len(obj.nodes))
  169. results.append((True, context, index))
  170. elif recursive and obj in node:
  171. contexts.extend(node.__children__())
  172. i -= 1
  173. if not results:
  174. if not recursive:
  175. raise ValueError(obj)
  176. results.append((False, self, slice(0, len(self.nodes))))
  177. return results
  178. def _get_tree(self, code, lines, marker, indent):
  179. """Build a tree to illustrate the way the Wikicode object was parsed.
  180. The method that builds the actual tree is ``__showtree__`` of ``Node``
  181. objects. *code* is the ``Wikicode`` object to build a tree for. *lines*
  182. is the list to append the tree to, which is returned at the end of the
  183. method. *marker* is some object to be used to indicate that the builder
  184. should continue on from the last line instead of starting a new one; it
  185. should be any object that can be tested for with ``is``. *indent* is
  186. the starting indentation.
  187. """
  188. def write(*args):
  189. """Write a new line following the proper indentation rules."""
  190. if lines and lines[-1] is marker: # Continue from the last line
  191. lines.pop() # Remove the marker
  192. last = lines.pop()
  193. lines.append(last + " ".join(args))
  194. else:
  195. lines.append(" " * 6 * indent + " ".join(args))
  196. get = lambda code: self._get_tree(code, lines, marker, indent + 1)
  197. mark = lambda: lines.append(marker)
  198. for node in code.nodes:
  199. node.__showtree__(write, get, mark)
  200. return lines
  201. @classmethod
  202. def _build_filter_methods(cls, **meths):
  203. """Given Node types, build the corresponding i?filter shortcuts.
  204. The should be given as keys storing the method's base name paired with
  205. values storing the corresponding :class:`.Node` type. For example, the
  206. dict may contain the pair ``("templates", Template)``, which will
  207. produce the methods :meth:`ifilter_templates` and
  208. :meth:`filter_templates`, which are shortcuts for
  209. :meth:`ifilter(forcetype=Template) <ifilter>` and
  210. :meth:`filter(forcetype=Template) <filter>`, respectively. These
  211. shortcuts are added to the class itself, with an appropriate docstring.
  212. """
  213. doc = """Iterate over {0}.
  214. This is equivalent to :meth:`{1}` with *forcetype* set to
  215. :class:`~{2.__module__}.{2.__name__}`.
  216. """
  217. make_ifilter = lambda ftype: (lambda self, *a, **kw:
  218. self.ifilter(forcetype=ftype, *a, **kw))
  219. make_filter = lambda ftype: (lambda self, *a, **kw:
  220. self.filter(forcetype=ftype, *a, **kw))
  221. for name, ftype in meths.items():
  222. ifilt = make_ifilter(ftype)
  223. filt = make_filter(ftype)
  224. ifilt.__doc__ = doc.format(name, "ifilter", ftype)
  225. filt.__doc__ = doc.format(name, "filter", ftype)
  226. setattr(cls, "ifilter_" + name, ifilt)
  227. setattr(cls, "filter_" + name, filt)
  228. @property
  229. def nodes(self):
  230. """A list of :class:`.Node` objects.
  231. This is the internal data actually stored within a :class:`.Wikicode`
  232. object.
  233. """
  234. return self._nodes
  235. @nodes.setter
  236. def nodes(self, value):
  237. if not isinstance(value, list):
  238. value = parse_anything(value).nodes
  239. self._nodes = value
  240. def get(self, index):
  241. """Return the *index*\\ th node within the list of nodes."""
  242. return self.nodes[index]
  243. def set(self, index, value):
  244. """Set the ``Node`` at *index* to *value*.
  245. Raises :exc:`IndexError` if *index* is out of range, or
  246. :exc:`ValueError` if *value* cannot be coerced into one :class:`.Node`.
  247. To insert multiple nodes at an index, use :meth:`get` with either
  248. :meth:`remove` and :meth:`insert` or :meth:`replace`.
  249. """
  250. nodes = parse_anything(value).nodes
  251. if len(nodes) > 1:
  252. raise ValueError("Cannot coerce multiple nodes into one index")
  253. if index >= len(self.nodes) or -1 * index > len(self.nodes):
  254. raise IndexError("List assignment index out of range")
  255. if nodes:
  256. self.nodes[index] = nodes[0]
  257. else:
  258. self.nodes.pop(index)
  259. def contains(self, obj):
  260. """Return whether this Wikicode object contains *obj*.
  261. If *obj* is a :class:`.Node` or :class:`.Wikicode` object, then we
  262. search for it exactly among all of our children, recursively.
  263. Otherwise, this method just uses :meth:`.__contains__` on the string.
  264. """
  265. if not isinstance(obj, (Node, Wikicode)):
  266. return obj in self
  267. try:
  268. self._do_strong_search(obj, recursive=True)
  269. except ValueError:
  270. return False
  271. return True
  272. def index(self, obj, recursive=False):
  273. """Return the index of *obj* in the list of nodes.
  274. Raises :exc:`ValueError` if *obj* is not found. If *recursive* is
  275. ``True``, we will look in all nodes of ours and their descendants, and
  276. return the index of our direct descendant node within *our* list of
  277. nodes. Otherwise, the lookup is done only on direct descendants.
  278. """
  279. strict = isinstance(obj, Node)
  280. equivalent = (lambda o, n: o is n) if strict else (lambda o, n: o == n)
  281. for i, node in enumerate(self.nodes):
  282. if recursive:
  283. for child in self._get_children(node):
  284. if equivalent(obj, child):
  285. return i
  286. elif equivalent(obj, node):
  287. return i
  288. raise ValueError(obj)
  289. def get_ancestors(self, obj):
  290. """Return a list of all ancestor nodes of the :class:`.Node` *obj*.
  291. The list is ordered from the most shallow ancestor (greatest great-
  292. grandparent) to the direct parent. The node itself is not included in
  293. the list. For example::
  294. >>> text = "{{a|{{b|{{c|{{d}}}}}}}}"
  295. >>> code = mwparserfromhell.parse(text)
  296. >>> node = code.filter_templates(matches=lambda n: n == "{{d}}")[0]
  297. >>> code.get_ancestors(node)
  298. ['{{a|{{b|{{c|{{d}}}}}}}}', '{{b|{{c|{{d}}}}}}', '{{c|{{d}}}}']
  299. Will return an empty list if *obj* is at the top level of this Wikicode
  300. object. Will raise :exc:`ValueError` if it wasn't found.
  301. """
  302. def _get_ancestors(code, needle):
  303. for node in code.nodes:
  304. if node is needle:
  305. return []
  306. for code in node.__children__():
  307. ancestors = _get_ancestors(code, needle)
  308. if ancestors is not None:
  309. return [node] + ancestors
  310. return None
  311. if isinstance(obj, Wikicode):
  312. obj = obj.get(0)
  313. elif not isinstance(obj, Node):
  314. raise ValueError(obj)
  315. ancestors = _get_ancestors(self, obj)
  316. if ancestors is None:
  317. raise ValueError(obj)
  318. return ancestors
  319. def get_parent(self, obj):
  320. """Return the direct parent node of the :class:`.Node` *obj*.
  321. This function is equivalent to calling :meth:`.get_ancestors` and
  322. taking the last element of the resulting list. Will return None if
  323. the node exists but does not have a parent; i.e., it is at the top
  324. level of the Wikicode object.
  325. """
  326. ancestors = self.get_ancestors(obj)
  327. return ancestors[-1] if ancestors else None
  328. def insert(self, index, value):
  329. """Insert *value* at *index* in the list of nodes.
  330. *value* can be anything parsable by :func:`.parse_anything`, which
  331. includes strings or other :class:`.Wikicode` or :class:`.Node` objects.
  332. """
  333. nodes = parse_anything(value).nodes
  334. for node in reversed(nodes):
  335. self.nodes.insert(index, node)
  336. def insert_before(self, obj, value, recursive=True):
  337. """Insert *value* immediately before *obj*.
  338. *obj* can be either a string, a :class:`.Node`, or another
  339. :class:`.Wikicode` object (as created by :meth:`get_sections`, for
  340. example). If *obj* is a string, we will operate on all instances of
  341. that string within the code, otherwise only on the specific instance
  342. given. *value* can be anything parsable by :func:`.parse_anything`. If
  343. *recursive* is ``True``, we will try to find *obj* within our child
  344. nodes even if it is not a direct descendant of this :class:`.Wikicode`
  345. object. If *obj* is not found, :exc:`ValueError` is raised.
  346. """
  347. if isinstance(obj, (Node, Wikicode)):
  348. context, index = self._do_strong_search(obj, recursive)
  349. context.insert(index.start, value)
  350. else:
  351. for exact, context, index in self._do_weak_search(obj, recursive):
  352. if exact:
  353. context.insert(index.start, value)
  354. else:
  355. obj = str(obj)
  356. self._slice_replace(context, index, obj, str(value) + obj)
  357. def insert_after(self, obj, value, recursive=True):
  358. """Insert *value* immediately after *obj*.
  359. *obj* can be either a string, a :class:`.Node`, or another
  360. :class:`.Wikicode` object (as created by :meth:`get_sections`, for
  361. example). If *obj* is a string, we will operate on all instances of
  362. that string within the code, otherwise only on the specific instance
  363. given. *value* can be anything parsable by :func:`.parse_anything`. If
  364. *recursive* is ``True``, we will try to find *obj* within our child
  365. nodes even if it is not a direct descendant of this :class:`.Wikicode`
  366. object. If *obj* is not found, :exc:`ValueError` is raised.
  367. """
  368. if isinstance(obj, (Node, Wikicode)):
  369. context, index = self._do_strong_search(obj, recursive)
  370. context.insert(index.stop, value)
  371. else:
  372. for exact, context, index in self._do_weak_search(obj, recursive):
  373. if exact:
  374. context.insert(index.stop, value)
  375. else:
  376. obj = str(obj)
  377. self._slice_replace(context, index, obj, obj + str(value))
  378. def replace(self, obj, value, recursive=True):
  379. """Replace *obj* with *value*.
  380. *obj* can be either a string, a :class:`.Node`, or another
  381. :class:`.Wikicode` object (as created by :meth:`get_sections`, for
  382. example). If *obj* is a string, we will operate on all instances of
  383. that string within the code, otherwise only on the specific instance
  384. given. *value* can be anything parsable by :func:`.parse_anything`.
  385. If *recursive* is ``True``, we will try to find *obj* within our child
  386. nodes even if it is not a direct descendant of this :class:`.Wikicode`
  387. object. If *obj* is not found, :exc:`ValueError` is raised.
  388. """
  389. if isinstance(obj, (Node, Wikicode)):
  390. context, index = self._do_strong_search(obj, recursive)
  391. for _ in range(index.start, index.stop):
  392. context.nodes.pop(index.start)
  393. context.insert(index.start, value)
  394. else:
  395. for exact, context, index in self._do_weak_search(obj, recursive):
  396. if exact:
  397. for _ in range(index.start, index.stop):
  398. context.nodes.pop(index.start)
  399. context.insert(index.start, value)
  400. else:
  401. self._slice_replace(context, index, str(obj), str(value))
  402. def append(self, value):
  403. """Insert *value* at the end of the list of nodes.
  404. *value* can be anything parsable by :func:`.parse_anything`.
  405. """
  406. nodes = parse_anything(value).nodes
  407. for node in nodes:
  408. self.nodes.append(node)
  409. def remove(self, obj, recursive=True):
  410. """Remove *obj* from the list of nodes.
  411. *obj* can be either a string, a :class:`.Node`, or another
  412. :class:`.Wikicode` object (as created by :meth:`get_sections`, for
  413. example). If *obj* is a string, we will operate on all instances of
  414. that string within the code, otherwise only on the specific instance
  415. given. If *recursive* is ``True``, we will try to find *obj* within our
  416. child nodes even if it is not a direct descendant of this
  417. :class:`.Wikicode` object. If *obj* is not found, :exc:`ValueError` is
  418. raised.
  419. """
  420. if isinstance(obj, (Node, Wikicode)):
  421. context, index = self._do_strong_search(obj, recursive)
  422. for _ in range(index.start, index.stop):
  423. context.nodes.pop(index.start)
  424. else:
  425. for exact, context, index in self._do_weak_search(obj, recursive):
  426. if exact:
  427. for _ in range(index.start, index.stop):
  428. context.nodes.pop(index.start)
  429. else:
  430. self._slice_replace(context, index, str(obj), "")
  431. def matches(self, other):
  432. """Do a loose equivalency test suitable for comparing page names.
  433. *other* can be any string-like object, including :class:`.Wikicode`, or
  434. an iterable of these. This operation is symmetric; both sides are
  435. adjusted. Specifically, whitespace and markup is stripped and the first
  436. letter's case is normalized. Typical usage is
  437. ``if template.name.matches("stub"): ...``.
  438. """
  439. normalize = lambda s: (s[0].upper() + s[1:]).replace("_", " ") if s else s
  440. this = normalize(self.strip_code().strip())
  441. if isinstance(other, (str, bytes, Wikicode, Node)):
  442. that = parse_anything(other).strip_code().strip()
  443. return this == normalize(that)
  444. for obj in other:
  445. that = parse_anything(obj).strip_code().strip()
  446. if this == normalize(that):
  447. return True
  448. return False
  449. def ifilter(self, recursive=True, matches=None, flags=FLAGS,
  450. forcetype=None):
  451. """Iterate over nodes in our list matching certain conditions.
  452. If *forcetype* is given, only nodes that are instances of this type (or
  453. tuple of types) are yielded. Setting *recursive* to ``True`` will
  454. iterate over all children and their descendants. ``RECURSE_OTHERS``
  455. will only iterate over children that are not the instances of
  456. *forcetype*. ``False`` will only iterate over immediate children.
  457. ``RECURSE_OTHERS`` can be used to iterate over all un-nested templates,
  458. even if they are inside of HTML tags, like so:
  459. >>> code = mwparserfromhell.parse("{{foo}}<b>{{foo|{{bar}}}}</b>")
  460. >>> code.filter_templates(code.RECURSE_OTHERS)
  461. ["{{foo}}", "{{foo|{{bar}}}}"]
  462. *matches* can be used to further restrict the nodes, either as a
  463. function (taking a single :class:`.Node` and returning a boolean) or a
  464. regular expression (matched against the node's string representation
  465. with :func:`re.search`). If *matches* is a regex, the flags passed to
  466. :func:`re.search` are :const:`re.IGNORECASE`, :const:`re.DOTALL`, and
  467. :const:`re.UNICODE`, but custom flags can be specified by passing
  468. *flags*.
  469. """
  470. gen = self._indexed_ifilter(recursive, matches, flags, forcetype)
  471. return (node for i, node in gen)
  472. def filter(self, *args, **kwargs):
  473. """Return a list of nodes within our list matching certain conditions.
  474. This is equivalent to calling :func:`list` on :meth:`ifilter`.
  475. """
  476. return list(self.ifilter(*args, **kwargs))
  477. def get_sections(self, levels=None, matches=None, flags=FLAGS, flat=False,
  478. include_lead=None, include_headings=True):
  479. """Return a list of sections within the page.
  480. Sections are returned as :class:`.Wikicode` objects with a shared node
  481. list (implemented using :class:`.SmartList`) so that changes to
  482. sections are reflected in the parent Wikicode object.
  483. Each section contains all of its subsections, unless *flat* is
  484. ``True``. If *levels* is given, it should be a iterable of integers;
  485. only sections whose heading levels are within it will be returned. If
  486. *matches* is given, it should be either a function or a regex; only
  487. sections whose headings match it (without the surrounding equal signs)
  488. will be included. *flags* can be used to override the default regex
  489. flags (see :meth:`ifilter`) if a regex *matches* is used.
  490. If *include_lead* is ``True``, the first, lead section (without a
  491. heading) will be included in the list; ``False`` will not include it;
  492. the default will include it only if no specific *levels* were given. If
  493. *include_headings* is ``True``, the section's beginning
  494. :class:`.Heading` object will be included; otherwise, this is skipped.
  495. """
  496. title_matcher = self._build_matcher(matches, flags)
  497. matcher = lambda heading: (title_matcher(heading.title) and
  498. (not levels or heading.level in levels))
  499. iheadings = self._indexed_ifilter(recursive=False, forcetype=Heading)
  500. sections = [] # Tuples of (index_of_first_node, section)
  501. open_headings = [] # Tuples of (index, heading), where index and
  502. # heading.level are both monotonically increasing
  503. # Add the lead section if appropriate:
  504. if include_lead or not (include_lead is not None or matches or levels):
  505. itr = self._indexed_ifilter(recursive=False, forcetype=Heading)
  506. try:
  507. first = next(itr)[0]
  508. sections.append((0, Wikicode(self.nodes[:first])))
  509. except StopIteration: # No headings in page
  510. sections.append((0, Wikicode(self.nodes[:])))
  511. # Iterate over headings, adding sections to the list as they end:
  512. for i, heading in iheadings:
  513. if flat: # With flat, all sections close at the next heading
  514. newly_closed, open_headings = open_headings, []
  515. else: # Otherwise, figure out which sections have closed, if any
  516. closed_start_index = len(open_headings)
  517. for j, (start, last_heading) in enumerate(open_headings):
  518. if heading.level <= last_heading.level:
  519. closed_start_index = j
  520. break
  521. newly_closed = open_headings[closed_start_index:]
  522. del open_headings[closed_start_index:]
  523. for start, closed_heading in newly_closed:
  524. if matcher(closed_heading):
  525. sections.append((start, Wikicode(self.nodes[start:i])))
  526. start = i if include_headings else (i + 1)
  527. open_headings.append((start, heading))
  528. # Add any remaining open headings to the list of sections:
  529. for start, heading in open_headings:
  530. if matcher(heading):
  531. sections.append((start, Wikicode(self.nodes[start:])))
  532. # Ensure that earlier sections are earlier in the returned list:
  533. return [section for i, section in sorted(sections)]
  534. def strip_code(self, normalize=True, collapse=True,
  535. keep_template_params=False):
  536. """Return a rendered string without unprintable code such as templates.
  537. The way a node is stripped is handled by the
  538. :meth:`~.Node.__strip__` method of :class:`.Node` objects, which
  539. generally return a subset of their nodes or ``None``. For example,
  540. templates and tags are removed completely, links are stripped to just
  541. their display part, headings are stripped to just their title.
  542. If *normalize* is ``True``, various things may be done to strip code
  543. further, such as converting HTML entities like ``&Sigma;``, ``&#931;``,
  544. and ``&#x3a3;`` to ``Σ``. If *collapse* is ``True``, we will try to
  545. remove excess whitespace as well (three or more newlines are converted
  546. to two, for example). If *keep_template_params* is ``True``, then
  547. template parameters will be preserved in the output (normally, they are
  548. removed completely).
  549. """
  550. kwargs = {
  551. "normalize": normalize,
  552. "collapse": collapse,
  553. "keep_template_params": keep_template_params
  554. }
  555. nodes = []
  556. for node in self.nodes:
  557. stripped = node.__strip__(**kwargs)
  558. if stripped:
  559. nodes.append(str(stripped))
  560. if collapse:
  561. stripped = "".join(nodes).strip("\n")
  562. while "\n\n\n" in stripped:
  563. stripped = stripped.replace("\n\n\n", "\n\n")
  564. return stripped
  565. return "".join(nodes)
  566. def get_tree(self):
  567. """Return a hierarchical tree representation of the object.
  568. The representation is a string makes the most sense printed. It is
  569. built by calling :meth:`_get_tree` on the :class:`.Wikicode` object and
  570. its children recursively. The end result may look something like the
  571. following::
  572. >>> text = "Lorem ipsum {{foo|bar|{{baz}}|spam=eggs}}"
  573. >>> print(mwparserfromhell.parse(text).get_tree())
  574. Lorem ipsum
  575. {{
  576. foo
  577. | 1
  578. = bar
  579. | 2
  580. = {{
  581. baz
  582. }}
  583. | spam
  584. = eggs
  585. }}
  586. """
  587. marker = object() # Random object we can find with certainty in a list
  588. return "\n".join(self._get_tree(self, [], marker, 0))
  589. Wikicode._build_filter_methods(
  590. arguments=Argument, comments=Comment, external_links=ExternalLink,
  591. headings=Heading, html_entities=HTMLEntity, tags=Tag, templates=Template,
  592. text=Text, wikilinks=Wikilink)