# 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 lists
import scalar
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: new = []
    else: new = lists.mk_seq(new)
    
    olds_parent = old.parentNode
    if not new: olds_parent.removeChild(old)
    else: # at least one node
        last = new.pop(-1)
        olds_parent.replaceChild(last, old) # note arg order reversed
        for node in reversed(new):
            olds_parent.insertBefore(node, last) # there is no insertAfter()

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 new != None and not isinstance(new, Node):
        if isinstance(new, bool): new = bool2str(new)
        else: new = strings.ustr(new)
        if new != None: 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 only_child(node):
    if not has_one_child(node):
        raise Exception('Must contain only one child:\n'+strings.ustr(node))
    return node.firstChild

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()
        values = list(NodeIter(entry))
        if not values: values = None
        elif len(values) == 1: values = values[0]
        return (entry.tagName, values)

##### 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: value_ = only_child(entry)
        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_parent(node):
    '''Removes node and then parent if it's empty'''
    parent = node.parentNode # save parent before detaching node
    remove(node)
    prune_empty(parent)

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(' '+strings.ustr(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