[wip] split exported zips

cogs/emote.py

@@ -55,6 +55,7 @@ class Emotes(commands.Cog):
 	TAR_MIMETYPES = {'application/x-tar'}
 	ZIP_MIMETYPES = {'application/zip', 'application/octet-stream', 'application/x-zip-compressed', 'multipart/x-zip'}
 	ARCHIVE_MIMETYPES = TAR_MIMETYPES | ZIP_MIMETYPES
+	ZIP_OVERHEAD_BYTES = 22
 
 	def __init__(self, bot):
 		self.bot = bot
@@ -247,23 +248,57 @@ class Emotes(commands.Cog):
 		if not emotes:
 			raise commands.BadArgument('No emotes of that type were found in this server.')
 
-		out = io.BytesIO()
 		async with context.typing():
+			async for zip_file in self.archive_emotes(context, emotes):
+				await context.send(file=zip_file)
+
+	async def archive_emotes(self, context, emotes):
+		d = collections.deque(emotes)
+		discrims = collections.defaultdict(int)
+		count = 1
+		while True:
+			out = io.BytesIO()
 			with zipfile.ZipFile(out, 'w', compression=zipfile.ZIP_STORED) as zip:
-				async def store(emote):
+				while True:
+					try:
+						emote = d.popleft()
+					except IndexError:
+						break
+
+					# don't put two files in the zip with the same name
+					discrims[emote.name] += 1
+					discrim = discrims[emote.name]
+					if discrim == 1:
+						name = emote.name
+					else:
+						name = f'{emote.name}-{discrim}'
+
 					# 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)
-					if type(data) is str:
+					name = f'{name}.{"gif" if emote.animated else "png"}'
+					if type(data) is str:  # error case
 						await context.send(f'{emote}: {data}')
-						return
+						continue
-					zinfo = zipfile.ZipInfo(
+
-						f'{emote.name}.{"gif" if emote.animated else "png"}',
+					if len(data) >= context.guild.filesize_limit:
-						date_time=emote.created_at.timetuple()[:6])
+						await context.send(f'{emote} could not be added because it alone exceeds the file size limit.')
+						continue
+					estimated_zip_overhead = len(name) + self.ZIP_OVERHEAD_BYTES
+					if out.tell() + len(data) + estimated_zip_overhead >= context.guild.filesize_limit:
+						# adding this emote would bring us over the file size limit
+						d.appendleft(emote)
+						break
+
+					zinfo = zipfile.ZipInfo(name, date_time=emote.created_at.timetuple()[:6])
 					zip.writestr(zinfo, data)
-				await utils.gather_or_cancel(*(store(emote) for emote in emotes))
+
+				if out.tell() == 0:
+					# no emotes were written
+					break
 
 			out.seek(0)
-		await context.send(file=discord.File(out, f'emotes-{context.guild.id}.zip'))
+			yield discord.File(out, f'emotes-{context.guild.id}-{count}.zip')
+			count += 1
 
 	@commands.command(name='import', aliases=['add-zip', 'add-tar', 'add-from-zip', 'add-from-tar'])
 	async def import_(self, context, url=None):