moved danbooru clone interaction to a function in utils instead of in the NSFW cog so it can be used in many commands.

roxbot/cogs/nsfw.py

@@ -23,8 +23,6 @@
 # SOFTWARE.
 
 
-import random
-
 import discord
 from discord.ext import commands
 
@@ -39,7 +37,7 @@ def tag_blacklist(guild):
 	return blacklist
 
 
-class NFSW():
+class NFSW:
 	"""The NSFW cog is a collection of commands that post images from popular NSFW sites. """
 	def __init__(self, bot_client):
 		self.bot = bot_client
@@ -52,43 +50,25 @@ class NFSW():
 			}
 		}
 
-	@roxbot.checks.is_nsfw()
-	@commands.command(hidden=True)
-	async def gelbooru_clone(self, ctx, base_url, post_url, tags):
-		limit = 150
+	async def gelbooru_clone(self, ctx, base_url, endpoint_url, tags):
 		if isinstance(ctx.channel, discord.TextChannel):
-			tags = tags + tag_blacklist(ctx.guild)
-		page = random.randrange(20)
-		url = base_url + tags + '&limit=' + str(limit) + '%pid=' + str(page)
-		if isinstance(ctx.channel, discord.DMChannel):
-			cache_id = ctx.author.id
+			banned_tags = tag_blacklist(ctx.guild)
 		else:
-			cache_id = ctx.guild.id
-		# IF ID is not in cache, create cache for ID
-		if not self.cache.get(cache_id, False):
-			self.cache[cache_id] = []
-
-		posts = await roxbot.http.api_request(url)
-
-		if posts is None:
+			banned_tags = ""
+
+		post = await roxbot.utils.danbooru_clone_api_req(
+			ctx.channel,
+			base_url,
+			endpoint_url,
+			tags=tags,
+			banned_tags=banned_tags,
+			cache=self.cache
+		)
+
+		if not post:
 			return await ctx.send("Nothing was found. *psst, check the tags you gave me.*")
-
-		post = None
-		counter = 0
-		while counter < 20:
-			post = random.choice(posts)
-			md5 = post.get("md5") or post.get("hash")
-			if md5 not in self.cache[cache_id]:
-				self.cache[cache_id].append(md5)
-				if len(self.cache[cache_id]) > 10:
-					self.cache[cache_id].pop(0)
-				break
-			counter += 1
-
-		url = post.get("file_url")
-		if not url:
-			url = post_url + "{0[directory]}/{0[image]}".format(post)
-		output = await ctx.send(url)
+		else:
+			output = await ctx.send(post)
 		await roxbot.utils.delete_option(self.bot, ctx, output, self.bot.get_emoji(444410658101002261) or "❌")
 
 	@roxbot.checks.is_nsfw()
@@ -102,7 +82,7 @@ class NFSW():
 			;e621 test
 		"""
 		base_url = "https://e621.net/post/index.json?tags="
-		return await ctx.invoke(self.gelbooru_clone, base_url=base_url, post_url="", tags=tags)
+		return await self.gelbooru_clone(ctx, base_url, "", tags)
 
 	@roxbot.checks.is_nsfw()
 	@commands.command()
@@ -115,8 +95,8 @@ class NFSW():
 			;rule34 test
 		"""
 		base_url = "https://rule34.xxx/index.php?page=dapi&s=post&q=index&json=1&tags="
-		post_url = "https://img.rule34.xxx/images/"
-		return await ctx.invoke(self.gelbooru_clone, base_url=base_url, post_url=post_url, tags=tags)
+		endpoint_url = "https://img.rule34.xxx/images/"
+		return await self.gelbooru_clone(ctx, base_url, endpoint_url, tags)
 
 	@roxbot.checks.is_nsfw()
 	@commands.command()
@@ -129,8 +109,8 @@ class NFSW():
 			;gelbooru test
 		"""
 		base_url = "https://gelbooru.com/index.php?page=dapi&s=post&q=index&json=1&tags="
-		post_url = "https://simg3.gelbooru.com/images/"
-		return await ctx.invoke(self.gelbooru_clone, base_url=base_url, post_url=post_url, tags=tags)
+		endpoint_url = "https://simg3.gelbooru.com/images/"
+		return await self.gelbooru_clone(ctx, base_url, endpoint_url, tags)
 
 	@commands.guild_only()
 	@commands.has_permissions(manage_channels=True)
@@ -176,5 +156,6 @@ class NFSW():
 			return await ctx.send("No valid option given.")
 		return settings.update(nsfw, "nsfw")
 
+
 def setup(bot_client):
 	bot_client.add_cog(NFSW(bot_client))

roxbot/utils.py

@@ -23,10 +23,12 @@
 # SOFTWARE.
 
 
+import random
 import asyncio
 import discord
 import argparse
 
+from roxbot import http
 from roxbot import guild_settings
 from roxbot.enums import EmbedColours
 
@@ -167,3 +169,60 @@ async def log(guild, command_name, **kwargs):
 			for key, value in kwargs.items():
 				embed.add_field(name=key, value=value)
 			return await channel.send(embed=embed)
+
+
+async def danbooru_clone_api_req(channel, base_url, endpoint_url, cache=None, tags="", banned_tags=""):
+	"""Utility function that deals with danbooru clone api interaction.
+	It also deals with cache management for these interactions.
+
+	Params
+	=======
+	channel: discord.Channel
+		Channel command has been invoked in
+	base_url: str
+		Base url of the site
+	endpoint_url: str
+		Endpoint of images in the API. This is used if the API does not give this in its response.
+	cache: dict (optional)
+		Post cache. Were channel ID's are keys with values that are lists of identifiable info.
+		Cache is handled in this function and will be updated so that other functions can access it.
+	tags: str (optional)
+		tags to use in the search. Separated by spaces.
+	banned_tags: str (optional)
+		banned tags to append to the search. Separated by spaces with a - in front to remove them from search results.
+	"""
+	limit = "150"
+	tags = tags + banned_tags
+	page_number = str(random.randrange(20))
+	url = base_url + tags + '&limit=' + limit + '%pid=' + page_number
+
+	if isinstance(channel, discord.DMChannel):
+		cache_id = channel.id
+	else:
+		cache_id = channel.guild.id
+
+	# IF ID is not in cache, create cache for ID
+	if not cache.get(cache_id, False):
+		cache[cache_id] = []
+
+	posts = await http.api_request(url)
+
+	if posts is None:
+		return None
+
+	post = None
+	counter = 0
+	while counter < 20:
+		post = random.choice(posts)
+		md5 = post.get("md5") or post.get("hash")
+		if md5 not in cache[cache_id]:
+			cache[cache_id].append(md5)
+			if len(cache[cache_id]) > 10:
+				cache[cache_id].pop(0)
+			break
+		counter += 1
+
+	url = post.get("file_url")
+	if not url:
+		url = endpoint_url + "{0[directory]}/{0[image]}".format(post)
+	return url