Python script to break large XML files

XML_breaker.py

import os
import sys
from xml.sax import parse
from xml.sax.saxutils import XMLGenerator

class CycleFile(object):

    def __init__(self, filename):
        self.basename, self.ext = os.path.splitext(filename)
        self.index = 0
        self.open_next_file()

    def open_next_file(self):
        self.index += 1
        self.file = open(self.name(), 'w')

    def name(self):
        return '%s%s%s' % (self.basename, self.index, self.ext)

    def cycle(self):
        self.file.close()
        self.open_next_file()

    def write(self, str):
        self.file.write(str)

    def close(self):
        self.file.close()

class XMLBreaker(XMLGenerator):
    
    def __init__(self, break_into=None, break_after=1000, out=None, *args, **kwargs):
        XMLGenerator.__init__(self, out, encoding='utf-8', *args, **kwargs)
        self.out_file = out
        self.break_into = break_into
        self.break_after = break_after
        self.context = []
        self.count = 0

    def startElement(self, name, attrs):
        XMLGenerator.startElement(self, name, attrs)
        self.context.append((name, attrs))

    def endElement(self, name):
        XMLGenerator.endElement(self, name)
        self.context.pop()
        if name == self.break_into:
            self.count += 1
            if self.count == self.break_after:
                self.count = 0
                for element in reversed(self.context):
                    self.out_file.write("\n")
                    XMLGenerator.endElement(self, element[0])
                self.out_file.cycle()
                XMLGenerator.startDocument(self)
                for element in self.context:
                    XMLGenerator.startElement(self, *element)

filename, break_into, break_after = sys.argv[1:]
parse(filename, XMLBreaker(break_into, int(break_after), out=CycleFile(filename)))

Thank you very much!

I think there is a bug in this script, in that it is not able to distinguish between elements named the same, that occur at different levels. If you feed it:

python38 my.xml name 25000

Where my.xml look like:

<root>
    <name>
        <cd>
            <name>
            </name>
        </cd>
       <cassette>
            <name>
             </name>
        </cassette>
    </name>
</root>

The script believes that there are three name elements, where you probably only want it to believe there is one. So it consequently splits the files in strange positions, and you'll probably end up with more records than you thought you had.

@positivity13
for those who are getting the following error ,

self.file.write(str)
TypeError: write() argument must be str, not bytes

change the write function into

    def write(self, str_text):
        self.file.write(str_text.decode("utf-8") if isinstance(str_text,bytes) else str_text)

Thank you so much. It works very well.

@positivity13 for those who are getting the following error ,

self.file.write(str)
TypeError: write() argument must be str, not bytes

change the write function into

    def write(self, str_text):
        self.file.write(str_text.decode("utf-8") if isinstance(str_text,bytes) else str_text)

Thanks!

Hello Nic,
A licence would be great. I would like to know whether or not I can reuse your code and what are the terms.
Cheers!
R

Thanks for sharing the script! One issue I got is that the last splitted file always failed to open, the error looked something like 'error on line {N} at column {C}: Premature end of data in tag debt line {N}'. Is there anywhere in the script that can be modified to resolve this issue?

Thanks for sharing this script! It works so well!