<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
from collections import defaultdict


def read_tab_file_handle_sorted(tab_file_handle, factor_index=0):
    """
    Parse a tab file (sorted by a column) and return a generator
    """
    previous_factor_id = ''
    factor_tab_list = list()
    # Reading tab file
    for line in tab_file_handle:
        l = line.strip()
        if l:
            tab = l.split()
            current_factor = tab[factor_index]
            # Yield the previous factor tab list
            if current_factor != previous_factor_id:
                if previous_factor_id:
                    yield factor_tab_list
                    factor_tab_list = list()
            factor_tab_list.append(tab)
            previous_factor_id = current_factor
    # Yield the last tab list
    yield factor_tab_list
    # Close tab file
    tab_file_handle.close()


if __name__ == '__main__':

    # Arguments parsing
    parser = argparse.ArgumentParser(description='')
    # -i / --input_tab
    parser.add_argument('-i', '--input_tab',
                        metavar='INBLAST',
                        type=argparse.FileType('r'),
                        default='-',
                        help='Input blast tab file. '
                             'Only best matches should be kept')
    # -o / --output_tab
    parser.add_argument('-o', '--output_tab',
                        metavar='OUTBLAST',
                        type=argparse.FileType('w'),
                        default='-',
                        help='Ouput filtered blast tab file')

    args = parser.parse_args()

    # Variables initialization
    tab_list_list = list()

    # Reading blast tab file
    for tab_list in read_tab_file_handle_sorted(args.input_tab, 0):
        query_id = tab_list[0][0]
        tab_list_list.append(sorted(tab_list, key=lambda x: x[1]))

    # Sort by specificity (decreasing match number) and then
    # by decreasing alignment length. So the first seen contigs
    # are the long specific ones
    tab_list_list.sort(key=lambda x: (len(x), -int(x[0][3])))

    #
    kept_references_ids_set = set()
    fillers_tab_list_list = list()
    dealt_contigs_num = 0
    dealt_contigs_id = set()

    # Count refs
    ref_count_dict = defaultdict(int)
    for tab_list in tab_list_list:
        for tab in tab_list:
            ref_count_dict[tab[1]] += 1

    # Init a buffer for the output file
    output_tab_buffer = list()

    # Start iterating
    while (len(tab_list_list)):

        # Deal with specific contigs
        tab_list_buffer = list()

        # Purely specific contigs
        current_conservation_count = len(tab_list_list[0])
        if current_conservation_count == 1:
            for tab_list in tab_list_list:
                if len(tab_list) == 1: # Purely specific contig
                    tab = tab_list[0]
                    kept_references_ids_set.add(tab[1]) # Add the matching ref id to the set
                    # Write the alignment
                    output_tab_buffer.append('\t'.join(tab))
                    dealt_contigs_id.add(tab[0])
                    dealt_contigs_num += 1
                else:
                    tab_list_buffer.append(tab_list)
        # Multiple-choices specific contigs
        else:
            # Get most specific contig so far
            selected_tab_list = tab_list_list[0]
            #~ print(selected_tab_list)
            tab_list_buffer = tab_list_list[1:]

            selected_tab_list.sort(key=lambda t: (-ref_count_dict[t[1]], t[1]))
            selected_tab = selected_tab_list[0]

            kept_references_ids_set.add(selected_tab[1]) # Add the matching ref id to the set
            # Write the alignment
            output_tab_buffer.append('\t'.join(selected_tab))
            dealt_contigs_id.add(selected_tab[0])
            dealt_contigs_num += 1

        tab_list_list = tab_list_buffer
        tab_list_buffer = list()

        # Uses already known fillers
        fillers_tab_list_buffer = list()
        for tab_list in fillers_tab_list_list:
            #
            references_ids_set = frozenset(x[1] for x in tab_list)
            references_intersection = references_ids_set & kept_references_ids_set
            #
            tab_buffer = list()
            for tab in tab_list:
                reference_id = tab[1]
                if reference_id in references_intersection:
                    output_tab_buffer.append('\t'.join(tab))
                    dealt_contigs_id.add(tab[0])
                else:
                    # Store the alignments not already used
                    tab_buffer.append(tab)
            if tab_buffer:
                fillers_tab_list_buffer.append(tab_buffer)
            else:
                dealt_contigs_num += 1
        fillers_tab_list_list = fillers_tab_list_buffer
        fillers_tab_list_buffer = list()

        # Deal with new fillers
        for tab_list in tab_list_list:
            #
            references_ids_set = frozenset(x[1] for x in tab_list)
            references_intersection = references_ids_set & kept_references_ids_set
            #
            if len(references_intersection):  # So fillers
                tab_buffer = list()
                for tab in tab_list:
                    reference_id = tab[1]
                    if reference_id in references_intersection:
                        output_tab_buffer.append('\t'.join(tab))
                        dealt_contigs_id.add(tab[0])
                    else:
                        # Store the alignments not already used
                        tab_buffer.append(tab)
                if tab_buffer:
                    fillers_tab_list_list.append(tab_buffer)
                else:
                    dealt_contigs_num += 1
            else: # Still specific
                tab_list_buffer.append(tab_list)
        tab_list_list = tab_list_buffer
        tab_list_buffer = list()

    # Write output file
    for line in output_tab_buffer:
        args.output_tab.write('{}\n'.format(line))<commit_msg>Corrected a bug. By trying to do the first iteration quick, we may miss some fillers contigs.<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
from collections import defaultdict


def read_tab_file_handle_sorted(tab_file_handle, factor_index=0):
    """
    Parse a tab file (sorted by a column) and return a generator
    """
    previous_factor_id = ''
    factor_tab_list = list()
    # Reading tab file
    for line in tab_file_handle:
        l = line.strip()
        if l:
            tab = l.split()
            current_factor = tab[factor_index]
            # Yield the previous factor tab list
            if current_factor != previous_factor_id:
                if previous_factor_id:
                    yield factor_tab_list
                    factor_tab_list = list()
            factor_tab_list.append(tab)
            previous_factor_id = current_factor
    # Yield the last tab list
    yield factor_tab_list
    # Close tab file
    tab_file_handle.close()


if __name__ == '__main__':

    # Arguments parsing
    parser = argparse.ArgumentParser(description='')
    # -i / --input_tab
    parser.add_argument('-i', '--input_tab',
                        metavar='INBLAST',
                        type=argparse.FileType('r'),
                        default='-',
                        help='Input blast tab file. '
                             'Only best matches should be kept')
    # -o / --output_tab
    parser.add_argument('-o', '--output_tab',
                        metavar='OUTBLAST',
                        type=argparse.FileType('w'),
                        default='-',
                        help='Ouput filtered blast tab file')

    args = parser.parse_args()

    # Variables initialization
    tab_list_list = list()

    # Reading blast tab file
    for tab_list in read_tab_file_handle_sorted(args.input_tab, 0):
        query_id = tab_list[0][0]
        tab_list_list.append(sorted(tab_list, key=lambda x: x[1]))

    # Sort by specificity (decreasing match number) and then
    # by decreasing alignment length. So the first seen contigs
    # are the long specific ones
    tab_list_list.sort(key=lambda x: (len(x), -int(x[0][3])))

    #
    kept_references_ids_set = set()
    fillers_tab_list_list = list()

    # Count refs
    ref_count_dict = defaultdict(int)
    for tab_list in tab_list_list:
        for tab in tab_list:
            ref_count_dict[tab[1]] += 1

    # Init a buffer for the output file
    output_tab_buffer = list()

    # Start iterating
    while (len(tab_list_list)):

        # Deal with specific contigs
        tab_list_buffer = list()

		# Get most specific contig so far
		selected_tab_list = tab_list_list[0]
		#~ print(selected_tab_list)
		tab_list_buffer = tab_list_list[1:]

		# Get the alignment against the reference with the most alignments
		selected_tab_list.sort(key=lambda t: (-ref_count_dict[t[1]], t[1]))
		selected_tab = selected_tab_list[0]

		kept_references_ids_set.add(selected_tab[1]) # Add the matching ref id to the set
		# Write the alignment
		output_tab_buffer.append('\t'.join(selected_tab))

        tab_list_list = tab_list_buffer
        tab_list_buffer = list()

        # Uses already known fillers
        fillers_tab_list_buffer = list()
        for tab_list in fillers_tab_list_list:
            #
            references_ids_set = frozenset(x[1] for x in tab_list)
            references_intersection = references_ids_set & kept_references_ids_set
            #
            tab_buffer = list()
            for tab in tab_list:
                reference_id = tab[1]
                if reference_id in references_intersection:
                    output_tab_buffer.append('\t'.join(tab))
                else:
                    # Store the alignments not already used
                    tab_buffer.append(tab)
            if tab_buffer:
                fillers_tab_list_buffer.append(tab_buffer)
        fillers_tab_list_list = fillers_tab_list_buffer
        fillers_tab_list_buffer = list()

        # Deal with new fillers
        for tab_list in tab_list_list:
            #
            references_ids_set = frozenset(x[1] for x in tab_list)
            references_intersection = references_ids_set & kept_references_ids_set
            #
            if len(references_intersection):  # So fillers
                tab_buffer = list()
                for tab in tab_list:
                    reference_id = tab[1]
                    if reference_id in references_intersection:
                        output_tab_buffer.append('\t'.join(tab))
                    else:
                        # Store the alignments not already used
                        tab_buffer.append(tab)
                if tab_buffer:
                    fillers_tab_list_list.append(tab_buffer)
            else: # Still specific
                tab_list_buffer.append(tab_list)
        tab_list_list = tab_list_buffer
        tab_list_buffer = list()

    # Write output file
    for line in output_tab_buffer:
        args.output_tab.write('{}\n'.format(line))
{'fraction_non_alphanumeric': 0.07835420393559929, 'fraction_numerical': 0.005008944543828265, 'mean_word_length': 1.8522959183673469, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3937238', 'n_tokens_mistral': 3287, 'n_tokens_neox': 3150, 'n_words': 750}