_io.py

# Copyright (C) 2009 by Eric Talevich (eric.talevich@gmail.com)
# This code is part of the Biopython distribution and governed by its
# license. Please see the LICENSE file that should have been included
# as part of this package.

"""I/O function wrappers for phylogenetic tree formats.

This API follows the same semantics as Biopython's `SeqIO` and `AlignIO`.
"""

from __future__ import print_function

from Bio import File
from Bio.Phylo import (
    BaseTree,
    NewickIO,
    NexusIO,
    PhyloXMLIO,
    NeXMLIO,
)

supported_formats = {
    'newick': NewickIO,
    'nexus': NexusIO,
    'phyloxml': PhyloXMLIO,
    'nexml': NeXMLIO,
}

try:
    from Bio.Phylo import CDAOIO
    supported_formats['cdao'] = CDAOIO
except ImportError:
    pass


def parse(file, format, **kwargs):
    """Parse a file iteratively, and yield each of the trees it contains.

    If a file only contains one tree, this still returns an iterable object that
    contains one element.

    Examples
    --------
    >>> trees = parse('../../Tests/PhyloXML/apaf.xml', 'phyloxml')
    >>> for tree in trees:
    ...     print(tree.rooted)
    True

    """
    with File.as_handle(file, 'r') as fp:
        for tree in getattr(supported_formats[format], 'parse')(fp, **kwargs):
            yield tree


def read(file, format, **kwargs):
    """Parse a file in the given format and return a single tree.

    Raises a `ValueError` if there are zero or multiple trees -- if this occurs,
    use `parse` instead to get the complete sequence of trees.
    """
    try:
        tree_gen = parse(file, format, **kwargs)
        tree = next(tree_gen)
    except StopIteration:
        raise ValueError("There are no trees in this file.")
    try:
        next(tree_gen)
    except StopIteration:
        return tree
    else:
        raise ValueError(
            "There are multiple trees in this file; use parse() instead.")


def write(trees, file, format, **kwargs):
    """Write a sequence of trees to file in the given format."""
    if isinstance(trees, (BaseTree.Tree, BaseTree.Clade)):
        # Passed a single tree instead of an iterable -- that's OK
        trees = [trees]
    with File.as_handle(file, 'w+') as fp:
        n = getattr(supported_formats[format], 'write')(trees, fp, **kwargs)
    return n


def convert(in_file, in_format, out_file, out_format, parse_args=None, **kwargs):
    """Convert between two tree file formats."""
    if parse_args is None:
        parse_args = {}
    trees = parse(in_file, in_format, **parse_args)
    return write(trees, out_file, out_format, **kwargs)