mirror of
https://github.com/uhIgnacio/EmoteManager.git
synced 2024-08-15 02:23:13 +00:00
split zips in parallel (fixes #3)
This commit is contained in:
parent
a175050e11
commit
3590a0162d
1 changed files with 39 additions and 25 deletions
|
@ -55,7 +55,7 @@ class Emotes(commands.Cog):
|
||||||
TAR_MIMETYPES = {'application/x-tar'}
|
TAR_MIMETYPES = {'application/x-tar'}
|
||||||
ZIP_MIMETYPES = {'application/zip', 'application/octet-stream', 'application/x-zip-compressed', 'multipart/x-zip'}
|
ZIP_MIMETYPES = {'application/zip', 'application/octet-stream', 'application/x-zip-compressed', 'multipart/x-zip'}
|
||||||
ARCHIVE_MIMETYPES = TAR_MIMETYPES | ZIP_MIMETYPES
|
ARCHIVE_MIMETYPES = TAR_MIMETYPES | ZIP_MIMETYPES
|
||||||
ZIP_OVERHEAD_BYTES = 22
|
ZIP_OVERHEAD_BYTES = 30
|
||||||
|
|
||||||
def __init__(self, bot):
|
def __init__(self, bot):
|
||||||
self.bot = bot
|
self.bot = bot
|
||||||
|
@ -253,18 +253,11 @@ class Emotes(commands.Cog):
|
||||||
await context.send(file=zip_file)
|
await context.send(file=zip_file)
|
||||||
|
|
||||||
async def archive_emotes(self, context, emotes):
|
async def archive_emotes(self, context, emotes):
|
||||||
d = collections.deque(emotes)
|
filesize_limit = 1024 ** 2 # XXX testing
|
||||||
discrims = collections.defaultdict(int)
|
|
||||||
count = 1
|
|
||||||
while True:
|
|
||||||
out = io.BytesIO()
|
|
||||||
with zipfile.ZipFile(out, 'w', compression=zipfile.ZIP_STORED) as zip:
|
|
||||||
while True:
|
|
||||||
try:
|
|
||||||
emote = d.popleft()
|
|
||||||
except IndexError:
|
|
||||||
break
|
|
||||||
|
|
||||||
|
discrims = collections.defaultdict(int)
|
||||||
|
downloaded = collections.deque()
|
||||||
|
async def download(emote):
|
||||||
# don't put two files in the zip with the same name
|
# don't put two files in the zip with the same name
|
||||||
discrims[emote.name] += 1
|
discrims[emote.name] += 1
|
||||||
discrim = discrims[emote.name]
|
discrim = discrims[emote.name]
|
||||||
|
@ -273,24 +266,45 @@ class Emotes(commands.Cog):
|
||||||
else:
|
else:
|
||||||
name = f'{emote.name}-{discrim}'
|
name = f'{emote.name}-{discrim}'
|
||||||
|
|
||||||
|
name = f'{name}.{"gif" if emote.animated else "png"}'
|
||||||
|
|
||||||
# place some level of trust on discord's CDN to actually give us images
|
# place some level of trust on discord's CDN to actually give us images
|
||||||
data = await self.fetch_safe(str(emote.url), validate_headers=False)
|
data = await self.fetch_safe(str(emote.url), validate_headers=False)
|
||||||
name = f'{name}.{"gif" if emote.animated else "png"}'
|
|
||||||
if type(data) is str: # error case
|
if type(data) is str: # error case
|
||||||
await context.send(f'{emote}: {data}')
|
await context.send(f'{emote}: {data}')
|
||||||
continue
|
return
|
||||||
|
|
||||||
if len(data) >= context.guild.filesize_limit:
|
est_zip_overhead = len(name) + self.ZIP_OVERHEAD_BYTES
|
||||||
await context.send(f'{emote} could not be added because it alone exceeds the file size limit.')
|
est_size_in_zip = est_zip_overhead + len(data)
|
||||||
continue
|
if est_size_in_zip >= filesize_limit:
|
||||||
estimated_zip_overhead = len(name) + self.ZIP_OVERHEAD_BYTES
|
self.bot.loop.create_task(
|
||||||
if out.tell() + len(data) + estimated_zip_overhead >= context.guild.filesize_limit:
|
context.send(f'{emote} could not be added because it alone would exceed the file size limit.')
|
||||||
# adding this emote would bring us over the file size limit
|
)
|
||||||
d.appendleft(emote)
|
return
|
||||||
|
|
||||||
|
downloaded.append((name, emote.created_at, est_size_in_zip, data))
|
||||||
|
|
||||||
|
await utils.gather_or_cancel(*map(download, emotes))
|
||||||
|
|
||||||
|
count = 1
|
||||||
|
while True:
|
||||||
|
out = io.BytesIO()
|
||||||
|
with zipfile.ZipFile(out, 'w', compression=zipfile.ZIP_STORED) as zip:
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
item = downloaded.popleft()
|
||||||
|
except IndexError:
|
||||||
break
|
break
|
||||||
|
|
||||||
zinfo = zipfile.ZipInfo(name, date_time=emote.created_at.timetuple()[:6])
|
name, created_at, est_size, image_data = item
|
||||||
zip.writestr(zinfo, data)
|
|
||||||
|
if out.tell() + est_size >= filesize_limit:
|
||||||
|
# adding this emote would bring us over the file size limit
|
||||||
|
downloaded.appendleft(item)
|
||||||
|
break
|
||||||
|
|
||||||
|
zinfo = zipfile.ZipInfo(name, date_time=created_at.timetuple()[:6])
|
||||||
|
zip.writestr(zinfo, image_data)
|
||||||
|
|
||||||
if out.tell() == 0:
|
if out.tell() == 0:
|
||||||
# no emotes were written
|
# no emotes were written
|
||||||
|
|
Loading…
Reference in a new issue