codekingpro/portable-devtools

1from typing import TYPE_CHECKING, Optional2 3from ...encoding import UTF84from ...plugins import FormatterPlugin5 6if TYPE_CHECKING:7    from xml.dom.minidom import Document8 9 10XML_DECLARATION_OPEN = '<?xml'11XML_DECLARATION_CLOSE = '?>'12 13 14def parse_xml(data: str) -> 'Document':15    """Parse given XML `data` string into an appropriate :class:`~xml.dom.minidom.Document` object."""16    from defusedxml.minidom import parseString17    return parseString(data)18 19 20def parse_declaration(raw_body: str) -> Optional[str]:21    body = raw_body.strip()22    # XMLDecl ::= '<?xml' DECL_CONTENT '?>'23    if body.startswith(XML_DECLARATION_OPEN):24        end = body.find(XML_DECLARATION_CLOSE)25        if end != -1:26            return body[:end + len(XML_DECLARATION_CLOSE)]27 28 29def pretty_xml(document: 'Document',30               declaration: Optional[str] = None,31               encoding: Optional[str] = UTF8,32               indent: int = 2) -> str:33    """Render the given :class:`~xml.dom.minidom.Document` `document` into a prettified string."""34    kwargs = {35        'encoding': encoding or UTF8,36        'indent': ' ' * indent,37    }38    body = document.toprettyxml(**kwargs).decode(kwargs['encoding'])39 40    # Remove blank lines automatically added by `toprettyxml()`.41    lines = [line for line in body.splitlines() if line.strip()]42 43    # xml.dom automatically adds the declaration, even if44    # it is not present in the actual body. Remove it.45    if len(lines) >= 1 and parse_declaration(lines[0]):46        lines.pop(0)47        if declaration:48            lines.insert(0, declaration)49 50    return '\n'.join(lines)51 52 53class XMLFormatter(FormatterPlugin):54 55    def __init__(self, **kwargs):56        super().__init__(**kwargs)57        self.enabled = self.format_options['xml']['format']58 59    def format_body(self, body: str, mime: str):60        if 'xml' not in mime:61            return body62 63        from xml.parsers.expat import ExpatError64        from defusedxml.common import DefusedXmlException65 66        declaration = parse_declaration(body)67        try:68            parsed_body = parse_xml(body)69        except ExpatError:70            pass  # Invalid XML, ignore.71        except DefusedXmlException:72            pass  # Unsafe XML, ignore.73        else:74            body = pretty_xml(parsed_body,75                              encoding=parsed_body.encoding,76                              indent=self.format_options['xml']['indent'],77                              declaration=declaration)78 79        return body80