<commit_before>"""
Classes that represent alignments between multiple sequences.
"""

import random
import string
from bx.misc.readlengths import read_lengths_file

# DNA reverse complement table
## DNA_COMP = "                                             -                  " \
##            " TVGH  CD  M KN   YSA BWXR       tvgh  cd  m kn   ysa bwxr      " \
##            "                                                                " \
##            "                                                                "

DNA_COMP = string.maketrans( "ACGTacgt", "TGCAtgca" )

class Alignment( object ):

    def __init__( self, score=0, attributes={}, species_to_lengths=None ):
        # species_to_lengths is needed only for file formats that don't provide
        # chromosome lengths;  it maps each species name to one of these:
        #   - the name of a file that contains a list of chromosome length pairs
        #   - a dict mapping chromosome names to their length
        #   - a single length value (useful when we just have one sequence and no chromosomes)
        # internally a file name is replaced by a dict, but only on an "as
        # needed" basis
        self.score = score
        self.text_size = 0
        self.attributes = attributes
        if species_to_lengths == None: self.species_to_lengths = {}
        else: self.species_to_lengths = species_to_lengths
        self.components = []

    def add_component( self, component ):
        component.alignment = self
        self.components.append( component )
        if component.text is not None:
            if self.text_size == 0: 
                self.text_size = len( component.text )
            elif self.text_size != len( component.text ): 
                raise Exception( "Components must have same text length" )

    def get_score( self ):
    	return self.__score
    def set_score( self,score ):
        if type( score ) == str:
            try:
                score = int(score)
            except:
                try:
                    score = float(score)
                except:
                    pass
        self.__score = score
    score = property( fget=get_score,fset=set_score )

    def __str__( self ):
        s = "a score=" + str( self.score )
        for key in self.attributes: 
            s += " %s=%s" % ( key, self.attributes[key] )
        s += "\n"
        # Components
        for c in self.components: 
            s += str( c )
            s += "\n"
        return s

    def src_size( self, src ):
        species,chrom = src_split( src )
        if species in self.species_to_lengths:
            chrom_to_length = self.species_to_lengths[species]
        elif chrom in self.species_to_lengths:
            chrom_to_length = self.species_to_lengths
        else:
            raise "no src_size (no length file for %s)" % species
        if type( chrom_to_length ) == int:         # (if it's a single length)
            return chrom_to_length
        if type( chrom_to_length ) == type( "" ):  # (if it's a file name)
            chrom_to_length = read_lengths_file( chrom_to_length )
            self.species_to_lengths[species] = chrom_to_length
        if chrom not in chrom_to_length: raise "no src_size (%s has no length for %s)" % ( species, chrom )
        return chrom_to_length[chrom]

    def get_component_by_src( self, src ):
        for c in self.components:
            if c.src == src: return c
        return None

    def get_component_by_src_start( self, src ):
        for c in self.components:
            if c.src.startswith( src ): return c
        return None    

    def slice( self, start, end ):
        new = Alignment( score=self.score, attributes=self.attributes )
        for component in self.components:
            # FIXME: Is this the right solution?
            if component.empty:
                continue
            new.components.append( component.slice( start, end ) )
        new.text_size = end - start
        return new
    
    def reverse_complement( self ):
        new = Alignment( score=self.score, attributes=self.attributes )
        for component in self.components:
            new.components.append( component.reverse_complement() )
        new.text_size = self.text_size
        return new
    
    def slice_by_component( self, component_index, start, end ):
        if type( component_index ) == type( 0 ):
            ref = self.components[ component_index ]
        elif type( component_index ) == type( "" ):
            ref = self.get_component_by_src( component_index )
        elif type( component_index ) == Component:
            ref = component_index
        else:
            raise ValueError( "can't figure out what to do" )
        start_col = ref.coord_to_col( start )  
        end_col = ref.coord_to_col( end )  
        return self.slice( start_col, end_col )
        
    def column_iter( self ):
        for i in range( self.text_size ):
            yield [ c.text[i] for c in self.components ]

    def limit_to_species( self, species ):
        new = Alignment( score=self.score, attributes=self.attributes )
        new.text_size = self.text_size
        for component in self.components:
            if component.src.split('.')[0] in species:
                new.add_component( component )
        return new

    def remove_all_gap_columns( self ):
        """
        Remove any columns containing only gaps from alignment components,
        text of components is modified IN PLACE.
        """
        try:
            seqs = [ list( c.text ) for c in self.components ]
        except TypeError:
            return
        i = 0
        text_size = self.text_size
        while i < text_size:
            all_gap = True
            for seq in seqs:
                if seq[i] != '-': all_gap = False
            if all_gap:
                for seq in seqs: del seq[i]
                text_size -= 1
            else:
                i += 1
        for i in range( len( self.components ) ):
            self.components[i].text = ''.join( seqs[i] )
        self.text_size = text_size
        
    def __eq__( self, other ):
        if other is None or type( other ) != type( self ):
            return False
        if self.score != other.score:
            return False
        if self.attributes != other.attributes:
            return False
        if len( self.components ) != len( other.components ):
            return False
        for c1, c2 in zip( self.components, other.components ):
            if c1 != c2:<commit_msg>Better (hopefully) fix for remove_all_gap_columns() where component.text is not of type sequence.<commit_after>"""
Classes that represent alignments between multiple sequences.
"""

import random
import string
from bx.misc.readlengths import read_lengths_file

# DNA reverse complement table
## DNA_COMP = "                                             -                  " \
##            " TVGH  CD  M KN   YSA BWXR       tvgh  cd  m kn   ysa bwxr      " \
##            "                                                                " \
##            "                                                                "

DNA_COMP = string.maketrans( "ACGTacgt", "TGCAtgca" )

class Alignment( object ):

    def __init__( self, score=0, attributes={}, species_to_lengths=None ):
        # species_to_lengths is needed only for file formats that don't provide
        # chromosome lengths;  it maps each species name to one of these:
        #   - the name of a file that contains a list of chromosome length pairs
        #   - a dict mapping chromosome names to their length
        #   - a single length value (useful when we just have one sequence and no chromosomes)
        # internally a file name is replaced by a dict, but only on an "as
        # needed" basis
        self.score = score
        self.text_size = 0
        self.attributes = attributes
        if species_to_lengths == None: self.species_to_lengths = {}
        else: self.species_to_lengths = species_to_lengths
        self.components = []

    def add_component( self, component ):
        component.alignment = self
        self.components.append( component )
        if component.text is not None:
            if self.text_size == 0: 
                self.text_size = len( component.text )
            elif self.text_size != len( component.text ): 
                raise Exception( "Components must have same text length" )

    def get_score( self ):
    	return self.__score
    def set_score( self,score ):
        if type( score ) == str:
            try:
                score = int(score)
            except:
                try:
                    score = float(score)
                except:
                    pass
        self.__score = score
    score = property( fget=get_score,fset=set_score )

    def __str__( self ):
        s = "a score=" + str( self.score )
        for key in self.attributes: 
            s += " %s=%s" % ( key, self.attributes[key] )
        s += "\n"
        # Components
        for c in self.components: 
            s += str( c )
            s += "\n"
        return s

    def src_size( self, src ):
        species,chrom = src_split( src )
        if species in self.species_to_lengths:
            chrom_to_length = self.species_to_lengths[species]
        elif chrom in self.species_to_lengths:
            chrom_to_length = self.species_to_lengths
        else:
            raise "no src_size (no length file for %s)" % species
        if type( chrom_to_length ) == int:         # (if it's a single length)
            return chrom_to_length
        if type( chrom_to_length ) == type( "" ):  # (if it's a file name)
            chrom_to_length = read_lengths_file( chrom_to_length )
            self.species_to_lengths[species] = chrom_to_length
        if chrom not in chrom_to_length: raise "no src_size (%s has no length for %s)" % ( species, chrom )
        return chrom_to_length[chrom]

    def get_component_by_src( self, src ):
        for c in self.components:
            if c.src == src: return c
        return None

    def get_component_by_src_start( self, src ):
        for c in self.components:
            if c.src.startswith( src ): return c
        return None    

    def slice( self, start, end ):
        new = Alignment( score=self.score, attributes=self.attributes )
        for component in self.components:
            # FIXME: Is this the right solution?
            if component.empty:
                continue
            new.components.append( component.slice( start, end ) )
        new.text_size = end - start
        return new
    
    def reverse_complement( self ):
        new = Alignment( score=self.score, attributes=self.attributes )
        for component in self.components:
            new.components.append( component.reverse_complement() )
        new.text_size = self.text_size
        return new
    
    def slice_by_component( self, component_index, start, end ):
        if type( component_index ) == type( 0 ):
            ref = self.components[ component_index ]
        elif type( component_index ) == type( "" ):
            ref = self.get_component_by_src( component_index )
        elif type( component_index ) == Component:
            ref = component_index
        else:
            raise ValueError( "can't figure out what to do" )
        start_col = ref.coord_to_col( start )  
        end_col = ref.coord_to_col( end )  
        return self.slice( start_col, end_col )
        
    def column_iter( self ):
        for i in range( self.text_size ):
            yield [ c.text[i] for c in self.components ]

    def limit_to_species( self, species ):
        new = Alignment( score=self.score, attributes=self.attributes )
        new.text_size = self.text_size
        for component in self.components:
            if component.src.split('.')[0] in species:
                new.add_component( component )
        return new

    def remove_all_gap_columns( self ):
        """
        Remove any columns containing only gaps from alignment components,
        text of components is modified IN PLACE.
        """
        seqs = []
        for c in self.components:
            try:
                seqs.append( list( c.text ) )
            except TypeError:
                seqs.append( None )
        i = 0
        text_size = self.text_size
        while i < text_size:
            all_gap = True
            for seq in seqs:
                if seq is None: continue
                if seq[i] != '-': all_gap = False
            if all_gap:
                for seq in seqs:
                    if seq is None: continue
                    del seq[i]
                text_size -= 1
            else:
                i += 1
        for i in range( len( self.components ) ):
            if seqs[i] is None: continue
            self.components[i].text = ''.join( seqs[i] )
        self.text_size = text_size
        
    def __eq__( self, other ):
        if other is None or type( other ) != type( self ):
            return False
        if self.score != other.score:
            return False
        if self.attributes != other.attributes:
            return False
        if len( self.components ) != len( other.components ):
            return False
        for c1, c2 in zip( self.components, other.components ):
            if c1 != c2:
{'fraction_non_alphanumeric': 0.08295344144682759, 'fraction_numerical': 0.0017935879231746506, 'mean_word_length': 1.6614956245027843, 'pattern_counts': {'":': 0, '<': 5, '<?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': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6918821', 'n_tokens_mistral': 3708, 'n_tokens_neox': 3563, 'n_words': 1155}