Merge chunks into complete FASTA bins

    Merge bin chunks by appending consecutive chunks to one another.

    Parameters
    ----------
    fasta_file : file, str or pathlib.Path
        The FASTA file containing the chunks to merge.
    output_dir : str or pathlib.Path
        The output directory to write the merged FASTA bin into.

def merge_fasta(fasta_file, output_dir):

    """Merge chunks into complete FASTA bins

    Merge bin chunks by appending consecutive chunks to one another.

    Parameters
    ----------
    fasta_file : file, str or pathlib.Path
        The FASTA file containing the chunks to merge.
    output_dir : str or pathlib.Path
        The output directory to write the merged FASTA bin into.
    """

    #   First, define some functions for ordering chunks and detecting
    #   consecutive chunk sequences

    def chunk_lexicographic_order(chunk):
        """A quick callback to sort chunk ids lexicographically
        (first on original names alphabetically, then on relative
        position on the original contig)
        """
        chunk_fields = chunk.split("_")
        chunk_name = chunk_fields[:-1]
        chunk_id = chunk_fields[-1]
        return (chunk_name, int(chunk_id))

    def are_consecutive(chunk1, chunk2):
        if None in {chunk1, chunk2}:
            return False
        else:
            ord1 = chunk_lexicographic_order(chunk1)
            ord2 = chunk_lexicographic_order(chunk2)
            return (ord1[0] == ord2[0]) and (ord1[1] == ord2[1] + 1)

    def consecutiveness(key_chunk_pair):
        """A callback for the groupby magic below
        """
        key, chunk = key_chunk_pair
        chunk_name, chunk_id = chunk_lexicographic_order(chunk)
        return (chunk_name, chunk_id - key)

    #   Read chunks and sort them
    genome = {
        record.id: record.seq for record in SeqIO.parse(fasta_file, "fasta")
    }

    sorted_ids = sorted(genome, key=chunk_lexicographic_order)

    #   Identify consecutive ranges and merge them
    new_genome = dict()
    for _, g in itertools.groupby(enumerate(sorted_ids), consecutiveness):
        chunk_range = map(operator.itemgetter(1), g)
        first_chunk = next(chunk_range)
        my_sequence = genome[first_chunk]
        my_chunk = None
        while "Reading chunk range":
            try:
                my_chunk = next(chunk_range)
                my_sequence += genome[my_chunk]
            except StopIteration:
                break

        try:
            last_chunk_id = my_chunk.split("_")[-1]
        except AttributeError:
            last_chunk_id = ""

        if last_chunk_id:
            new_chunk_id = "{}_{}".format(first_chunk, last_chunk_id)
        else:
            new_chunk_id = first_chunk
        new_genome[new_chunk_id] = my_sequence

    #   Write the result
    base_name = ".".join(os.path.basename(fasta_file).split(".")[:-1])
    output_name = "{}_merged.fa".format(base_name)
    merged_core_file = os.path.join(output_dir, output_name)
    with open(merged_core_file, "w") as output_handle:
        for my_id in sorted(new_genome, key=chunk_lexicographic_order):
            output_handle.write(">{}\n".format(my_id))
            output_handle.write("{}\n".format(new_genome[my_id]))