Write ElementTree directly to zip with utf-8 encoding

Question

I want to modify a large number of XMLs. They are stored in ZIP-files. The source-XMLs are utf-8 encoded (at least to the guesses of the file tool on Linux) and have a correct XML declaration: .

The target ZIPs and the XMLs contained therein should also have the correct XML declaration. However, the (at least to me) most obvious method (using ElementTree.tostring) fails.

Here is a self-contained example, that should work out of the box. Short walkthrough:

imports
preparations (creating src.zip, these ZIPs are a given in my actual application)
actual work of program (modifying XMLs), starting at # read XMLs from zip

Please focus on the lower part, especially # APPROACH 1, APPROACH 2, APPROACH 3:

import os
import tempfile
import zipfile
from xml.etree.ElementTree import Element, parse

src_1 = os.path.join(tempfile.gettempdir(), "one.xml")
src_2 = os.path.join(tempfile.gettempdir(), "two.xml")
src_zip = os.path.join(tempfile.gettempdir(), "src.zip")
trgt_appr1_zip = os.path.join(tempfile.gettempdir(), "trgt_appr1.zip")
trgt_appr2_zip = os.path.join(tempfile.gettempdir(), "trgt_appr2.zip")
trgt_appr3_zip = os.path.join(tempfile.gettempdir(), "trgt_appr3.zip")

# file on hard disk that must be used due to ElementTree insufficiencies
tmp_xml_name = os.path.join(tempfile.gettempdir(), "curr_xml.tmp")

# prepare src.zip
tree1 = ElementTree(Element('hello', {'beer': 'good'}))
tree1.write(os.path.join(tempfile.gettempdir(), "one.xml"), encoding="UTF-8", xml_declaration=True)
tree2 = ElementTree(Element('scnd', {'äkey': 'a value'}))
tree2.write(os.path.join(tempfile.gettempdir(), "two.xml"), encoding="UTF-8", xml_declaration=True)

with zipfile.ZipFile(src_zip, 'a') as src:
    with open(src_1, 'r', encoding="utf-8") as one:
        string_representation = one.read()
    # write to zip
    src.writestr(zinfo_or_arcname="one.xml", data=string_representation.encode("utf-8"))
    with open(src_2, 'r', encoding="utf-8") as two:
        string_representation = two.read()
    # write to zip
    src.writestr(zinfo_or_arcname="two.xml", data=string_representation.encode("utf-8"))
os.remove(src_1)
os.remove(src_2)

# read XMLs from zip
with zipfile.ZipFile(src_zip, 'r') as zfile:

    updated_trees = []

    for xml_name in zfile.namelist():

        curr_file = zfile.open(xml_name, 'r')
        tree = parse(curr_file)
        # modify tree
        updated_tree = tree
        updated_tree.getroot().append(Element('new', {'newkey': 'new value'}))
        updated_trees.append((xml_name, updated_tree))

    for xml_name, updated_tree in updated_trees:

        # write to target file
        with zipfile.ZipFile(trgt_appr1_zip, 'a') as trgt1_zip, zipfile.ZipFile(trgt_appr2_zip, 'a') as trgt2_zip, zipfile.ZipFile(trgt_appr3_zip, 'a') as trgt3_zip:

            #
            # APPROACH 1 [DESIRED, BUT DOES NOT WORK]: write tree to zip-file
            # encoding in XML declaration missing
            #
            # create byte representation of elementtree
            byte_representation = tostring(element=updated_tree.getroot(), encoding='UTF-8', method='xml')
            # write XML directly to zip
            trgt1_zip.writestr(zinfo_or_arcname=xml_name, data=byte_representation)

            #
            # APPROACH 2 [WORKS IN THEORY, BUT DOES NOT WORK]: write tree to zip-file
            # encoding in XML declaration is faulty (is 'utf8', should be 'utf-8' or 'UTF-8')
            #
            # create byte representation of elementtree
            byte_representation = tostring(element=updated_tree.getroot(), encoding='utf8', method='xml')
            # write XML directly to zip
            trgt2_zip.writestr(zinfo_or_arcname=xml_name, data=byte_representation)

            #
            # APPROACH 3 [WORKS, BUT LACKS PERFORMANCE]: write to file, then read from file, then write to zip
            #
            # write to file
            updated_tree.write(tmp_xml_name, encoding="UTF-8", method="xml", xml_declaration=True)
            # read from file
            with open(tmp_xml_name, 'r', encoding="utf-8") as tmp:
                string_representation = tmp.read()
            # write to zip
            trgt3_zip.writestr(zinfo_or_arcname=xml_name, data=string_representation.encode("utf-8"))

    os.remove(tmp_xml_name)

APPROACH 3 works, but it is much more resource-intensive than the other two.

APPROACH 2 is the only way I could get an ElementTree object to be written with an actual XML declaration -- which then turns out to be invalid (utf8 instead of UTF-8/utf-8).

APPROACH 1 would be most desired -- but fails during reading later in the pipeline, as the XML declaration is missing.

Question: How can I get rid of writing the whole XML to disk first, only to read it afterwards, write it to the zip and delete it after being done with the zip? What am I missing?

user_1698431 · Answer

You can use an io.BytesIO object. This allows using ElementTree.write, while avoiding exporting the tree to disk:

import zipfile
from io import BytesIO
from xml.etree.ElementTree import ElementTree, Element

tree = ElementTree(Element('hello', {'beer': 'good'}))
bio = BytesIO()
tree.write(bio, encoding='UTF-8', xml_declaration=True)
with zipfile.ZipFile('/tmp/test.zip', 'w') as z:
    z.writestr('test.xml', bio.getvalue())

If you are using Python 3.6 or higher, there's an even shorter solution: you can get a writable file object from the ZipFile object, which you can pass to ElementTree.write:

import zipfile
from xml.etree.ElementTree import ElementTree, Element

tree = ElementTree(Element('hello', {'beer': 'good'}))
with zipfile.ZipFile('/tmp/test.zip', 'w') as z:
    with z.open('test.xml', 'w') as f:
        tree.write(f, encoding='UTF-8', xml_declaration=True)

This also has the advantage that you don't store multiple copies of the tree in memory, which could be a relevant issue for large trees.

user_6105043 · Answer

The only thing that is really missing in approach one is the XML declaration header. For ElementTree.write(...) you can use the xml_declaration, unfortunately for your version this isn't available in ElementTree.tostring yet.

Starting with Python 3.8, the ElementTree.tostring method does have a xml_declaration argument, see: https://docs.python.org/3.8/library/xml.etree.elementtree.html

Even though that implementation is unavailable to you when using Python 3.6, you can easily copy the 3.8 implementation in your own Python file:

import io

def tostring(element, encoding=None, method=None, *,
             xml_declaration=None, default_namespace=None,
             short_empty_elements=True):
    """Generate string representation of XML element.
    All subelements are included.  If encoding is "unicode", a string
    is returned. Otherwise a bytestring is returned.
    *element* is an Element instance, *encoding* is an optional output
    encoding defaulting to US-ASCII, *method* is an optional output which can
    be one of "xml" (default), "html", "text" or "c14n", *default_namespace*
    sets the default XML namespace (for "xmlns").
    Returns an (optionally) encoded string containing the XML data.
    """
    stream = io.StringIO() if encoding == 'unicode' else io.BytesIO()
    ElementTree(element).write(stream, encoding,
                               xml_declaration=xml_declaration,
                               default_namespace=default_namespace,
                               method=method,
                               short_empty_elements=short_empty_elements)
    return stream.getvalue()

(See https://github.com/python/cpython/blob/v3.8.0/Lib/xml/etree/ElementTree.py#L1116)

In that case you can simply use approach one:

# create byte representation of elementtree
byte_representation = tostring(element=updated_tree.getroot(), encoding='UTF-8', method='xml', xml_declaration=True)
# write XML directly to zip
trgt1_zip.writestr(zinfo_or_arcname=xml_name, data=byte_representation)

Write ElementTree directly to zip with utf-8 encoding

2 Answers