[project @ zyp-20060516135739-925e0fe0819b13d6]

[project @ 32]
Replaced inline hash-code with hash-module.

2 files changed, 60 insertions, 53 deletions

diff --git a/anidb_add b/anidb_add
index 04dc1df..b0d06bd 100644
--- a/anidb_add
+++ b/anidb_add
@@ -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