A Python parser for MediaWiki wikicode https://mwparserfromhell.readthedocs.io/
25개 이상의 토픽을 선택하실 수 없습니다. Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

476 lines
20 KiB

  1. # -*- coding: utf-8 -*-
  2. #
  3. # Copyright (C) 2012 Ben Kurtovic <ben.kurtovic@verizon.net>
  4. #
  5. # Permission is hereby granted, free of charge, to any person obtaining a copy
  6. # of this software and associated documentation files (the "Software"), to deal
  7. # in the Software without restriction, including without limitation the rights
  8. # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  9. # copies of the Software, and to permit persons to whom the Software is
  10. # furnished to do so, subject to the following conditions:
  11. #
  12. # The above copyright notice and this permission notice shall be included in
  13. # all copies or substantial portions of the Software.
  14. #
  15. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  16. # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  17. # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  18. # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  19. # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  20. # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  21. # SOFTWARE.
  22. from __future__ import unicode_literals
  23. import re
  24. from .compat import maxsize, str
  25. from .nodes import Heading, Node, Tag, Template, Text, Wikilink
  26. from .string_mixin import StringMixIn
  27. from .utils import parse_anything
  28. __all__ = ["Wikicode"]
  29. FLAGS = re.IGNORECASE | re.DOTALL | re.UNICODE
  30. class Wikicode(StringMixIn):
  31. """A ``Wikicode`` is a container for nodes that operates like a string.
  32. Additionally, it contains methods that can be used to extract data from or
  33. modify the nodes, implemented in an interface similar to a list. For
  34. example, :py:meth:`index` can get the index of a node in the list, and
  35. :py:meth:`insert` can add a new node at that index. The :py:meth:`filter()
  36. <ifilter>` series of functions is very useful for extracting and iterating
  37. over, for example, all of the templates in the object.
  38. """
  39. def __init__(self, nodes):
  40. super(Wikicode, self).__init__()
  41. self._nodes = nodes
  42. def __unicode__(self):
  43. return "".join([str(node) for node in self.nodes])
  44. def _get_children(self, node):
  45. """Iterate over all descendants of a given *node*, including itself.
  46. This is implemented by the ``__iternodes__()`` generator of ``Node``
  47. classes, which by default yields itself and nothing more.
  48. """
  49. for context, child in node.__iternodes__(self._get_all_nodes):
  50. yield child
  51. def _get_context(self, node, obj):
  52. """Return a ``Wikicode`` that contains *obj* in its descendants.
  53. The closest (shortest distance from *node*) suitable ``Wikicode`` will
  54. be returned, or ``None`` if the *obj* is the *node* itself.
  55. Raises ``ValueError`` if *obj* is not within *node*.
  56. """
  57. for context, child in node.__iternodes__(self._get_all_nodes):
  58. if child is obj:
  59. return context
  60. raise ValueError(obj)
  61. def _get_all_nodes(self, code):
  62. """Iterate over all of our descendant nodes.
  63. This is implemented by calling :py:meth:`_get_children` on every node
  64. in our node list (:py:attr:`self.nodes <nodes>`).
  65. """
  66. for node in code.nodes:
  67. for child in self._get_children(node):
  68. yield child
  69. def _is_equivalent(self, obj, node):
  70. """Return ``True`` if *obj* and *node* are equivalent, else ``False``.
  71. If *obj* is a ``Node``, the function will test whether they are the
  72. same object, otherwise it will compare them with ``==``.
  73. """
  74. if isinstance(obj, Node):
  75. if node is obj:
  76. return True
  77. else:
  78. if node == obj:
  79. return True
  80. return False
  81. def _contains(self, nodes, obj):
  82. """Return ``True`` if *obj* is inside of *nodes*, else ``False``.
  83. If *obj* is a ``Node``, we will only return ``True`` if *obj* is
  84. actually in the list (and not just a node that equals it). Otherwise,
  85. the test is simply ``obj in nodes``.
  86. """
  87. if isinstance(obj, Node):
  88. for node in nodes:
  89. if node is obj:
  90. return True
  91. return False
  92. return obj in nodes
  93. def _do_search(self, obj, recursive, callback, context, *args, **kwargs):
  94. """Look within *context* for *obj*, executing *callback* if found.
  95. If *recursive* is ``True``, we'll look within context and its
  96. descendants, otherwise we'll just execute callback. We raise
  97. :py:exc:`ValueError` if *obj* isn't in our node list or context. If
  98. found, *callback* is passed the context, the index of the node within
  99. the context, and whatever were passed as ``*args`` and ``**kwargs``.
  100. """
  101. if recursive:
  102. for i, node in enumerate(context.nodes):
  103. if self._is_equivalent(obj, node):
  104. return callback(context, i, *args, **kwargs)
  105. if self._contains(self._get_children(node), obj):
  106. context = self._get_context(node, obj)
  107. return self._do_search(obj, recursive, callback, context,
  108. *args, **kwargs)
  109. raise ValueError(obj)
  110. callback(context, self.index(obj, recursive=False), *args, **kwargs)
  111. def _get_tree(self, code, lines, marker, indent):
  112. """Build a tree to illustrate the way the Wikicode object was parsed.
  113. The method that builds the actual tree is ``__showtree__`` of ``Node``
  114. objects. *code* is the ``Wikicode`` object to build a tree for. *lines*
  115. is the list to append the tree to, which is returned at the end of the
  116. method. *marker* is some object to be used to indicate that the builder
  117. should continue on from the last line instead of starting a new one; it
  118. should be any object that can be tested for with ``is``. *indent* is
  119. the starting indentation.
  120. """
  121. def write(*args):
  122. """Write a new line following the proper indentation rules."""
  123. if lines and lines[-1] is marker: # Continue from the last line
  124. lines.pop() # Remove the marker
  125. last = lines.pop()
  126. lines.append(last + " ".join(args))
  127. else:
  128. lines.append(" " * 6 * indent + " ".join(args))
  129. get = lambda code: self._get_tree(code, lines, marker, indent + 1)
  130. mark = lambda: lines.append(marker)
  131. for node in code.nodes:
  132. node.__showtree__(write, get, mark)
  133. return lines
  134. @property
  135. def nodes(self):
  136. """A list of :py:class:`~.Node` objects.
  137. This is the internal data actually stored within a
  138. :py:class:`~.Wikicode` object.
  139. """
  140. return self._nodes
  141. @nodes.setter
  142. def nodes(self, value):
  143. self._nodes = value
  144. def get(self, index):
  145. """Return the *index*\ th node within the list of nodes."""
  146. return self.nodes[index]
  147. def set(self, index, value):
  148. """Set the ``Node`` at *index* to *value*.
  149. Raises :py:exc:`IndexError` if *index* is out of range, or
  150. :py:exc:`ValueError` if *value* cannot be coerced into one
  151. :py:class:`~.Node`. To insert multiple nodes at an index, use
  152. :py:meth:`get` with either :py:meth:`remove` and :py:meth:`insert` or
  153. :py:meth:`replace`.
  154. """
  155. nodes = parse_anything(value).nodes
  156. if len(nodes) > 1:
  157. raise ValueError("Cannot coerce multiple nodes into one index")
  158. if index >= len(self.nodes) or -1 * index > len(self.nodes):
  159. raise IndexError("List assignment index out of range")
  160. self.nodes.pop(index)
  161. if nodes:
  162. self.nodes[index] = nodes[0]
  163. def index(self, obj, recursive=False):
  164. """Return the index of *obj* in the list of nodes.
  165. Raises :py:exc:`ValueError` if *obj* is not found. If *recursive* is
  166. ``True``, we will look in all nodes of ours and their descendants, and
  167. return the index of our direct descendant node within *our* list of
  168. nodes. Otherwise, the lookup is done only on direct descendants.
  169. """
  170. if recursive:
  171. for i, node in enumerate(self.nodes):
  172. if self._contains(self._get_children(node), obj):
  173. return i
  174. raise ValueError(obj)
  175. for i, node in enumerate(self.nodes):
  176. if self._is_equivalent(obj, node):
  177. return i
  178. raise ValueError(obj)
  179. def insert(self, index, value):
  180. """Insert *value* at *index* in the list of nodes.
  181. *value* can be anything parasable by :py:func:`.parse_anything`, which
  182. includes strings or other :py:class:`~.Wikicode` or :py:class:`~.Node`
  183. objects.
  184. """
  185. nodes = parse_anything(value).nodes
  186. for node in reversed(nodes):
  187. self.nodes.insert(index, node)
  188. def insert_before(self, obj, value, recursive=True):
  189. """Insert *value* immediately before *obj* in the list of nodes.
  190. *obj* can be either a string or a :py:class:`~.Node`. *value* can be
  191. anything parasable by :py:func:`.parse_anything`. If *recursive* is
  192. ``True``, we will try to find *obj* within our child nodes even if it
  193. is not a direct descendant of this :py:class:`~.Wikicode` object. If
  194. *obj* is not in the node list, :py:exc:`ValueError` is raised.
  195. """
  196. callback = lambda self, i, value: self.insert(i, value)
  197. self._do_search(obj, recursive, callback, self, value)
  198. def insert_after(self, obj, value, recursive=True):
  199. """Insert *value* immediately after *obj* in the list of nodes.
  200. *obj* can be either a string or a :py:class:`~.Node`. *value* can be
  201. anything parasable by :py:func:`.parse_anything`. If *recursive* is
  202. ``True``, we will try to find *obj* within our child nodes even if it
  203. is not a direct descendant of this :py:class:`~.Wikicode` object. If
  204. *obj* is not in the node list, :py:exc:`ValueError` is raised.
  205. """
  206. callback = lambda self, i, value: self.insert(i + 1, value)
  207. self._do_search(obj, recursive, callback, self, value)
  208. def replace(self, obj, value, recursive=True):
  209. """Replace *obj* with *value* in the list of nodes.
  210. *obj* can be either a string or a :py:class:`~.Node`. *value* can be
  211. anything parasable by :py:func:`.parse_anything`. If *recursive* is
  212. ``True``, we will try to find *obj* within our child nodes even if it
  213. is not a direct descendant of this :py:class:`~.Wikicode` object. If
  214. *obj* is not in the node list, :py:exc:`ValueError` is raised.
  215. """
  216. def callback(self, i, value):
  217. self.nodes.pop(i)
  218. self.insert(i, value)
  219. self._do_search(obj, recursive, callback, self, value)
  220. def append(self, value):
  221. """Insert *value* at the end of the list of nodes.
  222. *value* can be anything parasable by :py:func:`.parse_anything`.
  223. """
  224. nodes = parse_anything(value).nodes
  225. for node in nodes:
  226. self.nodes.append(node)
  227. def remove(self, obj, recursive=True):
  228. """Remove *obj* from the list of nodes.
  229. *obj* can be either a string or a :py:class:`~.Node`. If *recursive* is
  230. ``True``, we will try to find *obj* within our child nodes even if it
  231. is not a direct descendant of this :py:class:`~.Wikicode` object. If
  232. *obj* is not in the node list, :py:exc:`ValueError` is raised.
  233. """
  234. callback = lambda self, i: self.nodes.pop(i)
  235. self._do_search(obj, recursive, callback, self)
  236. def ifilter(self, recursive=False, matches=None, flags=FLAGS,
  237. forcetype=None):
  238. """Iterate over nodes in our list matching certain conditions.
  239. If *recursive* is ``True``, we will iterate over our children and all
  240. descendants of our children, otherwise just our immediate children. If
  241. *matches* is given, we will only yield the nodes that match the given
  242. regular expression (with :py:func:`re.search`). The default flags used
  243. are :py:const:`re.IGNORECASE`, :py:const:`re.DOTALL`, and
  244. :py:const:`re.UNICODE`, but custom flags can be specified by passing
  245. *flags*. If *forcetype* is given, only nodes that are instances of this
  246. type are yielded.
  247. """
  248. if recursive:
  249. nodes = self._get_all_nodes(self)
  250. else:
  251. nodes = self.nodes
  252. for node in nodes:
  253. if not forcetype or isinstance(node, forcetype):
  254. if not matches or re.search(matches, str(node), flags):
  255. yield node
  256. def ifilter_links(self, recursive=False, matches=None, flags=FLAGS):
  257. """Iterate over wikilink nodes.
  258. This is equivalent to :py:meth:`ifilter` with *forcetype* set to
  259. :py:class:`~.Wikilink`.
  260. """
  261. return self.ifilter(recursive, matches, flags, forcetype=Wikilink)
  262. def ifilter_templates(self, recursive=False, matches=None, flags=FLAGS):
  263. """Iterate over template nodes.
  264. This is equivalent to :py:meth:`ifilter` with *forcetype* set to
  265. :py:class:`~.Template`.
  266. """
  267. return self.filter(recursive, matches, flags, forcetype=Template)
  268. def ifilter_text(self, recursive=False, matches=None, flags=FLAGS):
  269. """Iterate over text nodes.
  270. This is equivalent to :py:meth:`ifilter` with *forcetype* set to
  271. :py:class:`~.nodes.Text`.
  272. """
  273. return self.filter(recursive, matches, flags, forcetype=Text)
  274. def ifilter_tags(self, recursive=False, matches=None, flags=FLAGS):
  275. """Iterate over tag nodes.
  276. This is equivalent to :py:meth:`ifilter` with *forcetype* set to
  277. :py:class:`~.Tag`.
  278. """
  279. return self.ifilter(recursive, matches, flags, forcetype=Tag)
  280. def filter(self, recursive=False, matches=None, flags=FLAGS,
  281. forcetype=None):
  282. """Return a list of nodes within our list matching certain conditions.
  283. This is equivalent to calling :py:func:`list` on :py:meth:`ifilter`.
  284. """
  285. return list(self.ifilter(recursive, matches, flags, forcetype))
  286. def filter_links(self, recursive=False, matches=None, flags=FLAGS):
  287. """Return a list of wikilink nodes.
  288. This is equivalent to calling :py:func:`list` on
  289. :py:meth:`ifilter_links`.
  290. """
  291. return list(self.ifilter_links(recursive, matches, flags))
  292. def filter_templates(self, recursive=False, matches=None, flags=FLAGS):
  293. """Return a list of template nodes.
  294. This is equivalent to calling :py:func:`list` on
  295. :py:meth:`ifilter_templates`.
  296. """
  297. return list(self.ifilter_templates(recursive, matches, flags))
  298. def filter_text(self, recursive=False, matches=None, flags=FLAGS):
  299. """Return a list of text nodes.
  300. This is equivalent to calling :py:func:`list` on
  301. :py:meth:`ifilter_text`.
  302. """
  303. return list(self.ifilter_text(recursive, matches, flags))
  304. def filter_tags(self, recursive=False, matches=None, flags=FLAGS):
  305. """Return a list of tag nodes.
  306. This is equivalent to calling :py:func:`list` on
  307. :py:meth:`ifilter_tags`.
  308. """
  309. return list(self.ifilter_tags(recursive, matches, flags))
  310. def get_sections(self, flat=True, matches=None, levels=None, flags=FLAGS,
  311. include_headings=True):
  312. """Return a list of sections within the page.
  313. Sections are returned as :py:class:`~.Wikicode` objects with a shared
  314. node list (implemented using :py:class:`~.SmartList`) so that changes
  315. to sections are reflected in the parent Wikicode object.
  316. With *flat* as ``True``, each returned section contains all of its
  317. subsections within the :py:class:`~.Wikicode`; otherwise, the returned
  318. sections contain only the section up to the next heading, regardless of
  319. its size. If *matches* is given, it should be a regex to matched
  320. against the titles of section headings; only sections whose headings
  321. match the regex will be included. If *levels* is given, it should be a =
  322. list of integers; only sections whose heading levels are within the
  323. list will be returned. If *include_headings* is ``True``, the section's
  324. literal :py:class:`~.Heading` object will be included in returned
  325. :py:class:`~.Wikicode` objects; otherwise, this is skipped.
  326. """
  327. if matches:
  328. matches = r"^(=+?)\s*" + matches + r"\s*\1$"
  329. headings = self.filter(recursive=True, matches=matches, flags=flags,
  330. forcetype=Heading)
  331. if levels:
  332. headings = [head for head in headings if head.level in levels]
  333. sections = []
  334. buffers = [[maxsize, 0]]
  335. i = 0
  336. while i < len(self.nodes):
  337. if self.nodes[i] in headings:
  338. this = self.nodes[i].level
  339. for (level, start) in buffers:
  340. if not flat or this <= level:
  341. buffers.remove([level, start])
  342. sections.append(Wikicode(self.nodes[start:i]))
  343. buffers.append([this, i])
  344. if not include_headings:
  345. i += 1
  346. i += 1
  347. for (level, start) in buffers:
  348. if start != i:
  349. sections.append(Wikicode(self.nodes[start:i]))
  350. return sections
  351. def strip_code(self, normalize=True, collapse=True):
  352. """Return a rendered string without unprintable code such as templates.
  353. The way a node is stripped is handled by the
  354. :py:meth:`~.Node.__showtree__` method of :py:class:`~.Node` objects,
  355. which generally return a subset of their nodes or ``None``. For
  356. example, templates and tags are removed completely, links are stripped
  357. to just their display part, headings are stripped to just their title.
  358. If *normalize* is ``True``, various things may be done to strip code
  359. further, such as converting HTML entities like ``&Sigma;``, ``&#931;``,
  360. and ``&#x3a3;`` to ``Σ``. If *collapse* is ``True``, we will try to
  361. remove excess whitespace as well (three or more newlines are converted
  362. to two, for example).
  363. """
  364. nodes = []
  365. for node in self.nodes:
  366. stripped = node.__strip__(normalize, collapse)
  367. if stripped:
  368. nodes.append(str(stripped))
  369. if collapse:
  370. stripped = "".join(nodes).strip("\n")
  371. while "\n\n\n" in stripped:
  372. stripped = stripped.replace("\n\n\n", "\n\n")
  373. return stripped
  374. else:
  375. return "".join(nodes)
  376. def get_tree(self):
  377. """Return a hierarchical tree representation of the object.
  378. The representation is a string makes the most sense printed. It is
  379. built by calling :py:meth:`_get_tree` on the
  380. :py:class:`~.Wikicode` object and its children recursively. The end
  381. result may look something like the following::
  382. >>> text = "Lorem ipsum {{foo|bar|{{baz}}|spam=eggs}}"
  383. >>> print mwparserfromhell.parse(text).get_tree()
  384. Lorem ipsum
  385. {{
  386. foo
  387. | 1
  388. = bar
  389. | 2
  390. = {{
  391. baz
  392. }}
  393. | spam
  394. = eggs
  395. }}
  396. """
  397. marker = object() # Random object we can find with certainty in a list
  398. return "\n".join(self._get_tree(self, [], marker, 0))