# XML DOM tree manipulation

import cgi

from HTMLParser import HTMLParser

import itertools

import re

from xml.dom import Node

import xml.dom.minidom as minidom

import strings

import util

##### Escaping input

def escape(str_):

    return strings.to_unicode(cgi.escape(str_, True)).encode('ascii',

        'xmlcharrefreplace')

def unescape(str_): return HTMLParser().unescape(str_)

##### Names

def strip_namespace(name):

    namespace, sep, base = name.partition(':')

    if sep != '': return base

    else: return name

##### Nodes

def is_node(value): return isinstance(value, Node)

##### Replacing a node

def remove(node): node.parentNode.removeChild(node)

def replace(old, new):

    '''@param new Node|None'''

    assert old.parentNode != None # not removed from parent tree

    if new == None: old.parentNode.removeChild(old)

    else: old.parentNode.replaceChild(new, old) # note order reversed

def bool2str(val):

    '''For use with replace_with_text()'''

    if val: return '1'

    else: return None # remove node

def replace_with_text(node, new):

    '''

    @return The *new* node

    '''

    if isinstance(new, bool): new = bool2str(new)

    if util.is_str(new): new = node.ownerDocument.createTextNode(new)

    replace(node, new)

    return new

##### Element node contents

def is_elem(node): return node.nodeType == Node.ELEMENT_NODE

def is_completely_empty(node): return node.firstChild == None

def has_one_child(node):

    return node.firstChild != None and node.firstChild.nextSibling == None

def is_simple(node):

    '''Whether every child recursively has no more than one child'''

    return (not is_elem(node) or is_completely_empty(node)

        or (has_one_child(node) and is_simple(node.firstChild)))

class NodeIter:

    def __init__(self, node): self.child = node.firstChild

    def __iter__(self): return self

    def curr(self):

        if self.child != None: return self.child

        raise StopIteration

    def next(self):

        child = self.curr()

        self.child = self.child.nextSibling

        return child

##### Comments

def is_comment(node): return node.nodeType == Node.COMMENT_NODE

def is_empty(node):

    for child in NodeIter(node):

        if not (is_whitespace(child) or is_comment(child)): return False

    return True

def clean_comment(str_):

    '''Sanitizes comment node contents. Strips invalid strings.'''

    return re.sub(r'-{2,}', r'-', str_) # comments can't contain '--'

def mk_comment(doc, str_): return doc.createComment(clean_comment(str_))

##### Child nodes that are elements

class NodeElemIter:

    def __init__(self, node): self.child = node.firstChild

    def __iter__(self): return self

    def curr(self):

        while self.child != None:

            if is_elem(self.child): return self.child

            self.child = self.child.nextSibling

        raise StopIteration

    def next(self):

        child = self.curr()

        self.child = self.child.nextSibling

        return child

def first_elem(node): return NodeElemIter(node).next()

def has_elems(node):

    try: first_elem(node); return True

    except StopIteration: return False

class NodeElemReverseIter:

    def __init__(self, node): self.child = node.lastChild

    def __iter__(self): return self

    def curr(self):

        while self.child != None:

            if is_elem(self.child): return self.child

            self.child = self.child.previousSibling

        raise StopIteration

    def next(self):

        child = self.curr()

        self.child = self.child.previousSibling

        return child

def last_elem(node): return NodeElemReverseIter(node).next()

class NodeEntryIter:

    def __init__(self, node): self.iter_ = NodeElemIter(node)

    def __iter__(self): return self

    def next(self):

        entry = self.iter_.next()

        if is_empty(entry): value = None

        else:

            assert has_one_child(entry) # TODO: convert to an exception

            value = entry.firstChild

        return (entry.tagName, value)

##### Parent nodes

def parent(node):

    '''Does not treat the document object as the root node's parent, since it's

    not a true element node'''

    parent_ = node.parentNode

    if parent_ != None and is_elem(parent_): return parent_

    else: return None

class NodeParentIter:

    '''See parent() for special treatment of root node.

    Note that the first element returned is the current node, not its parent.'''

    def __init__(self, node): self.node = node

    def __iter__(self): return self

    def curr(self):

        if self.node != None: return self.node

        else: raise StopIteration

    def next(self):

        node = self.curr()

        self.node = parent(self.node)

        return node

##### Element nodes containing text

def is_text_node(node): return node.nodeType == Node.TEXT_NODE

def is_text(node): return has_one_child(node) and is_text_node(node.firstChild)

def value_node(node):

    if is_elem(node):

        iter_ = NodeIter(node)

        util.skip(iter_, is_comment)

        try: return iter_.next()

        except StopIteration: return None

    else: return node

def value(node):

    value_ = util.do_ignore_none(lambda n: n.nodeValue, value_node(node))

    if value_ == strings.isspace_none_str: value_ = None # None equiv

    return value_

def is_whitespace(node):

    return is_text_node(node) and (node.nodeValue == ''

        or node.nodeValue.isspace())

def set_value(node, value):

    value_node_ = value_node(node)

    if value != None:

        if value_node_ != None:

            value_node_.nodeValue = value

        else:

            assert is_elem(node)

            node.appendChild(node.ownerDocument.createTextNode(value))

    elif value_node_ != None:

        if is_elem(node): remove(value_node_)

        else:

            if is_text_node(node): value = strings.isspace_none_str # None equiv

            node.nodeValue = value

class NodeTextEntryIter:

    def __init__(self, node): self.iter_ = NodeElemIter(node)

    def __iter__(self): return self

    def next(self):

        entry = self.iter_.next()

        if is_empty(entry): value_ = None

        elif is_text(entry): value_ = value(entry)

        else:

            assert has_one_child(entry) # TODO: convert to an exception

            value_ = entry.firstChild

        return (entry.tagName, value_)

def is_text_node_entry(val): return util.is_str(val[1])

def non_empty(iterable):

    return itertools.ifilter(lambda i: i[1] != None, iterable)

class TextEntryOnlyIter(util.CheckedIter):

    def __init__(self, iterable):

        util.CheckedIter.__init__(self, is_text_node_entry, non_empty(iterable))

##### IDs

def get_id(node):

    '''If the node doesn't have an ID, assumes the node itself is the ID.

    @return None if the node doesn't have an ID or a value

    '''

    id_ = node.getAttribute('id')

    if id_ != '': return id_

    else: return value(node) # assume the node itself is the ID

def set_id(node, id_): node.setAttribute('id', id_)

##### Child nodes

def prune_empty(node):

    '''Removes node if it's empty'''

    if is_empty(node): remove(node)

def prune_children(node):

    '''Removes empty children'''

    for child in NodeElemIter(node): prune_empty(child)

def prune(node):

    '''Removes empty children and then node if it's empty'''

    prune_children(node)

    prune_empty(node)

def set_child(node, name, value):

    '''Note: does not remove any existing child of the same name'''

    child = node.ownerDocument.createElement(name)

    set_value(child, value)

    node.appendChild(child)

def by_tag_name(node, name, last_only=False, ignore_namespace=False):

    '''last_only optimization returns only the last matching node'''

    if ignore_namespace: filter_name = strip_namespace

    else: filter_name = lambda name: name

    name = filter_name(name)

    children = []

    if last_only: iter_ = NodeElemReverseIter(node)

    else: iter_ = NodeElemIter(node)

    for child in iter_:

        if filter_name(child.tagName) == name:

            children.append(child)

            if last_only: break

    return children

def merge(from_, into):

    '''Merges two nodes of the same tag name and their newly-adjacent children.

    @post The into node is saved; the from_ node is deleted.

    '''

    if from_ == None or into == None: return # base case

    if from_.tagName != into.tagName: return # not mergeable

    from_first = from_.firstChild # save before merge

    for child in NodeIter(from_): into.appendChild(child)

    remove(from_)

    # Recurse

    merge(from_first, from_first.previousSibling) # = into.lastChild

def merge_by_name(root, name):

    '''Merges siblings in root with the given name'''

    children = by_tag_name(root, name)

    child0 = children.pop(0)

    for child in children: merge(child, child0)

##### XML documents

def create_doc(root='_'):

    return minidom.getDOMImplementation().createDocument(None, root, None)

##### Printing XML

prettyxml_config = dict(addindent='    ', newl='\n')

toprettyxml_config = prettyxml_config.copy()

util.rename_key(toprettyxml_config, 'addindent', 'indent')

##### minidom modifications

#### Module

minidom._write_data = lambda writer, data: writer.write(escape(data))

minidom.Node.__iter__ = lambda self: NodeIter(self)

def __Node_str(self):

    return strings.remove_extra_newl(self.toprettyxml(**toprettyxml_config))

minidom.Node.__str__ = __Node_str

minidom.Node.__repr__ = __Node_str

minidom.Element.__repr__ = __Node_str

#### Node

minidom.Node.pop = lambda self: self.removeChild(self.lastChild)

def __Node_clear(self):

    while not is_empty(self): self.pop()

minidom.Node.clear = __Node_clear

#### Text

__Text_writexml_orig = minidom.Text.writexml

def __Text_writexml(self, *args, **kw_args):

    if is_whitespace(self): pass # we add our own whitespace

    else: __Text_writexml_orig(self, *args, **kw_args)

minidom.Text.writexml = __Text_writexml

#### Attr

def __Attr_str(self): return escape(self.name)+'="'+escape(self.value)+'"'

minidom.Attr.__str__ = __Attr_str

minidom.Attr.__repr__ = __Attr_str

#### Element

def __Element_write_opening(self, writer, indent='', addindent='', newl=''):

    writer.write(indent+'<'+escape(self.tagName))

    for attr_idx in xrange(self.attributes.length):

        writer.write(' '+str(self.attributes.item(attr_idx)))

    writer.write('>'+newl)

minidom.Element.write_opening = __Element_write_opening

def __Element_write_closing(self, writer, indent='', addindent='', newl=''):

    writer.write('</'+escape(self.tagName)+'>'+newl)

minidom.Element.write_closing = __Element_write_closing

__Element_writexml_orig = minidom.Element.writexml

def __Element_writexml(self, writer, indent='', addindent='', newl=''):

    if isinstance(indent, int): indent = addindent*indent

    if is_simple(self):

        writer.write(indent)

        __Element_writexml_orig(self, writer)

        writer.write(newl)

    else: __Element_writexml_orig(self, writer, indent, addindent, newl)

minidom.Element.writexml = __Element_writexml

#### Document

def __Document_write_opening(self, writer, indent='', addindent='', newl='',

    encoding=None):

    xmlDecl = '<?xml version="1.0" '

    if encoding != None: xmlDecl += 'encoding="'+escape(encoding)+'"'

    xmlDecl += '?>'+newl

    writer.write(xmlDecl)

    assert has_one_child(self)

    assert is_elem(self.firstChild)

    self.firstChild.write_opening(writer, indent, addindent, newl)

minidom.Document.write_opening = __Document_write_opening

def __Document_write_closing(self, writer, indent='', addindent='', newl=''):

    self.firstChild.write_closing(writer, indent, addindent, newl)

minidom.Document.write_closing = __Document_write_closing