Manage hashes in database files. (Python recipe) by eysi

ActiveState Code (http://code.activestate.com/recipes/578081/)

"""
This module can be used to calculate file hashes, store them in a database file
and retrieve them at a later date.

It uses the files modify time stamp to know if it can use the hash stored in
the db or if it has to re-calculate it. So the user will not have to worry
about the hash being incorrect if the file changes in between runs.

usage:
    hdb = HashDB()
    hdb.Load('dbfile.hashdb') # load hashes from db file
    md5 = hdb.Get('somefile', hashtype='md5') # retrieve hash
    hdb.Clean(maxAge=24*60*60) # removes old and invalid entries
    hdb.Write('dbfile.hashdb') # write hashes to db file
"""

__author__ = 'Eysteinn Torfi Kristinsson'
__copyright__ = "Copyright 2012, MIT License"

import hashlib
import cPickle
import os
import time

def CalcFileHash(path, hashtype, buffer=8192):
    '''
    Calculate the hash of a file.
    '''
    hash = getattr(hashlib, hashtype)()
    with open(path, 'rb') as fp:
        while 1:
            data = fp.read(buffer)
            if not data:
                break
            hash.update(data)
    return hash.digest()


class HashDB(object):
    '''
    Calculate and manage file hashes using database files.
    '''

    def __init__(self):
        self.db = {}

    def Load(self, dbFile):
        '''
        Load DB from a file.
        '''
        with open(dbFile, 'rb') as fp:
            self.db.update(cPickle.load(fp))

    def Cleanup(self, maxAge=None):
        '''
        Deletes entries that are no longer valid.
        Checks if file exists and if it has the same modify time stamp, if not
        then the entry is deleted from the DB.
        maxAge (in seconds) can also be passed in, it is used to check against
        the last time the entry was accessed in the DB and if the entry is
        older than maxAge, then it is deleted from the DB.
        '''
        deletes = []
        for path, (mtime, atime, cache) in self.db.iteritems():
            if not os.path.isfile(path) or \
                        (maxAge is not None and atime < time.time()-maxAge) or \
                        mtime != int(os.stat(path).st_mtime):
                deletes.append(path)
        for path in deletes:
            del self.db[path]

    def Get(self, path, hashtype='md5'):
        '''
        Retrieve hash from DB. If not found in the DB or the file has been
        modified since last it's hash was retrieved, a new one will be
        calculated and stored in the DB.
        '''
        # self.db[path][0] is modify time of file
        # self.db[path][1] is access time of db entry
        # self.db[path][2] is dict containing file hashes
        mtime = int(os.stat(path).st_mtime)
        if path in self.db:
            if self.db[path][0] != mtime:
                self.db[path] = [mtime, 0, {}]
            if not hashtype in self.db[path][2]:
                self.db[path][2][hashtype] = CalcFileHash(path, hashtype)
        else:
            self.db[path] = [mtime, 0,
                {hashtype : CalcFileHash(path, hashtype)}]
        self.db[path][1] = time.time() # update access time
        return self.db[path][2][hashtype]

    def Write(self, dbFile):
        '''
        Write DB to file.
        '''
        with open(dbFile, 'wb') as fp:
            cPickle.dump(self.db, fp)


def test():
    import sys, tempfile, shutil
    startTime = time.time()
    testfolder = '.'
    dbFile = 'data.hashdb'
    hdb = HashDB()
    if os.path.isfile(dbFile):
        print 'Loading from file'
        hdb.Load(dbFile)
    try:
        for root, dirs, files in os.walk(testfolder):
            for fileName in files:
                fileAbs = os.path.abspath(os.path.join(root, fileName))
                if os.path.isfile(fileAbs):
                    hashtype = 'sha1'
                    print '%s (%s): %r' % (
                        fileAbs, 
                        hashtype, 
                        hdb.Get(fileAbs, hashtype),
                        )
    finally:
        print 'Cleaning DB'
        hdb.Cleanup(60*30) # 30 minutes
        print 'writing to file'
        # incase KeyboardInterrupt, write to tempfile so we don't destroy
        # the database we had previously 
        dbFilePath, dbFileName = os.path.split(dbFile)
        tmpFile = tempfile.mktemp(
            dir = dbFilePath,
            prefix = dbFileName+'_',
            suffix='.tmp',
            )
        try:
            hdb.Write(tmpFile)
            shutil.move(tmpFile, dbFile)
        except:
            os.unlink(tmpFile)
            raise
    td = time.time()-startTime
    print td
    print 'total runtime: %dm%02ds' % (td/60, td%60)

if __name__ == '__main__':
    test()