COMPRESSION PAR LA METHODE DE HUFFMAN

# -*- coding: cp1252 -*-
import Arbre, struct, sys

class HuffmanError(Exception):
    pass

class Huffman:
    """classe de compression/decompression par la methode d'Huffman
    """
    def __loop(self, arbre, node, carTable, cur):
        dirs=arbre.GetDir(node)
        for i in range(len(dirs)):
            if(dirs[i].getType()=='feuille'):
                carTable[dirs[i].getValue()]=''.join(cur)+str(i)
            else:
                newcur=[k for k in cur] #pour que ça se chamboule âs ds la recursivite
                newcur.append(str(i))
                carTable=self.__loop(arbre, dirs[i], carTable, newcur)
        return carTable
            
    def __makeArbreFromTable(self, table, tot):
        n=0
        while 1:
            #on prend les deux plus petites occurences
            oc=[(tot, None, None), (tot, None, None)]
            for key, dat in table.items():
                data=dat[0]
                for i in range(2):
                    if(data<=oc[i][0]):
                        if(i==0) :
                            ancien=oc[0]
                            oc[1]=ancien
                            oc[0]=(dat[0], dat[1], key)
                            break
                        else:
                            oc[1]=(dat[0], dat[1], key)
                            break
            node=Arbre.Node(1, oc[0][1], oc[1][1])
            del(table[oc[0][2]])
            del(table[oc[1][2]])
            table['node'+str(n)]=[oc[0][0]+oc[1][0], node]
            if(len(table)==1) : break
            n+=1
        a=Arbre.Arbre()
        a.SetRacine(table['node'+str(n)][1])
        
        #parcourt de larbre pour noter le chemin pour chaque car
        chem=self.__loop(a, table['node'+str(n)][1], {}, [])
        
        return a, chem

    def getBits(self, chem, car):
        return chem[car]

    def getOctet(self, bits, full=0):
        """si full a 1 on rempli la fin par des zero"""
        if(len(bits)!=8 and full==1) : bits+='0'*(8-len(bits))
        m=128
        octet=0
        for bit in bits:
            octet+=int(bit)*m
            m=m/2
        return struct.pack('B', octet)
                        
    def compress(self, f1, f2):
        """f1 doit avoir la methode read et f2 la methode write"""
        #creation de la table
        table={}
        tot=0.0
        while 1:
            buf=f1.read(1024*256)
            if not buf :
                table['eof']=[1, 'eof']
                tot+=1
                break
            else:
                for car in buf:
                    if table.has_key(car)==True : table[car][0]+=1
                    else : table[car]=[1, car]
                    tot+=1
        
        #creation de larbre et du chemin des car
        arbre, chem=self.__makeArbreFromTable(table.copy(), tot)

        #fabrication du fichier, header:
        header=''
        header+=struct.pack('c', 'H') #identification du fichier
        header+=struct.pack('c', 'U')
        header+=struct.pack('L', len(table)-1) #elements de la table (sans le eof)
        header+=struct.pack('L', tot) #total cars
        header+=struct.pack('L', 14+len(table)*5-5) #offset data
        f2.write(header)
        
        #ecriture de la table sans le eof qui est tjrs a 1 de tte facon
        Table=''
        for key, data in table.items():
            if(key!='eof'):
                Table+=struct.pack('c', key)
                Table+=struct.pack('L', data[0])
        f2.write(Table)
        
        #ecriture du fichier compresse
        f1.seek(0)
        curbit=''
        outbuf=''
        while 1:
            buf=f1.read(1024*16)
            if not buf :
                bits=self.getBits(chem, 'eof')
                if(len(curbit+bits)<=8):
                    f2.write(self.getOctet(curbit+bits, 1))
                elif(len(curbit+bits)>8):
                    towrite=curbit+bits
                    for i in range(0, (len(towrite)/8)+1):
                        if(i==len(towrite)/8) : f2.write(self.getOctet(towrite[i*8:], 1))
                        else: f2.write(self.getOctet(towrite[i*8:i*8+8]))
                break
            else:
                for c in buf:
                    bits=self.getBits(chem, c)
                    if(len(curbit+bits)<=8):
                        curbit=curbit+bits
                    elif(len(curbit+bits)>8):
                        towrite=curbit+bits
                        for i in range(0, (len(towrite)/8)+1):
                            if(i==len(towrite)/8) : curbit=towrite[i*8:]
                            else: outbuf+=self.getOctet(towrite[i*8:i*8+8])
                f2.write(outbuf)
                outbuf=''
        f1.close()
        f2.close()
    def decompress(self, f1, f2):
        #dabord lecture des infos:
        if(f1.read(2)!='HU') : raise HuffmanError, 'This file has not been compressed by me'
        elem=struct.unpack('L', f1.read(4))[0]
        totcars=struct.unpack('L', f1.read(4))[0]
        offset=struct.unpack('L', f1.read(4))[0]
        
        ###lecture de la table
        table={}
        for i in range(elem):
            el=f1.read(5)
            key=struct.unpack('c', el[0])[0]
            table[key]=[struct.unpack('L', el[1:])[0], key]
        table['eof']=[1, 'eof']
        
        #creation de larbre et des chemins
        arbre, chem=self.__makeArbreFromTable(table.copy(), totcars)
        
        racine=arbre.getRacine()
        curnode=racine
        f1.seek(offset)
        while 1:
            buf=f1.read(1024*256)
            if not buf : break #preventif
            else:
                for car in buf:
                    for i in range(8):
                        bit=ord(car)&(0x80>>i)
                        if(bit!=0) : bit=1
                        item=arbre.GetDir(curnode)[bit]
                        if(item.getType()=='feuille') :
                            if(item.getValue()!='eof') : f2.write(item.getValue()); curnode=racine
                            else : break
                        else : curnode=item
        f1.close()
        f2.close()

if __name__=='__main__':
    if(len(sys.argv)!=4) : print 'usage : Huffman.py mode source dest
    h=Huffman()
    a=open(sys.argv[2], 'rb')
    b=open(sys.argv[3], 'wb')
    if(sys.argv[1]=='c') : h.compress(a, b)
    else : h.decompress(a, b)
    raw_input('operation finie')