1# Use of this source code is governed by the MIT license.
2__license__ = "MIT"
3
4__all__ = [
5    'LXMLTreeBuilderForXML',
6    'LXMLTreeBuilder',
7    ]
8
9try:
10    from collections.abc import Callable # Python 3.6
11except ImportError as e:
12    from collections import Callable
13
14from io import BytesIO
15from io import StringIO
16from lxml import etree
17from bs4.element import (
18    Comment,
19    Doctype,
20    NamespacedAttribute,
21    ProcessingInstruction,
22    XMLProcessingInstruction,
23)
24from bs4.builder import (
25    FAST,
26    HTML,
27    HTMLTreeBuilder,
28    PERMISSIVE,
29    ParserRejectedMarkup,
30    TreeBuilder,
31    XML)
32from bs4.dammit import EncodingDetector
33
34LXML = 'lxml'
35
36def _invert(d):
37    "Invert a dictionary."
38    return dict((v,k) for k, v in list(d.items()))
39
40class LXMLTreeBuilderForXML(TreeBuilder):
41    DEFAULT_PARSER_CLASS = etree.XMLParser
42
43    is_xml = True
44    processing_instruction_class = XMLProcessingInstruction
45
46    NAME = "lxml-xml"
47    ALTERNATE_NAMES = ["xml"]
48
49    # Well, it's permissive by XML parser standards.
50    features = [NAME, LXML, XML, FAST, PERMISSIVE]
51
52    CHUNK_SIZE = 512
53
54    # This namespace mapping is specified in the XML Namespace
55    # standard.
56    DEFAULT_NSMAPS = dict(xml='http://www.w3.org/XML/1998/namespace')
57
58    DEFAULT_NSMAPS_INVERTED = _invert(DEFAULT_NSMAPS)
59
60    # NOTE: If we parsed Element objects and looked at .sourceline,
61    # we'd be able to see the line numbers from the original document.
62    # But instead we build an XMLParser or HTMLParser object to serve
63    # as the target of parse messages, and those messages don't include
64    # line numbers.
65    # See: https://bugs.launchpad.net/lxml/+bug/1846906
66
67    def initialize_soup(self, soup):
68        """Let the BeautifulSoup object know about the standard namespace
69        mapping.
70
71        :param soup: A `BeautifulSoup`.
72        """
73        super(LXMLTreeBuilderForXML, self).initialize_soup(soup)
74        self._register_namespaces(self.DEFAULT_NSMAPS)
75
76    def _register_namespaces(self, mapping):
77        """Let the BeautifulSoup object know about namespaces encountered
78        while parsing the document.
79
80        This might be useful later on when creating CSS selectors.
81
82        :param mapping: A dictionary mapping namespace prefixes to URIs.
83        """
84        for key, value in list(mapping.items()):
85            if key and key not in self.soup._namespaces:
86                # Let the BeautifulSoup object know about a new namespace.
87                # If there are multiple namespaces defined with the same
88                # prefix, the first one in the document takes precedence.
89                self.soup._namespaces[key] = value
90
91    def default_parser(self, encoding):
92        """Find the default parser for the given encoding.
93
94        :param encoding: A string.
95        :return: Either a parser object or a class, which
96          will be instantiated with default arguments.
97        """
98        if self._default_parser is not None:
99            return self._default_parser
100        return etree.XMLParser(
101            target=self, strip_cdata=False, recover=True, encoding=encoding)
102
103    def parser_for(self, encoding):
104        """Instantiate an appropriate parser for the given encoding.
105
106        :param encoding: A string.
107        :return: A parser object such as an `etree.XMLParser`.
108        """
109        # Use the default parser.
110        parser = self.default_parser(encoding)
111
112        if isinstance(parser, Callable):
113            # Instantiate the parser with default arguments
114            parser = parser(
115                target=self, strip_cdata=False, recover=True, encoding=encoding
116            )
117        return parser
118
119    def __init__(self, parser=None, empty_element_tags=None, **kwargs):
120        # TODO: Issue a warning if parser is present but not a
121        # callable, since that means there's no way to create new
122        # parsers for different encodings.
123        self._default_parser = parser
124        if empty_element_tags is not None:
125            self.empty_element_tags = set(empty_element_tags)
126        self.soup = None
127        self.nsmaps = [self.DEFAULT_NSMAPS_INVERTED]
128        super(LXMLTreeBuilderForXML, self).__init__(**kwargs)
129
130    def _getNsTag(self, tag):
131        # Split the namespace URL out of a fully-qualified lxml tag
132        # name. Copied from lxml's src/lxml/sax.py.
133        if tag[0] == '{':
134            return tuple(tag[1:].split('}', 1))
135        else:
136            return (None, tag)
137
138    def prepare_markup(self, markup, user_specified_encoding=None,
139                       exclude_encodings=None,
140                       document_declared_encoding=None):
141        """Run any preliminary steps necessary to make incoming markup
142        acceptable to the parser.
143
144        lxml really wants to get a bytestring and convert it to
145        Unicode itself. So instead of using UnicodeDammit to convert
146        the bytestring to Unicode using different encodings, this
147        implementation uses EncodingDetector to iterate over the
148        encodings, and tell lxml to try to parse the document as each
149        one in turn.
150
151        :param markup: Some markup -- hopefully a bytestring.
152        :param user_specified_encoding: The user asked to try this encoding.
153        :param document_declared_encoding: The markup itself claims to be
154            in this encoding.
155        :param exclude_encodings: The user asked _not_ to try any of
156            these encodings.
157
158        :yield: A series of 4-tuples:
159         (markup, encoding, declared encoding,
160          has undergone character replacement)
161
162         Each 4-tuple represents a strategy for converting the
163         document to Unicode and parsing it. Each strategy will be tried
164         in turn.
165        """
166        is_html = not self.is_xml
167        if is_html:
168            self.processing_instruction_class = ProcessingInstruction
169        else:
170            self.processing_instruction_class = XMLProcessingInstruction
171
172        if isinstance(markup, str):
173            # We were given Unicode. Maybe lxml can parse Unicode on
174            # this system?
175            yield markup, None, document_declared_encoding, False
176
177        if isinstance(markup, str):
178            # No, apparently not. Convert the Unicode to UTF-8 and
179            # tell lxml to parse it as UTF-8.
180            yield (markup.encode("utf8"), "utf8",
181                   document_declared_encoding, False)
182
183        # This was provided by the end-user; treat it as a known
184        # definite encoding per the algorithm laid out in the HTML5
185        # spec.  (See the EncodingDetector class for details.)
186        known_definite_encodings = [user_specified_encoding]
187
188        # This was found in the document; treat it as a slightly lower-priority
189        # user encoding.
190        user_encodings = [document_declared_encoding]
191        detector = EncodingDetector(
192            markup, known_definite_encodings=known_definite_encodings,
193            user_encodings=user_encodings, is_html=is_html,
194            exclude_encodings=exclude_encodings
195        )
196        for encoding in detector.encodings:
197            yield (detector.markup, encoding, document_declared_encoding, False)
198
199    def feed(self, markup):
200        if isinstance(markup, bytes):
201            markup = BytesIO(markup)
202        elif isinstance(markup, str):
203            markup = StringIO(markup)
204
205        # Call feed() at least once, even if the markup is empty,
206        # or the parser won't be initialized.
207        data = markup.read(self.CHUNK_SIZE)
208        try:
209            self.parser = self.parser_for(self.soup.original_encoding)
210            self.parser.feed(data)
211            while len(data) != 0:
212                # Now call feed() on the rest of the data, chunk by chunk.
213                data = markup.read(self.CHUNK_SIZE)
214                if len(data) != 0:
215                    self.parser.feed(data)
216            self.parser.close()
217        except (UnicodeDecodeError, LookupError, etree.ParserError) as e:
218            raise ParserRejectedMarkup(e)
219
220    def close(self):
221        self.nsmaps = [self.DEFAULT_NSMAPS_INVERTED]
222
223    def start(self, name, attrs, nsmap={}):
224        # Make sure attrs is a mutable dict--lxml may send an immutable dictproxy.
225        attrs = dict(attrs)
226        nsprefix = None
227        # Invert each namespace map as it comes in.
228        if len(nsmap) == 0 and len(self.nsmaps) > 1:
229                # There are no new namespaces for this tag, but
230                # non-default namespaces are in play, so we need a
231                # separate tag stack to know when they end.
232                self.nsmaps.append(None)
233        elif len(nsmap) > 0:
234            # A new namespace mapping has come into play.
235
236            # First, Let the BeautifulSoup object know about it.
237            self._register_namespaces(nsmap)
238
239            # Then, add it to our running list of inverted namespace
240            # mappings.
241            self.nsmaps.append(_invert(nsmap))
242
243            # Also treat the namespace mapping as a set of attributes on the
244            # tag, so we can recreate it later.
245            attrs = attrs.copy()
246            for prefix, namespace in list(nsmap.items()):
247                attribute = NamespacedAttribute(
248                    "xmlns", prefix, "http://www.w3.org/2000/xmlns/")
249                attrs[attribute] = namespace
250
251        # Namespaces are in play. Find any attributes that came in
252        # from lxml with namespaces attached to their names, and
253        # turn then into NamespacedAttribute objects.
254        new_attrs = {}
255        for attr, value in list(attrs.items()):
256            namespace, attr = self._getNsTag(attr)
257            if namespace is None:
258                new_attrs[attr] = value
259            else:
260                nsprefix = self._prefix_for_namespace(namespace)
261                attr = NamespacedAttribute(nsprefix, attr, namespace)
262                new_attrs[attr] = value
263        attrs = new_attrs
264
265        namespace, name = self._getNsTag(name)
266        nsprefix = self._prefix_for_namespace(namespace)
267        self.soup.handle_starttag(name, namespace, nsprefix, attrs)
268
269    def _prefix_for_namespace(self, namespace):
270        """Find the currently active prefix for the given namespace."""
271        if namespace is None:
272            return None
273        for inverted_nsmap in reversed(self.nsmaps):
274            if inverted_nsmap is not None and namespace in inverted_nsmap:
275                return inverted_nsmap[namespace]
276        return None
277
278    def end(self, name):
279        self.soup.endData()
280        completed_tag = self.soup.tagStack[-1]
281        namespace, name = self._getNsTag(name)
282        nsprefix = None
283        if namespace is not None:
284            for inverted_nsmap in reversed(self.nsmaps):
285                if inverted_nsmap is not None and namespace in inverted_nsmap:
286                    nsprefix = inverted_nsmap[namespace]
287                    break
288        self.soup.handle_endtag(name, nsprefix)
289        if len(self.nsmaps) > 1:
290            # This tag, or one of its parents, introduced a namespace
291            # mapping, so pop it off the stack.
292            self.nsmaps.pop()
293
294    def pi(self, target, data):
295        self.soup.endData()
296        self.soup.handle_data(target + ' ' + data)
297        self.soup.endData(self.processing_instruction_class)
298
299    def data(self, content):
300        self.soup.handle_data(content)
301
302    def doctype(self, name, pubid, system):
303        self.soup.endData()
304        doctype = Doctype.for_name_and_ids(name, pubid, system)
305        self.soup.object_was_parsed(doctype)
306
307    def comment(self, content):
308        "Handle comments as Comment objects."
309        self.soup.endData()
310        self.soup.handle_data(content)
311        self.soup.endData(Comment)
312
313    def test_fragment_to_document(self, fragment):
314        """See `TreeBuilder`."""
315        return '<?xml version="1.0" encoding="utf-8"?>\n%s' % fragment
316
317
318class LXMLTreeBuilder(HTMLTreeBuilder, LXMLTreeBuilderForXML):
319
320    NAME = LXML
321    ALTERNATE_NAMES = ["lxml-html"]
322
323    features = ALTERNATE_NAMES + [NAME, HTML, FAST, PERMISSIVE]
324    is_xml = False
325    processing_instruction_class = ProcessingInstruction
326
327    def default_parser(self, encoding):
328        return etree.HTMLParser
329
330    def feed(self, markup):
331        encoding = self.soup.original_encoding
332        try:
333            self.parser = self.parser_for(encoding)
334            self.parser.feed(markup)
335            self.parser.close()
336        except (UnicodeDecodeError, LookupError, etree.ParserError) as e:
337            raise ParserRejectedMarkup(e)
338
339
340    def test_fragment_to_document(self, fragment):
341        """See `TreeBuilder`."""
342        return '<html><body>%s</body></html>' % fragment
343