Read draw.io diagrams in Python and read/write text fields

drawioGettext.py

#!/usr/bin/python3

# Provides gettext-like functionality for drawio files:
# - Extract translatable strings from drawio diagram files (svg, png, xml) to PO files
# - Replace translatable strings in drawio diagram file with translations from a PO file
# - Regenerate the diagram in another language

# Reads .drawio.svg or .drawio.xml file, prints (or modifies) the `value` fields of the diagram, then writes a new file.
# Only changes the metadata, not the SVG. Use drawio to generate the SVG from the new metadata, e.g.:
# drawio -x -e -f svg -o regenerated.drawio.svg changed.drawio.svg

import xml.etree.ElementTree as ET

def getCompressedDiagramFromSVG(filename):
    """Get the diagram data from SVG metadata (it is compressed)."""
    svgtree = ET.parse(filename)
    if svgtree is None:
        print("Could not read SVG at " + filename + ". Filename incorrect?")
        return None, None, None
    svgroot = svgtree.getroot()
    mxstring = svgroot.attrib['content']
    mxelem = ET.fromstring(mxstring)
    for element in mxelem:
        if element.tag == 'diagram':
            b64string = element.text
            return b64string, mxelem, svgtree
    print(filename + " contains no drawio diagram!")
    return None, None, None

# TODO: Add PNG reading

def getCompressedDiagramFromXML(filename):
    """Get the diagram data from an XML or drawio file"""
    xmltree = ET.parse(filename)
    root = xmltree.getroot()
    # TODO: 
    mxelem = root[0]  # This can go wrong if there's more than one element. It needs to be the diagram element.
    b64string = mxelem.text
    return b64string, mxelem, xmltree

def decodeDiagram(b64string):
    """Decodes diagram from a base64-encoded and compressed string.
    
    Returns: 
        xml.etree.ElementTree
    """
    ## Decode diagram
    # 1) From base64 to encode bytestring
    # 2) Inflate (decompress with DEFLATE algorithm)
    # 3) Decode URI Encoding

    import base64
    string1 = base64.b64decode(b64string)

    import zlib
    # See https://stackoverflow.com/questions/46351275/using-pako-deflate-with-python
    # pako_inflate_raw
    c = zlib.decompressobj(-15)
    string2 = c.decompress(string1) + c.flush()

    import urllib.parse
    string3 = string2.decode('utf-8')
    graphstring = urllib.parse.unquote(string3)

    mxGraphTree = ET.fromstring(graphstring)
    return mxGraphTree


def changeDiagram(mxGraphTree):
    """Walks through value fields of a diagram.

    Returns:
        xml.etree.ElementTree: The changed diagram
    """

    ## Read text in diagram (and/or modify it)
    for cell in mxGraphTree[0]:
        if cell.attrib.get('value', None):
            pass
            # print(cell.attrib['value'])
            cell.attrib['value'] = "TEST"
            # TODO: Retrieve strings as translation string.

    return mxGraphTree


def encodeDiagram(mxGraphTree):
    """Encodes diagram into a base64-encoded and compressed string.
    
    Args:
        xml.etree.ElementTree: 

    Returns: 
        string
    """

    newgraphstring = ET.tostring(mxGraphTree, encoding='unicode')

    ## Re-encode diagram
    # Revert 3) Encode URI Encoding
    import urllib.parse
    newstring3 = urllib.parse.quote(newgraphstring, safe='~()*!.\'')
    newstring2 = newstring3.encode('utf-8')

    # Revert 2) To base64
    import zlib
    # See https://stackoverflow.com/questions/46351275/using-pako-deflate-with-python
    # pako_deflate_raw
    c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -15, memLevel=8, strategy=zlib.Z_DEFAULT_STRATEGY)
    newstring1 = c.compress(newstring2) + c.flush()

    # Revert 1) Deflate (decompress with DEFLATE algorithm)
    import base64
    newb64string = base64.b64encode(newstring1).decode('utf-8')

    return newb64string

def parseAndReturnDiagram(b64string):
    mxGraphTree = decodeDiagram(b64string)
    mxGraphTree = changeDiagram(mxGraphTree)
    return encodeDiagram(mxGraphTree)

def writeDiagramToXML(b64string, mxelem, xmltree, filename='changed.drawio.xml'):
    """Write base64-encoded and compressed diagram to XML"""
    mxelem.text = b64string
    xmltree.getroot()[0] = mxelem
    xmltree.write(filename)

def writeDiagramToSVG(b64string, mxelem, svgtree, filename='changed.drawio.xml'):
    """Write base64-encoded and compressed diagram into the SVG metadata"""
    ET.register_namespace("", "http://www.w3.org/2000/svg")  # Not perfect, but sufficient
    for element in mxelem:
        if element.tag == 'diagram':
            element.text = b64string
            break
    newmxstring = ET.tostring(mxelem, encoding='unicode')
    svgtree.getroot().attrib['content'] = newmxstring
    svgtree.write(filename)
    # The new file only has changed metadata, the SVG is still the same. Use drawio to generate the SVG from the new metadata, e.g.:
    # drawio -x -e -f svg -o regenerated.drawio.svg changed.drawio.svg

if __name__ == "__main__":
    b64string, mxelem, tree = getCompressedDiagramFromSVG("exampleEZ.drawio.svg")
    newb64string = parseAndReturnDiagram(b64string)
    writeDiagramToXML(newb64string, mxelem, tree, 'test.drawio.svg')

    # b64string, mxelem, tree = getCompressedDiagramFromXML("exampleEX.drawio.xml")
    # newb64string = parseAndReturnDiagram(b64string)
    # writeDiagramToXML(newb64string, mxelem, tree, 'test.drawio.xml')