diff --git a/karma.py b/karma.py index 9a8f82db3e04173079bfc0de966a8a0505c1bf77..41b4be2f326f8c3162d571c37cc6e5178462d119 100644 --- a/karma.py +++ b/karma.py @@ -60,14 +60,27 @@ class Karma(BotPlugin): return "Karma for '{}' is {}".format(args, self.karma[args]) + def _get_num_from_args(self, args): + try: + num = int(args[0]) + return num + except ValueError: + return None + except IndexError: + return 5 + @botcmd def karma_top(self, message, args): """Gets the top 5 karmaed items""" if not self.karma: return "There are no karmaed items" - keys = sorted(self.karma, key=self.karma.get, reverse=True)[:5] - values = sorted(self.karma.values(), reverse=True)[:5] + num = self._get_num_from_args(args) + if not num: + return "**Error**: Couldn't convert to an integer" + + keys = sorted(self.karma, key=self.karma.get, reverse=True)[:num] + values = sorted(self.karma.values(), reverse=True)[:num] return ", ".join(["{}: {}".format(k, v) for k, v in zip(keys, values)]) @@ -77,7 +90,11 @@ class Karma(BotPlugin): if not self.karma: return "There are no karmaed items" - keys = sorted(self.karma, key=self.karma.get)[:5] - values = sorted(self.karma.values())[:5] + num = self._get_num_from_args(args) + if not num: + return "**Error**: Couldn't convert to an integer" + + keys = sorted(self.karma, key=self.karma.get)[:num] + values = sorted(self.karma.values())[:num] return ", ".join(["{}: {}".format(k, v) for k, v in zip(keys, values)]) diff --git a/test_karma.py b/test_karma.py index 40382fcd1ebb1a89c1cb58ab122ff4894e570fa6..c9f56554eb35b2724de841d9c51da730772851fc 100644 --- a/test_karma.py +++ b/test_karma.py @@ -32,3 +32,6 @@ def test_top_bottom(testbot): testbot.assertCommand("!karma top", "high: 2, medium: 1, low: -1") testbot.assertCommand("!karma bottom", "low: -1, medium: 1, high: 2") + + testbot.assertCommand("!karma top 1", "high: 2") + testbot.assertCommand("!karma top 2", "high: 2, medium: 1")