diff options
| -rw-r--r-- | anidb_add | 73 | ||||
| -rw-r--r-- | pyanidb/hash.py | 40 |
2 files changed, 60 insertions, 53 deletions
@@ -1,40 +1,7 @@ #!/usr/bin/env python -import pyanidb as anidb -import ConfigParser, os, sys, thread, time, getpass, multihash - -num_threads = 0 - -def file_hash(name): - e = multihash.Ed2k() - f = open(name) - data = f.read(32768) - while data: - e.update(data) - data = f.read(32768) - f.close() - return e.digest() - -def hash_file(name): - if not os.access(name, os.R_OK): - print 'Invalid file: %s' % (name) - return - size = os.stat(name).st_size - hash = file_hash(name) - print 'Hashed: ed2k://|file|%s|%d|%s|' % (name, size, hash) - return name, size, hash - -def hash_thread(filelist, hashlist): - global num_threads - num_threads += 1 - try: - while filelist: - h = hash_file(filelist.pop(0)) - if h: - hashlist.append(h) - except IndexError: - pass - num_threads -= 1 +import pyanidb, pyanidb.hash +import ConfigParser, os, sys, getpass, multihash def auth(): try: @@ -50,7 +17,7 @@ def auth(): username, password = auth() try: - a = anidb.AniDB(username, password) + a = pyanidb.AniDB(username, password) #t = a.ping() #if t: # print 'AniDB is reachable, %.3fs' % (t) @@ -62,35 +29,35 @@ try: if a.new_version: print 'New version available.' - filelist = sys.argv[1:] - hashlist = [] - - thread.start_new_thread(hash_thread, (filelist, hashlist)) - - while hashlist or num_threads or filelist: - if not hashlist: - time.sleep(0.1) - continue - name, size, hash = hashlist.pop(0) + files = [] + for name in sys.argv[1:]: + if os.access(name, os.R_OK): + files.append(name) + else: + print 'Invalid file: %s' % (name) + for filename, hash in pyanidb.hash.hash_files(files): + hash = hash.ed2k() + size = os.stat(filename).st_size + print 'Hashed: ed2k://|file|%s|%d|%s|' % (filename, size, hash) try: while 1: try: a.add_hash(size, hash) - except anidb.AniDBTimeout: + except pyanidb.AniDBTimeout: print 'Connection timed out, retrying.' continue break - except anidb.AniDBUnknownFile: - print 'Unknown file: %s' % (name) + except pyanidb.AniDBUnknownFile: + print 'Unknown file: %s' % (filename) continue - print 'Added file: %s' % (name) + print 'Added file: %s' % (filename) print 'All operations finished.' -except anidb.AniDBUserError: +except pyanidb.AniDBUserError: print 'Invalid username/password.' sys.exit(1) -except anidb.AniDBTimeout: +except pyanidb.AniDBTimeout: print 'Connection timed out.' sys.exit(1) -except anidb.AniDBError, e: +except pyanidb.AniDBError, e: print 'Fatal error:', e sys.exit(1) diff --git a/pyanidb/hash.py b/pyanidb/hash.py new file mode 100644 index 0000000..35c05ae --- /dev/null +++ b/pyanidb/hash.py @@ -0,0 +1,40 @@ +import multihash, threading, time + +def file_hash(name, algorithms): + h = multihash.Multihash(*algorithms) + f = open(name) + data = f.read(32768) + while data: + h.update(data) + data = f.read(32768) + f.close() + return h + +class Hashthread(threading.Thread): + def __init__(self, filelist, hashlist, algorithms, *args, **kwargs): + self.filelist = filelist + self.hashlist = hashlist + self.algorithms = algorithms + threading.Thread.__init__(self, *args, **kwargs) + def run(self): + try: + while 1: + f = self.filelist.pop(0) + h = file_hash(f, self.algorithms) + self.hashlist.append((f, h)) + except IndexError: + return + +def hash_files(files, num_threads = 1, algorithms = ('ed2k',)): + hashlist = [] + threads = [] + for x in xrange(num_threads): + thread = Hashthread(files, hashlist, algorithms) + thread.start() + threads.append(thread) + while hashlist or sum([thread.isAlive() for thread in threads]): + try: + yield hashlist.pop(0) + except IndexError: + time.sleep(0.1) + raise StopIteration |