from cgi import escape
from re import compile
from exceptions import Exception

invalid_namespace_chars = compile('[^0-9A-Za-z_:.-]')

class XMLWriterError(Exception):
    pass

def _checkValidXMLNameSpace(name):
    """
    Verifies that the given name is a valid XML namespace identifier.
    Currently is not very thorough (":" may be present multiple times).
    """
    # TODO: This check should be more thorough.
    if len(invalid_namespace_chars.findall(name)) > 0:
        raise XMLWriterError, "%s is not a valid namespace identifier."\
            % name

class XMLWriter:
    """
    Class that implements an XML output interface that handles
    formatting and checks for valid XML. Currently does not understand
    character encoding.
    """

    def __init__(self, file, doctype = None, eol='\n', indent='  '):
        self.file = file
        self.eol = eol
        self.indent = indent
        self.tags = []
        self.level = 0
        # The needs_indent variable is used by the write function below.
        self.needs_indent = 0
        # TODO: Write doctype and allow for more flexible xml tag here.
        self.file.write('<?xml version="1.0" encoding="utf-7"?>')
        self.file.write(self.eol)

    def _makeAttributeString(self, attrs):
        """
        Makes a string out of the given attributes.
        """
        # Check for data validatity -- quotes (") need to be quoted.
        for key in attrs:
            _checkValidXMLNameSpace(key)
            attrs[key] = str(attrs[key]).replace('"', '&quot;')
        return ' '.join(['%s="%s"' % (key, attrs[key])
            for key in attrs.keys()])

    def openTag(self, tag, attrsdict = {}, **attrs):
        """
        Opens a tag. Attributes may be specified either as a
        dictionary (for names that are invalid Python names) or as
        keyword arguments (for convenience).
        """
        attrs.update(attrsdict)
        attrstr = self._makeAttributeString(attrs)
        _checkValidXMLNameSpace(tag)
        self.file.write('%s<%s' % (self.indent * self.level, tag))
        if attrstr != '':
            self.file.write(' %s' % attrstr)
        self.file.write('>%s' % self.eol)

        self.level += 1
        self.tags.append(tag)
        self.needs_indent = 1

    def writeTaglet(self, tag, attrsdict = {}, **attrs):
        """
        Writes a dataless tag. Attributes may be specified the same
        way as with openTag.
        """
        attrs.update(attrsdict)
        attrstr = self._makeAttributeString(attrs)
        # XXX: This function duplicates too much from the previous.
        _checkValidXMLNameSpace(tag)
        self.file.write('%s<%s' % (self.indent * self.level, tag))
        if attrstr != '':
            self.file.write(' %s' % attrstr)
        self.file.write(' />%s' % self.eol)
        self.needs_indent = 1

    def write(self, data):
        """
        Writes data. Currently doesn't check for newlines embedded in
        data.
        """
        if self.needs_indent:
            self.file.write(self.indent * self.level)
        self.file.write(escape(data))
        self.needs_indent = 0

    def writeln(self, line):
        """
        Writes a single line of data.
        """
        self.write(line)
        self.file.write(self.eol)
        self.needs_indent = 1

    def writelines(self, lines):
        """
        Write several lines of data.
        """
        for line in lines:
            self.writeln(line)

    def closeTag(self, tag = None):
        """
        Closes the given tag.
        """
        if tag is None:
            tag = self.tags[-1]
        else:
            if tag != self.tags[-1]:
                raise XMLWriterError, \
                    'Closing tag "%s" instead of "%s"' % (tag,
                    self.tags[-1])

        self.level -= 1
        self.tags.pop(-1)

        self.file.write('%s</%s>%s' % (self.indent * self.level, tag,
            self.eol))

    # Convenience shorthand functions:
    o = openTag
    t = writeTaglet
    w = writeln
    c = closeTag

if __name__=='__main__':
    # Test case
    from StringIO import StringIO
    f = StringIO()
    x = XMLWriter(f)
    x.openTag('root', {'one': '1', 'two': '2'}, three='3', four=4)
    x.write('root data')
    x.write(' and more data')
    x.writeln(' and some more data.')
    x.writeln('second line of data.')
    x.openTag('first')
    x.writeln('First internal level.')
    x.writeTaglet('second')
    x.writeTaglet('third', pone = '1')
    x.o('fourth', test=1, test2='2')
    x.w('Fourth level now.')
    x.c()
    x.closeTag()
    x.closeTag('root')

    print f.getvalue()
    f.close()