def run(self,func,tasks,func2=None):
'''run will send a list of tasks, a tuple with arguments, through a function.
the arguments should be ordered correctly.
:param func: the function to run with multiprocessing.pool
:param tasks: a list of tasks, each a tuple of arguments to process
:param func2: filter function to run result from func through (optional)
'''
# Keep track of some progress for the user
progress = 1
total = len(tasks)
# If two functions are run per task, double total jobs
if func2 is not None:
total = total * 2
finished = []
level1 = []
results = []
try:
prefix = "[%s/%s]" %(progress,total)
bot.show_progress(0,total,length=35,prefix=prefix)
pool = multiprocessing.Pool(self.workers,init_worker)
self.start()
for task in tasks:
result = pool.apply_async(multi_wrapper,
multi_package(func,[task]))
results.append(result)
level1.append(result._job)
while len(results) > 0:
result = results.pop()
result.wait()
bot.show_progress(progress,total,length=35,prefix=prefix)
progress+=1
prefix = "[%s/%s]" %(progress,total)
# Pass the result through a second function?
if func2 is not None and result._job in level1:
result = pool.apply_async(multi_wrapper,
multi_package(func2,[(result.get(),)]))
results.append(result)
else:
finished.append(result.get())
self.end()
pool.close()
pool.join()
except (KeyboardInterrupt, SystemExit):
bot.error("Keyboard interrupt detected, terminating workers!")
pool.terminate()
sys.exit(1)
except Exception as e:
bot.error(e)
return finished
def stream(self,url,file_name,data=None,headers=None,default_headers=True,
show_progress=False):
'''stream is a get that will stream to file_name
:param data: a dictionary of key:value items to add to the data args variable
:param url: the url to get
:param show_progress: if True, show a progress bar with the bot
:returns response: the requests response object, or stream
'''
bot.debug("GET (stream) %s" %url)
# If we use default headers, start with client's
request_headers = dict()
if default_headers and len(self.headers) > 0:
request_headers = self.headers
if headers is not None:
request_headers.update(headers)
request = self.prepare_request(headers=request_headers,
data=data,
url=url)
response = self.submit_request(request)
# Keep user updated with Progress Bar?
if show_progress:
content_size = None
if 'Content-Length' in response.headers and response.code not in [400,401]:
progress = 0
content_size = int(response.headers['Content-Length'])
bot.show_progress(progress,content_size,length=35)
chunk_size = 1 << 20
with open(file_name, 'wb') as filey:
while True:
chunk = response.read(chunk_size)
if not chunk:
break
try:
filey.write(chunk)
if show_progress:
if content_size is not None:
progress+=chunk_size
bot.show_progress(iteration=progress,
total=content_size,
length=35,
carriage_return=False)
except Exception as error:
bot.error("Error writing to %s: %s exiting" %(file_name,error))
sys.exit(1)
# Newline to finish download
if show_progress:
sys.stdout.write('\n')
return file_name