|
|
|
|
@ -1,3 +1,7 @@
|
|
|
|
|
#!/usr/bin/env python3
|
|
|
|
|
|
|
|
|
|
# coding: utf-8
|
|
|
|
|
# author: u/Trivernis
|
|
|
|
|
import os
|
|
|
|
|
import shutil
|
|
|
|
|
import yaml
|
|
|
|
|
@ -8,6 +12,7 @@ import urllib.request as urlreq
|
|
|
|
|
|
|
|
|
|
user_agent = 'python:riddle:3.0 (by u/Trivernis)' # the reddit api user-agent
|
|
|
|
|
img_ext = ['jpg', 'jpeg', 'png'] # default used extensions to filter for images
|
|
|
|
|
min_size = 5 # minimum size in kilobytes. changeable in settings
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def assert_dir_exist(dirpath):
|
|
|
|
|
@ -20,32 +25,56 @@ def assert_dir_exist(dirpath):
|
|
|
|
|
os.mkdir(dirpath)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def download_file(url: str, dest: str):
|
|
|
|
|
def download_file(url: str, dest: str, progressbar = None):
|
|
|
|
|
"""
|
|
|
|
|
Downloads a url to a file
|
|
|
|
|
:param url: download url
|
|
|
|
|
:param dest: download destination
|
|
|
|
|
:param progressbar: The progressbar instance to clear it before writing an error message
|
|
|
|
|
:return: Success?
|
|
|
|
|
"""
|
|
|
|
|
f = open(dest, "wb")
|
|
|
|
|
req = urlreq.Request(url)
|
|
|
|
|
success = False
|
|
|
|
|
try:
|
|
|
|
|
image = urlreq.urlopen(req)
|
|
|
|
|
f.write(image.read())
|
|
|
|
|
f.close()
|
|
|
|
|
return True
|
|
|
|
|
success = True
|
|
|
|
|
except ConnectionError:
|
|
|
|
|
print('\r[-] Connection Error \r')
|
|
|
|
|
return False
|
|
|
|
|
if progressbar:
|
|
|
|
|
progressbar.clear()
|
|
|
|
|
print('\r[-] Connection Error')
|
|
|
|
|
except urlreq.HTTPError as err:
|
|
|
|
|
print('\r[-] HTTPError for %s: %s \r' % (url, err))
|
|
|
|
|
return False
|
|
|
|
|
if progressbar:
|
|
|
|
|
progressbar.clear()
|
|
|
|
|
print('\r[-] HTTPError for %s: %s' % (url, err))
|
|
|
|
|
except urlreq.URLError as err:
|
|
|
|
|
if progressbar:
|
|
|
|
|
progressbar.clear()
|
|
|
|
|
print('\r[-] URLError for %s: %s' % (url, err))
|
|
|
|
|
f.close()
|
|
|
|
|
try:
|
|
|
|
|
file_size = round(os.path.getsize(dest) / 1000)
|
|
|
|
|
if not success:
|
|
|
|
|
os.remove(dest)
|
|
|
|
|
elif file_size < min_size:
|
|
|
|
|
os.remove(dest)
|
|
|
|
|
success = False
|
|
|
|
|
if progressbar:
|
|
|
|
|
progressbar.clear()
|
|
|
|
|
print('\r[-] Removed %s: Too small (%s kb)' % (dest, file_size))
|
|
|
|
|
except IOError as err:
|
|
|
|
|
if progressbar:
|
|
|
|
|
progressbar.clear()
|
|
|
|
|
print('\r[-] Error when removing file %s: %s' % (dest, err))
|
|
|
|
|
return success
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ProgressBar:
|
|
|
|
|
"""
|
|
|
|
|
A simple progressbar.
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
def __init__(self, total=100, prefix='', suffix='', length=50, fill='█'):
|
|
|
|
|
self.prefix = prefix
|
|
|
|
|
self.suffix = suffix
|
|
|
|
|
@ -53,12 +82,22 @@ class ProgressBar:
|
|
|
|
|
self.length = length
|
|
|
|
|
self.total = total
|
|
|
|
|
self.progress = 0
|
|
|
|
|
self.textlength = 0
|
|
|
|
|
|
|
|
|
|
def tick(self):
|
|
|
|
|
"""
|
|
|
|
|
Next step of the progressbar. The stepwidth is always 1.
|
|
|
|
|
:return:
|
|
|
|
|
"""
|
|
|
|
|
self.progress += 1
|
|
|
|
|
self._print_progress()
|
|
|
|
|
|
|
|
|
|
def setprogress(self, progress):
|
|
|
|
|
def setprogress(self, progress: float):
|
|
|
|
|
"""
|
|
|
|
|
Set the progress of the bar.
|
|
|
|
|
:param progress: progress in percent
|
|
|
|
|
:return: None
|
|
|
|
|
"""
|
|
|
|
|
self.progress = progress
|
|
|
|
|
self._print_progress()
|
|
|
|
|
|
|
|
|
|
@ -71,11 +110,20 @@ class ProgressBar:
|
|
|
|
|
percent = ("{0:." + str(1) + "f}").format(100 * (iteration / float(total)))
|
|
|
|
|
filled_length = int(self.length * iteration // total)
|
|
|
|
|
bar = self.fill * filled_length + '-' * (self.length - filled_length)
|
|
|
|
|
print('\r%s |%s| %s%% %s' % (prefix, bar, percent, suffix), end='\r')
|
|
|
|
|
# Print New Line on Complete
|
|
|
|
|
textout = '\r%s |%s| %s%% %s' % (prefix, bar, percent, suffix)
|
|
|
|
|
print(textout, end='\r')
|
|
|
|
|
self.textlength = len(textout)
|
|
|
|
|
# Print new line on complete
|
|
|
|
|
if iteration == total:
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
def clear(self):
|
|
|
|
|
"""
|
|
|
|
|
clear last progress output
|
|
|
|
|
:return:
|
|
|
|
|
"""
|
|
|
|
|
print(' '*self.textlength, end='\r')
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def parser_init():
|
|
|
|
|
"""
|
|
|
|
|
@ -89,13 +137,17 @@ def parser_init():
|
|
|
|
|
If not set it is the maximum fetchable number.""")
|
|
|
|
|
parser.add_option('-o', '--output', dest='output',
|
|
|
|
|
type='str', default=None,
|
|
|
|
|
help='The name of the output folder. If none is specified, it\'s the subreddits name.')
|
|
|
|
|
help="""The name of the output folder.
|
|
|
|
|
If none is specified, it\'s the subreddits name.""")
|
|
|
|
|
parser.add_option('-z', '--zip', dest='zip',
|
|
|
|
|
action='store_true', default=False,
|
|
|
|
|
help='Stores the images in a zip file if true')
|
|
|
|
|
parser.add_option('-n', '--nsfw', dest='nsfw',
|
|
|
|
|
parser.add_option('--nsfw', dest='nsfw',
|
|
|
|
|
action='store_true', default=False,
|
|
|
|
|
help='If set nsfw-content is also downloaded.')
|
|
|
|
|
parser.add_option('--lzma', dest='lzma',
|
|
|
|
|
action='store_true', default=False,
|
|
|
|
|
help='If set the lzma-compression module is used.')
|
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ -110,7 +162,7 @@ def get_images(reddit_client: praw.Reddit, subreddit: str, limit: int, nsfw: boo
|
|
|
|
|
"""
|
|
|
|
|
print('[~] Fetching images for r/%s...' % subreddit)
|
|
|
|
|
urls = [submission.url for submission in reddit_client.subreddit(subreddit).hot(limit=limit)
|
|
|
|
|
if not submission.over_18 or nsfw] # fetches hot images and filters by nsfw if nsfw not set to true
|
|
|
|
|
if not submission.over_18 or nsfw] # fetches hot images and filters nsfw if set to false
|
|
|
|
|
return [url for url in urls if url.split('.')[-1] in img_ext]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ -128,16 +180,16 @@ def download_images(images: list, dl_dir: str):
|
|
|
|
|
assert_dir_exist(dl_dir)
|
|
|
|
|
|
|
|
|
|
for img in images: # download each image if it doesn't exist
|
|
|
|
|
pb.tick()
|
|
|
|
|
success = False
|
|
|
|
|
imgname = img.split('/')[-1]
|
|
|
|
|
name = os.path.join(dl_dir, imgname)
|
|
|
|
|
if not os.path.isfile(name):
|
|
|
|
|
success = download_file(img, name)
|
|
|
|
|
success = download_file(img, name, pb)
|
|
|
|
|
else:
|
|
|
|
|
preexist += 1
|
|
|
|
|
if success:
|
|
|
|
|
realcount += 1
|
|
|
|
|
pb.tick()
|
|
|
|
|
print('[+] Successfully downloaded %s out of %s images to %s (%s already existed)' %
|
|
|
|
|
(realcount, imgcount, dl_dir, preexist))
|
|
|
|
|
|
|
|
|
|
@ -158,11 +210,12 @@ def filter_zip_files(images: list, zip_fname: str):
|
|
|
|
|
return images
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def compress_folder(folder: str, zip_fname: str):
|
|
|
|
|
def compress_folder(folder: str, zip_fname: str, compression: int):
|
|
|
|
|
"""
|
|
|
|
|
Zips the contents of a folder to the destination zipfile name.
|
|
|
|
|
:param folder: the folder to zip
|
|
|
|
|
:param zip_fname: the name of the destination zipfile
|
|
|
|
|
:param compression: The compression method (constant from zipfile module)
|
|
|
|
|
:return: None
|
|
|
|
|
"""
|
|
|
|
|
print('[~] Compressing folder...')
|
|
|
|
|
@ -171,7 +224,7 @@ def compress_folder(folder: str, zip_fname: str):
|
|
|
|
|
if os.path.isfile(zip_fname): # append to the zipfile if it already exists
|
|
|
|
|
mode = 'a'
|
|
|
|
|
|
|
|
|
|
zfile = zipfile.ZipFile(zip_fname, mode)
|
|
|
|
|
zfile = zipfile.ZipFile(zip_fname, mode, compression=compression)
|
|
|
|
|
|
|
|
|
|
for _, _, files in os.walk(folder): # add all files of the folder to the zipfile
|
|
|
|
|
for file in files:
|
|
|
|
|
@ -181,16 +234,24 @@ def compress_folder(folder: str, zip_fname: str):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def main():
|
|
|
|
|
"""
|
|
|
|
|
Main entry method. Loads the settings and iterates through subreddits and downloads all images it fetched.
|
|
|
|
|
If the --zip flag is set, the images will be downloaded in a .cache directory and then compressed.
|
|
|
|
|
"""
|
|
|
|
|
options, subreddits = parser_init()
|
|
|
|
|
with open('config.yaml', 'r') as file: # loads the config.yaml file
|
|
|
|
|
config_fname = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'config.yaml')
|
|
|
|
|
with open(config_fname, 'r') as file: # loads the config.yaml file
|
|
|
|
|
try:
|
|
|
|
|
settings = yaml.safe_load(file)
|
|
|
|
|
except yaml.YAMLError as err:
|
|
|
|
|
print(err)
|
|
|
|
|
if settings:
|
|
|
|
|
if 'image-extensions' in settings: # uses image extensions specified in config.yaml fallback to default
|
|
|
|
|
if 'image-extensions' in settings:
|
|
|
|
|
global img_ext
|
|
|
|
|
img_ext = settings['image-extensions']
|
|
|
|
|
if 'min-size' in settings:
|
|
|
|
|
global min_size
|
|
|
|
|
min_size = int(settings['min-size'])
|
|
|
|
|
credentials = settings['credentials']
|
|
|
|
|
client = praw.Reddit(
|
|
|
|
|
client_id=credentials['client_id'],
|
|
|
|
|
@ -200,13 +261,18 @@ def main():
|
|
|
|
|
for subreddit in subreddits:
|
|
|
|
|
dldest = subreddit
|
|
|
|
|
if options.output:
|
|
|
|
|
dldest = options.output # uses the -o output destination instead of a folder with the subreddit name
|
|
|
|
|
images = get_images(client, subreddit, limit=options.count, nsfw=options.nsfw)
|
|
|
|
|
dldest = options.output # uses the -o output destination
|
|
|
|
|
images = get_images(client, subreddit, limit=options.count,
|
|
|
|
|
nsfw=options.nsfw)
|
|
|
|
|
if options.zip: # downloads to a cache-folder first before compressing it to zip
|
|
|
|
|
comp_mode = zipfile.ZIP_STORED
|
|
|
|
|
if options.lzma:
|
|
|
|
|
comp_mode = zipfile.ZIP_LZMA
|
|
|
|
|
cachedir = '.cache-' + dldest.split('/')[-1]
|
|
|
|
|
images = filter_zip_files(images, dldest+'.zip')
|
|
|
|
|
download_images(images, '.cache')
|
|
|
|
|
compress_folder('.cache', dldest+'.zip')
|
|
|
|
|
shutil.rmtree('.cache')
|
|
|
|
|
download_images(images, cachedir)
|
|
|
|
|
compress_folder(cachedir, dldest+'.zip', compression=comp_mode)
|
|
|
|
|
shutil.rmtree(cachedir)
|
|
|
|
|
else:
|
|
|
|
|
download_images(images, dldest)
|
|
|
|
|
print('[+] All downloads finished')
|
|
|
|
|
|
Trivernis
5 years ago
committed by
GitHub