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)))

@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!