from lxml import etree
from app import formatting


def get_namespace_from_top(fn, key='xmlns'):
    ac, el = next(etree.iterparse(fn))
    return {'xmlns': el.nsmap[key]}


def generate_tags_multiple_files(input_files, tag, ignore_tags, ns=None):
    """
    Calls xmltag generator for multiple files.
    """
    # Deprecate?
    for fn in input_files:
        return generate_xmltags(fn, tag, ignore_tags, ns)


def generate_tags_multiple_files_strings(input_files, ns, tag, ignore_tags):
    """
    Creates stringified xml output of elements with certain tag.
    """
    for el in generate_tags_multiple_files(input_files, tag, ignore_tags, ns):
        yield formatting.string_and_clear(el, ns)


def generate_xmltags(fn, tag, ignore_tags, ns=None):
    """
    Base generator for percolator xml psm, peptide, protein output,
    as well as for mzML, mzIdentML.
    ignore_tags are the ones that are not cleared when met by parser.
    """
    if ns is None:
        xmlns = ''
    else:
        xmlns = '{%s}' % ns['xmlns']
    for ac, el in etree.iterparse(fn):
        if el.tag == '{0}{1}'.format(xmlns, tag):
            yield el
        elif el.tag in ['{0}{1}'.format(xmlns, x) for x in
                        ignore_tags]:
            formatting.clear_el(el)

from lxml import etree
import itertools
from app import formatting


def get_namespace_from_top(fn, key='xmlns'):
    ac, el = next(etree.iterparse(fn))
    return {'xmlns': el.nsmap[key]}


def generate_tags_multiple_files(input_files, tag, ignore_tags, ns=None):
    """
    Calls xmltag generator for multiple files.
    """
    return itertools.chain.from_iterable([generate_xmltags(
        fn, tag, ignore_tags, ns) for fn in input_files])


def generate_tags_multiple_files_strings(input_files, ns, tag, ignore_tags):
    """
    Creates stringified xml output of elements with certain tag.
    """
    for el in generate_tags_multiple_files(input_files, tag, ignore_tags, ns):
        yield formatting.string_and_clear(el, ns)


def generate_xmltags(fn, tag, ignore_tags, ns=None):
    """
    Base generator for percolator xml psm, peptide, protein output,
    as well as for mzML, mzIdentML.
    ignore_tags are the ones that are not cleared when met by parser.
    """
    if ns is None:
        xmlns = ''
    else:
        xmlns = '{%s}' % ns['xmlns']
    for ac, el in etree.iterparse(fn):
        if el.tag == '{0}{1}'.format(xmlns, tag):
            yield el
        elif el.tag in ['{0}{1}'.format(xmlns, x) for x in
                        ignore_tags]:
            formatting.clear_el(el)

Return chained iterators instead of only first of multiple iterators
Return chained iterators instead of only first of multiple iterators
