Merged .husbando and .waifu to a single function, improved docs

bot/plugins/useless.py

@@ -129,39 +129,7 @@ class Useless(DatabasePlugin):
 
         %%waifu [<nick>...]
         """
-        nick = ' '.join(args.get('<nick>',  mask.nick))
+        return self.husbando_waifu('waifu', mask, target, args)
-
-        if nick.startswith('='):
-            waifu = nick[1:]
-
-            try:
-                self.cur.execute('''
-                    INSERT INTO
-                      users (nick, waifu)
-                    VALUES
-                      (lower(%s), %s)
-                    ON CONFLICT (nick) DO UPDATE SET
-                      waifu = excluded.waifu
-                    ''', [mask.nick, waifu])
-                self.con.commit()
-
-                self.bot.notice(mask.nick, 'Waifu set to: {}'.format(waifu))
-            except Error as ex:
-                self.log.error(ex)
-                self.con.rollback()
-        else:
-            self.cur.execute('''
-                SELECT
-                  waifu
-                FROM
-                  users
-                WHERE
-                  lower(nick) = lower(%s)
-                ''', [nick])
-            result = self.cur.fetchone()['waifu']
-
-            if result:
-                return '\x02[Waifu]\x02 {}: {}'.format(nick, result)
 
     @command
     def husbando(self, mask: IrcString, target: IrcString, args: DocOptDict):
@@ -169,43 +137,10 @@ class Useless(DatabasePlugin):
 
         %%husbando [<nick>...]
         """
-        nick = ' '.join(args.get('<nick>', mask.nick))
+        return self.husbando_waifu('husbando', mask, target, args)
-
-        if nick.startswith('='):
-            nick = nick[1:]
-
-            try:
-                self.cur.execute('''
-                    INSERT INTO
-                      users (nick, husbando)
-                    VALUES
-                      (lower(%s), %s)
-                    ON CONFLICT (nick) DO UPDATE SET
-                      husbando = excluded.husbando
-                    ''', [mask.nick, nick])
-                self.con.commit()
-
-                self.bot.notice(mask.nick, 'Husbando set to: {}'.format(nick))
-            except Error as ex:
-                self.log.error(ex)
-                self.con.rollback()
-        else:
-            self.cur.execute('''
-                SELECT
-                  husbando
-                FROM
-                  users
-                WHERE
-                  lower(nick) = lower(%s)
-                ''', [nick])
-            result = self.cur.fetchone()['husbando']
-
-            if result:
-                return '\x02[Husbando]\x02 {}: {}'.format(nick, result)
 
     @command
-    def storyofpomfface(self, mask: IrcString, target: IrcString,
+    def storyofpomfface(self, mask: IrcString, target: IrcString, args: DocOptDict):
-                        args: DocOptDict):
         """Story of pomf face
 
         %%storyofpomfface
@@ -250,7 +185,7 @@ class Useless(DatabasePlugin):
 
     @command
     def bier(self, mask: IrcString, target: IrcString, args: DocOptDict):
-        """Gives nick a beer
+        """Gives a user a beer
 
         %%bier [<nick>]
         """
@@ -259,7 +194,7 @@ class Useless(DatabasePlugin):
 
     @command
     def fucken(self, mask: IrcString, target: IrcString, args: DocOptDict):
-        """Kills and fucks a nick
+        """Kills and fucks a user
 
         %%fucken [<nick>]
         """
@@ -268,31 +203,15 @@ class Useless(DatabasePlugin):
 
     @command(aliases=['anhero', 'sudoku'])
     def seppuku(self, mask: IrcString, target: IrcString, args: DocOptDict):
-        """Kicks a nick
+        """Kicks a user
 
         %%seppuku
         """
         self.bot.kick(target, mask.nick, 'Sayonara bonzai-chan...')
 
-    # @command
-    # def anhero(self, mask: IrcString, target: IrcString, args: DocOptDict):
-    #     """Kicks a nick
-    #
-    #     %%anhero
-    #     """
-    #     self.bot.kick(target, mask.nick, 'Sayonara bonzai-chan...')
-    #
-    # @command
-    # def sudoku(self, mask: IrcString, target: IrcString, args: DocOptDict):
-    #     """Kicks a nick
-    #
-    #     %%sudoku
-    #     """
-    #     self.anhero(mask, target, args)
-
     @command
     def hack(self, mask: IrcString, target: IrcString, args: DocOptDict):
-        """Hacks (a nick)
+        """Hacks (a user)
 
         %%hack [<nick>]
         """
@@ -371,6 +290,43 @@ class Useless(DatabasePlugin):
 
         return 'Assuming {}\'s gender... they\'re a {}.'.format(nick, gender)
 
+    def husbando_waifu(self, field: str, mask: IrcString, target: IrcString, args: DocOptDict):
+        nick = args.get('<nick>', mask.nick)
+        if isinstance(nick, list):
+            nick = ' '.join(nick)
+
+        if nick.startswith('='):
+            nick = nick[1:]
+
+            try:
+                self.cur.execute('''
+                    INSERT INTO
+                      users (nick, {0})
+                    VALUES
+                      (lower(%s), %s)
+                    ON CONFLICT (nick) DO UPDATE SET
+                      {0} = excluded.{0}
+                    '''.format(field), [mask.nick, nick])
+                self.con.commit()
+
+                self.bot.notice(mask.nick, '{} set to: {}'.format(field.title(), nick))
+            except Error as ex:
+                self.log.error(ex)
+                self.con.rollback()
+        else:
+            self.cur.execute('''
+                SELECT
+                  {}
+                FROM
+                  users
+                WHERE
+                  lower(nick) = lower(%s)
+                '''.format(field), [nick])
+            result = self.cur.fetchone()
+
+            if result and result[field]:
+                return '\x02[{}]\x02 {}: {}'.format(field.title(), nick, result[field])
+
     # noinspection PyMethodMayBeStatic
     def _rainbow(self, i, char):
         return '\x03{0:02d}{1}'.format(RAINBOW[i % RAINBOW_LEN], char)