selector.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405
  1. """
  2. XPath selectors based on lxml
  3. """
  4. import sys
  5. import six
  6. from lxml import etree, html
  7. from .utils import flatten, iflatten, extract_regex, shorten
  8. from .csstranslator import HTMLTranslator, GenericTranslator
  9. class CannotRemoveElementWithoutRoot(Exception):
  10. pass
  11. class CannotRemoveElementWithoutParent(Exception):
  12. pass
  13. class SafeXMLParser(etree.XMLParser):
  14. def __init__(self, *args, **kwargs):
  15. kwargs.setdefault('resolve_entities', False)
  16. super(SafeXMLParser, self).__init__(*args, **kwargs)
  17. _ctgroup = {
  18. 'html': {'_parser': html.HTMLParser,
  19. '_csstranslator': HTMLTranslator(),
  20. '_tostring_method': 'html'},
  21. 'xml': {'_parser': SafeXMLParser,
  22. '_csstranslator': GenericTranslator(),
  23. '_tostring_method': 'xml'},
  24. }
  25. def _st(st):
  26. if st is None:
  27. return 'html'
  28. elif st in _ctgroup:
  29. return st
  30. else:
  31. raise ValueError('Invalid type: %s' % st)
  32. def create_root_node(text, parser_cls, base_url=None):
  33. """Create root node for text using given parser class.
  34. """
  35. body = text.strip().replace('\x00', '').encode('utf8') or b'<html/>'
  36. parser = parser_cls(recover=True, encoding='utf8')
  37. root = etree.fromstring(body, parser=parser, base_url=base_url)
  38. if root is None:
  39. root = etree.fromstring(b'<html/>', parser=parser, base_url=base_url)
  40. return root
  41. class SelectorList(list):
  42. """
  43. The :class:`SelectorList` class is a subclass of the builtin ``list``
  44. class, which provides a few additional methods.
  45. """
  46. # __getslice__ is deprecated but `list` builtin implements it only in Py2
  47. def __getslice__(self, i, j):
  48. o = super(SelectorList, self).__getslice__(i, j)
  49. return self.__class__(o)
  50. def __getitem__(self, pos):
  51. o = super(SelectorList, self).__getitem__(pos)
  52. return self.__class__(o) if isinstance(pos, slice) else o
  53. def __getstate__(self):
  54. raise TypeError("can't pickle SelectorList objects")
  55. def xpath(self, xpath, namespaces=None, **kwargs):
  56. """
  57. Call the ``.xpath()`` method for each element in this list and return
  58. their results flattened as another :class:`SelectorList`.
  59. ``query`` is the same argument as the one in :meth:`Selector.xpath`
  60. ``namespaces`` is an optional ``prefix: namespace-uri`` mapping (dict)
  61. for additional prefixes to those registered with ``register_namespace(prefix, uri)``.
  62. Contrary to ``register_namespace()``, these prefixes are not
  63. saved for future calls.
  64. Any additional named arguments can be used to pass values for XPath
  65. variables in the XPath expression, e.g.::
  66. selector.xpath('//a[href=$url]', url="http://www.example.com")
  67. """
  68. return self.__class__(flatten([x.xpath(xpath, namespaces=namespaces, **kwargs) for x in self]))
  69. def css(self, query):
  70. """
  71. Call the ``.css()`` method for each element in this list and return
  72. their results flattened as another :class:`SelectorList`.
  73. ``query`` is the same argument as the one in :meth:`Selector.css`
  74. """
  75. return self.__class__(flatten([x.css(query) for x in self]))
  76. def re(self, regex, replace_entities=True):
  77. """
  78. Call the ``.re()`` method for each element in this list and return
  79. their results flattened, as a list of unicode strings.
  80. By default, character entity references are replaced by their
  81. corresponding character (except for ``&amp;`` and ``&lt;``.
  82. Passing ``replace_entities`` as ``False`` switches off these
  83. replacements.
  84. """
  85. return flatten([x.re(regex, replace_entities=replace_entities) for x in self])
  86. def re_first(self, regex, default=None, replace_entities=True):
  87. """
  88. Call the ``.re()`` method for the first element in this list and
  89. return the result in an unicode string. If the list is empty or the
  90. regex doesn't match anything, return the default value (``None`` if
  91. the argument is not provided).
  92. By default, character entity references are replaced by their
  93. corresponding character (except for ``&amp;`` and ``&lt;``.
  94. Passing ``replace_entities`` as ``False`` switches off these
  95. replacements.
  96. """
  97. for el in iflatten(x.re(regex, replace_entities=replace_entities) for x in self):
  98. return el
  99. return default
  100. def getall(self):
  101. """
  102. Call the ``.get()`` method for each element is this list and return
  103. their results flattened, as a list of unicode strings.
  104. """
  105. return [x.get() for x in self]
  106. extract = getall
  107. def get(self, default=None):
  108. """
  109. Return the result of ``.get()`` for the first element in this list.
  110. If the list is empty, return the default value.
  111. """
  112. for x in self:
  113. return x.get()
  114. return default
  115. extract_first = get
  116. @property
  117. def attrib(self):
  118. """Return the attributes dictionary for the first element.
  119. If the list is empty, return an empty dict.
  120. """
  121. for x in self:
  122. return x.attrib
  123. return {}
  124. def remove(self):
  125. """
  126. Remove matched nodes from the parent for each element in this list.
  127. """
  128. for x in self:
  129. x.remove()
  130. class Selector(object):
  131. """
  132. :class:`Selector` allows you to select parts of an XML or HTML text using CSS
  133. or XPath expressions and extract data from it.
  134. ``text`` is a ``unicode`` object in Python 2 or a ``str`` object in Python 3
  135. ``type`` defines the selector type, it can be ``"html"``, ``"xml"`` or ``None`` (default).
  136. If ``type`` is ``None``, the selector defaults to ``"html"``.
  137. ``base_url`` allows setting a URL for the document. This is needed when looking up external entities with relative paths.
  138. See [`lxml` documentation](https://lxml.de/api/index.html) ``lxml.etree.fromstring`` for more information.
  139. """
  140. __slots__ = ['text', 'namespaces', 'type', '_expr', 'root',
  141. '__weakref__', '_parser', '_csstranslator', '_tostring_method']
  142. _default_type = None
  143. _default_namespaces = {
  144. "re": "http://exslt.org/regular-expressions",
  145. # supported in libxslt:
  146. # set:difference
  147. # set:has-same-node
  148. # set:intersection
  149. # set:leading
  150. # set:trailing
  151. "set": "http://exslt.org/sets"
  152. }
  153. _lxml_smart_strings = False
  154. selectorlist_cls = SelectorList
  155. def __init__(self, text=None, type=None, namespaces=None, root=None,
  156. base_url=None, _expr=None):
  157. self.type = st = _st(type or self._default_type)
  158. self._parser = _ctgroup[st]['_parser']
  159. self._csstranslator = _ctgroup[st]['_csstranslator']
  160. self._tostring_method = _ctgroup[st]['_tostring_method']
  161. if text is not None:
  162. if not isinstance(text, six.text_type):
  163. msg = "text argument should be of type %s, got %s" % (
  164. six.text_type, text.__class__)
  165. raise TypeError(msg)
  166. root = self._get_root(text, base_url)
  167. elif root is None:
  168. raise ValueError("Selector needs either text or root argument")
  169. self.namespaces = dict(self._default_namespaces)
  170. if namespaces is not None:
  171. self.namespaces.update(namespaces)
  172. self.root = root
  173. self._expr = _expr
  174. def __getstate__(self):
  175. raise TypeError("can't pickle Selector objects")
  176. def _get_root(self, text, base_url=None):
  177. return create_root_node(text, self._parser, base_url=base_url)
  178. def xpath(self, query, namespaces=None, **kwargs):
  179. """
  180. Find nodes matching the xpath ``query`` and return the result as a
  181. :class:`SelectorList` instance with all elements flattened. List
  182. elements implement :class:`Selector` interface too.
  183. ``query`` is a string containing the XPATH query to apply.
  184. ``namespaces`` is an optional ``prefix: namespace-uri`` mapping (dict)
  185. for additional prefixes to those registered with ``register_namespace(prefix, uri)``.
  186. Contrary to ``register_namespace()``, these prefixes are not
  187. saved for future calls.
  188. Any additional named arguments can be used to pass values for XPath
  189. variables in the XPath expression, e.g.::
  190. selector.xpath('//a[href=$url]', url="http://www.example.com")
  191. """
  192. try:
  193. xpathev = self.root.xpath
  194. except AttributeError:
  195. return self.selectorlist_cls([])
  196. nsp = dict(self.namespaces)
  197. if namespaces is not None:
  198. nsp.update(namespaces)
  199. try:
  200. result = xpathev(query, namespaces=nsp,
  201. smart_strings=self._lxml_smart_strings,
  202. **kwargs)
  203. except etree.XPathError as exc:
  204. msg = u"XPath error: %s in %s" % (exc, query)
  205. msg = msg if six.PY3 else msg.encode('unicode_escape')
  206. six.reraise(ValueError, ValueError(msg), sys.exc_info()[2])
  207. if type(result) is not list:
  208. result = [result]
  209. result = [self.__class__(root=x, _expr=query,
  210. namespaces=self.namespaces,
  211. type=self.type)
  212. for x in result]
  213. return self.selectorlist_cls(result)
  214. def css(self, query):
  215. """
  216. Apply the given CSS selector and return a :class:`SelectorList` instance.
  217. ``query`` is a string containing the CSS selector to apply.
  218. In the background, CSS queries are translated into XPath queries using
  219. `cssselect`_ library and run ``.xpath()`` method.
  220. .. _cssselect: https://pypi.python.org/pypi/cssselect/
  221. """
  222. return self.xpath(self._css2xpath(query))
  223. def _css2xpath(self, query):
  224. return self._csstranslator.css_to_xpath(query)
  225. def re(self, regex, replace_entities=True):
  226. """
  227. Apply the given regex and return a list of unicode strings with the
  228. matches.
  229. ``regex`` can be either a compiled regular expression or a string which
  230. will be compiled to a regular expression using ``re.compile(regex)``.
  231. By default, character entity references are replaced by their
  232. corresponding character (except for ``&amp;`` and ``&lt;``).
  233. Passing ``replace_entities`` as ``False`` switches off these
  234. replacements.
  235. """
  236. return extract_regex(regex, self.get(), replace_entities=replace_entities)
  237. def re_first(self, regex, default=None, replace_entities=True):
  238. """
  239. Apply the given regex and return the first unicode string which
  240. matches. If there is no match, return the default value (``None`` if
  241. the argument is not provided).
  242. By default, character entity references are replaced by their
  243. corresponding character (except for ``&amp;`` and ``&lt;``).
  244. Passing ``replace_entities`` as ``False`` switches off these
  245. replacements.
  246. """
  247. return next(iflatten(self.re(regex, replace_entities=replace_entities)), default)
  248. def get(self):
  249. """
  250. Serialize and return the matched nodes in a single unicode string.
  251. Percent encoded content is unquoted.
  252. """
  253. try:
  254. return etree.tostring(self.root,
  255. method=self._tostring_method,
  256. encoding='unicode',
  257. with_tail=False)
  258. except (AttributeError, TypeError):
  259. if self.root is True:
  260. return u'1'
  261. elif self.root is False:
  262. return u'0'
  263. else:
  264. return six.text_type(self.root)
  265. extract = get
  266. def getall(self):
  267. """
  268. Serialize and return the matched node in a 1-element list of unicode strings.
  269. """
  270. return [self.get()]
  271. def register_namespace(self, prefix, uri):
  272. """
  273. Register the given namespace to be used in this :class:`Selector`.
  274. Without registering namespaces you can't select or extract data from
  275. non-standard namespaces. See :ref:`selector-examples-xml`.
  276. """
  277. self.namespaces[prefix] = uri
  278. def remove_namespaces(self):
  279. """
  280. Remove all namespaces, allowing to traverse the document using
  281. namespace-less xpaths. See :ref:`removing-namespaces`.
  282. """
  283. for el in self.root.iter('*'):
  284. if el.tag.startswith('{'):
  285. el.tag = el.tag.split('}', 1)[1]
  286. # loop on element attributes also
  287. for an in el.attrib.keys():
  288. if an.startswith('{'):
  289. el.attrib[an.split('}', 1)[1]] = el.attrib.pop(an)
  290. # remove namespace declarations
  291. etree.cleanup_namespaces(self.root)
  292. def remove(self):
  293. """
  294. Remove matched nodes from the parent element.
  295. """
  296. try:
  297. parent = self.root.getparent()
  298. except AttributeError:
  299. # 'str' object has no attribute 'getparent'
  300. raise CannotRemoveElementWithoutRoot(
  301. "The node you're trying to remove has no root, "
  302. "are you trying to remove a pseudo-element? "
  303. "Try to use 'li' as a selector instead of 'li::text' or "
  304. "'//li' instead of '//li/text()', for example."
  305. )
  306. try:
  307. parent.remove(self.root)
  308. except AttributeError:
  309. # 'NoneType' object has no attribute 'remove'
  310. raise CannotRemoveElementWithoutParent(
  311. "The node you're trying to remove has no parent, "
  312. "are you trying to remove a root element?"
  313. )
  314. @property
  315. def attrib(self):
  316. """Return the attributes dictionary for underlying element.
  317. """
  318. return dict(self.root.attrib)
  319. def __bool__(self):
  320. """
  321. Return ``True`` if there is any real content selected or ``False``
  322. otherwise. In other words, the boolean value of a :class:`Selector` is
  323. given by the contents it selects.
  324. """
  325. return bool(self.get())
  326. __nonzero__ = __bool__
  327. def __str__(self):
  328. data = repr(shorten(self.get(), width=40))
  329. return "<%s xpath=%r data=%s>" % (type(self).__name__, self._expr, data)
  330. __repr__ = __str__