class XMLGenerator(BaseGenerator):
def __init__(self, pretty_print: bool = False):
self.pretty_print = pretty_print
self.renderer = XMLRenderer()
self.stream = StringIO()
self.xml = SimplerXMLGenerator(self.stream, XMLRenderer.charset)
def start(self):
self.xml.startDocument()
self.xml.startElement(XMLRenderer.root_tag_name, {})
def add_list_item(self, data: dict):
self.xml.startElement(XMLRenderer.item_tag_name, {})
self.renderer._to_xml(self.xml, data)
self.xml.endElement(XMLRenderer.item_tag_name)
def finish(self):
self.xml.endElement(XMLRenderer.root_tag_name)
self.xml.endDocument()
def get_data(self):
if not self.pretty_print:
return self.stream.getvalue()
else:
tree = etree.XML(self.stream.getvalue().encode('utf-8'))
string = etree.tostring(tree,
encoding='utf-8',
pretty_print=True,
xml_declaration=True)
return string.decode('utf-8')
