Merge pull request #142 from tekulvw/alias-patch

Alias patch

cogs/alias.py

@@ -1,4 +1,3 @@
-import discord
 from discord.ext import commands
 from .utils.chat_formatting import *
 from .utils.dataIO import fileIO
@@ -6,38 +5,44 @@ from .utils import checks
 from __main__ import send_cmd_help
 import os
 
+
 class Alias:
-    def __init__(self,bot):
+    def __init__(self, bot):
         self.bot = bot
-        self.aliases = fileIO("data/alias/aliases.json","load")
+        self.aliases = fileIO("data/alias/aliases.json", "load")
 
     @commands.group(pass_context=True)
     @checks.mod_or_permissions(manage_server=True)
-    async def alias(self,ctx):
+    async def alias(self, ctx):
         """Manage per-server aliases for commands"""
         if ctx.invoked_subcommand is None:
             await send_cmd_help(ctx)
 
-    @alias.command(name="add",pass_context=True)
+    @alias.command(name="add", pass_context=True)
-    async def _add_alias(self,ctx,command : str,*,to_execute):
+    async def _add_alias(self, ctx, command: str, *, to_execute):
         """Add an alias for a command
 
            Example: !alias add test flip @Twentysix"""
         server = ctx.message.server
-        if self.get_prefix(to_execute) == False:
+        if self.part_of_existing_command(command, server.id):
-            to_execute = self.bot.command_prefix[0] + to_execute
+            await self.bot.say('I can\'t safely add an alias that starts with '
+                               'an existing command or alias. Sry <3')
+            return
+        prefix = self.get_prefix(to_execute)
+        if prefix is not None:
+            to_execute = to_execute[len(prefix):]
         if server.id not in self.aliases:
             self.aliases[server.id] = {}
-        #curr_aliases = self.aliases[server.id]
         if command not in self.bot.commands:
             self.aliases[server.id][command] = to_execute
-            fileIO("data/alias/aliases.json","save",self.aliases)
+            fileIO("data/alias/aliases.json", "save", self.aliases)
             await self.bot.say("Alias '{}' added.".format(command))
         else:
-            await self.bot.say("Cannot add '{}' because it's a real bot command.".format(command))        
+            await self.bot.say("Cannot add '{}' because it's a real bot "
+                               "command.".format(command))
 
-    @alias.command(name="help",pass_context=True)
+    @alias.command(name="help", pass_context=True)
-    async def _help_alias(self,ctx,command):
+    async def _help_alias(self, ctx, command):
         """Tries to execute help for the base command of the alias"""
         server = ctx.message.server
         if server.id in self.aliases:
@@ -53,8 +58,8 @@ class Alias:
             else:
                 await self.bot.say("That alias doesn't exist.")
 
-    @alias.command(name="show",pass_context=True)
+    @alias.command(name="show", pass_context=True)
-    async def _show_alias(self,ctx,command):
+    async def _show_alias(self, ctx, command):
         """Shows what command the alias executes."""
         server = ctx.message.server
         if server.id in self.aliases:
@@ -64,17 +69,18 @@ class Alias:
             else:
                 await self.bot.say("That alias doesn't exist.")
 
-    @alias.command(name="del",pass_context=True)
+    @alias.command(name="del", pass_context=True)
-    async def _del_alias(self,ctx,command : str):
+    async def _del_alias(self, ctx, command: str):
         """Deletes an alias"""
         server = ctx.message.server
         if server.id in self.aliases:
-            self.aliases[server.id].pop(command,None)
+            self.aliases[server.id].pop(command, None)
-            fileIO("data/alias/aliases.json","save",self.aliases)
+            fileIO("data/alias/aliases.json", "save", self.aliases)
         await self.bot.say("Alias '{}' deleted.".format(command))
 
-    async def check_aliases(self,message):
+    async def check_aliases(self, message):
-        if message.author.id == self.bot.user.id or len(message.content) < 2 or message.channel.is_private:
+        if message.author.id == self.bot.user.id or \
+                len(message.content) < 2 or message.channel.is_private:
             return
 
         msg = message.content
@@ -82,29 +88,48 @@ class Alias:
         prefix = self.get_prefix(msg)
 
         if prefix and server.id in self.aliases:
-            aliaslist = self.aliases[server.id]
+            for alias in self.aliases[server.id]:
-            alias = msg[len(prefix):].split(" ")[0]
+                if msg[len(prefix):].startswith(alias):
-            args = msg[len(self.first_word(message.content)):]
+                    new_command = self.aliases[server.id][alias]
-            if alias in aliaslist.keys():
+                    args = message.content[len(prefix+alias):]
-                content = aliaslist[alias] + args
+                    message.content = prefix + new_command + args
-                new_message = message
+                    await self.bot.process_commands(message)
-                new_message.content = content
-                await self.bot.process_commands(new_message)
 
-    def first_word(self,msg):
+    def part_of_existing_command(self, alias, server):
+        '''Command or alias'''
+        for command in self.bot.commands:
+            if alias.startswith(command):
+                return True
+        if server not in self.aliases:
+            return False
+        if alias.split(" ")[0] in self.aliases[server]:
+            return True
+        return False
+
+    def remove_old(self):
+        for sid in self.aliases:
+            for aliasname, alias in self.aliases[sid].items():
+                prefix = self.get_prefix(alias)
+                if prefix is not None:
+                    self.aliases[sid][aliasname] = alias[len(prefix):]
+        fileIO("data/alias/aliases.json", "save", self.aliases)
+
+    def first_word(self, msg):
         return msg.split(" ")[0]
 
     def get_prefix(self, msg):
         for p in self.bot.command_prefix:
             if msg.startswith(p):
                 return p
-        return False
+        return None
+
 
 def check_folder():
     if not os.path.exists("data/alias"):
         print("Creating data/alias folder...")
         os.makedirs("data/alias")
 
+
 def check_file():
     aliases = {}
 
@@ -113,9 +138,11 @@ def check_file():
         print("Creating default alias's aliases.json...")
         fileIO(f, "save", aliases)
 
+
 def setup(bot):
     check_folder()
     check_file()
     n = Alias(bot)
+    n.remove_old()
     bot.add_listener(n.check_aliases, "on_message")
     bot.add_cog(n)