From 4923ffe98a998f546118f517e6c6a83f4a1bb678 Mon Sep 17 00:00:00 2001 From: palmtree5 Date: Wed, 9 Aug 2017 17:23:41 -0800 Subject: [PATCH] [Economy] [WIP] rewrite (#781) * [Economy][Bank] redo branch * WIP WIP * Implement all current bank commands API calls * Set dunder all and put into bot * make core change to economy * Add is_global method to bank WIP * Add extra bank API commands * Update bank UI Update some imports Remove bank UI errors file Typing thing * Update bank get_global_accounts and touch up economy some more Do some more economy updates * Remove bank from bot * Another passing test FINALLY * Fixy type things Last fixes for now Fix arg to toggle global RJM Invalid bid amount handler cooldown msg currency name fix Fix fun bug ANother bug And payday limit * PEP8 stuff * Docstring change * Fix this thing * [Economy][Bank] redo branch * [Economy][Bank] modify guild owner or bot owner check, add admin or bot owner check for global vs local bank * [Economy] apply admin or bot owner check to [p]economyset * Make some public things private * [Economy] lots of refactoring for conditional permission checks and guild checks + supporting global economy * And working stuff * Fix Kowlin's bug * Fix slot bugs --- cogs/bank/__init__.py | 5 + cogs/bank/bank.py | 67 +++++ cogs/bank/errors.py | 37 +++ cogs/economy/__init__.py | 6 + cogs/economy/economy.py | 528 +++++++++++++++++++++++++++++++++++++ core/bank.py | 401 ++++++++++++++++++++++++++++ tests/cogs/test_economy.py | 55 ++++ tests/conftest.py | 5 +- 8 files changed, 1102 insertions(+), 2 deletions(-) create mode 100644 cogs/bank/__init__.py create mode 100644 cogs/bank/bank.py create mode 100644 cogs/bank/errors.py create mode 100644 cogs/economy/__init__.py create mode 100644 cogs/economy/economy.py create mode 100644 core/bank.py create mode 100644 tests/cogs/test_economy.py diff --git a/cogs/bank/__init__.py b/cogs/bank/__init__.py new file mode 100644 index 000000000..e467452e2 --- /dev/null +++ b/cogs/bank/__init__.py @@ -0,0 +1,5 @@ +from .bank import Bank, check_global_setting_guildowner, check_global_setting_admin + + +def setup(bot): + bot.add_cog(Bank(bot)) diff --git a/cogs/bank/bank.py b/cogs/bank/bank.py new file mode 100644 index 000000000..dbdab2684 --- /dev/null +++ b/cogs/bank/bank.py @@ -0,0 +1,67 @@ +from discord.ext import commands + +from core import checks, bank +from core.bot import Red # Only used for type hints + + +def check_global_setting_guildowner(): + async def pred(ctx: commands.Context): + if bank.is_global(): + return checks.is_owner() + else: + return checks.guildowner_or_permissions(administrator=True) + return commands.check(pred) + + +def check_global_setting_admin(): + async def pred(ctx: commands.Context): + if bank.is_global(): + return checks.is_owner() + else: + return checks.admin_or_permissions(manage_guild=True) + return commands.check(pred) + + +class Bank: + """Bank""" + + def __init__(self, bot: Red): + self.bot = bot + + # SECTION commands + + @commands.group() + @checks.guildowner_or_permissions(administrator=True) + async def bankset(self, ctx: commands.Context): + """Base command for bank settings""" + if ctx.invoked_subcommand is None: + await self.bot.send_cmd_help(ctx) + + @bankset.command(name="toggleglobal") + @checks.is_owner() + async def bankset_toggleglobal(self, ctx: commands.Context): + """Toggles whether the bank is global or not + If the bank is global, it will become per-guild + If the bank is per-guild, it will become global""" + cur_setting = bank.is_global() + await bank.set_global(not cur_setting, ctx.author) + + word = "per-guild" if cur_setting else "global" + + await ctx.send("The bank is now {}.".format(word)) + + @bankset.command(name="bankname") + @check_global_setting_guildowner() + async def bankset_bankname(self, ctx: commands.Context, *, name: str): + """Set the bank's name""" + await bank.set_bank_name(name, ctx.guild) + await ctx.send("Bank's name has been set to {}".format(name)) + + @bankset.command(name="creditsname") + @check_global_setting_guildowner() + async def bankset_creditsname(self, ctx: commands.Context, *, name: str): + """Set the name for the bank's currency""" + await bank.set_currency_name(name, ctx.guild) + await ctx.send("Currency name has been set to {}".format(name)) + + # ENDSECTION diff --git a/cogs/bank/errors.py b/cogs/bank/errors.py new file mode 100644 index 000000000..8cc962363 --- /dev/null +++ b/cogs/bank/errors.py @@ -0,0 +1,37 @@ +class BankError(Exception): + pass + +class BankNotGlobal(BankError): + pass + + +class BankIsGlobal(BankError): + pass + + +class AccountAlreadyExists(BankError): + pass + + +class NoAccount(BankError): + pass + + +class NoSenderAccount(NoAccount): + pass + + +class NoReceiverAccount(NoAccount): + pass + + +class InsufficientBalance(BankError): + pass + + +class NegativeValue(BankError): + pass + + +class SameSenderAndReceiver(BankError): + pass \ No newline at end of file diff --git a/cogs/economy/__init__.py b/cogs/economy/__init__.py new file mode 100644 index 000000000..43f6d5eb5 --- /dev/null +++ b/cogs/economy/__init__.py @@ -0,0 +1,6 @@ +from .economy import Economy +from core.bot import Red + + +def setup(bot: Red): + bot.add_cog(Economy(bot)) diff --git a/cogs/economy/economy.py b/cogs/economy/economy.py new file mode 100644 index 000000000..b0ac80a53 --- /dev/null +++ b/cogs/economy/economy.py @@ -0,0 +1,528 @@ +import calendar +import logging +import random +from collections import defaultdict, deque +from enum import Enum + +import discord +from discord.ext import commands +from core import checks, Config, bank +from core.utils.chat_formatting import pagify, box +from core.bot import Red +from cogs.bank import check_global_setting_guildowner, check_global_setting_admin + +logger = logging.getLogger("red.economy") + +NUM_ENC = "\N{COMBINING ENCLOSING KEYCAP}" + + +class SMReel(Enum): + cherries = "\N{CHERRIES}" + cookie = "\N{COOKIE}" + two = "\N{DIGIT TWO}" + NUM_ENC + flc = "\N{FOUR LEAF CLOVER}" + cyclone = "\N{CYCLONE}" + sunflower = "\N{SUNFLOWER}" + six = "\N{DIGIT SIX}" + NUM_ENC + mushroom = "\N{MUSHROOM}" + heart = "\N{HEAVY BLACK HEART}" + snowflake = "\N{SNOWFLAKE}" + + +PAYOUTS = { + (SMReel.two, SMReel.two, SMReel.six): { + "payout": lambda x: x * 2500 + x, + "phrase": "JACKPOT! 226! Your bid has been multiplied * 2500!" + }, + (SMReel.flc, SMReel.flc, SMReel.flc): { + "payout": lambda x: x + 1000, + "phrase": "4LC! +1000!" + }, + (SMReel.cherries, SMReel.cherries, SMReel.cherries): { + "payout": lambda x: x + 800, + "phrase": "Three cherries! +800!" + }, + (SMReel.two, SMReel.six): { + "payout": lambda x: x * 4 + x, + "phrase": "2 6! Your bid has been multiplied * 4!" + }, + (SMReel.cherries, SMReel.cherries): { + "payout": lambda x: x * 3 + x, + "phrase": "Two cherries! Your bid has been multiplied * 3!" + }, + "3 symbols": { + "payout": lambda x: x + 500, + "phrase": "Three symbols! +500!" + }, + "2 symbols": { + "payout": lambda x: x * 2 + x, + "phrase": "Two consecutive symbols! Your bid has been multiplied * 2!" + }, +} + +SLOT_PAYOUTS_MSG = ("Slot machine payouts:\n" + "{two.value} {two.value} {six.value} Bet * 2500\n" + "{flc.value} {flc.value} {flc.value} +1000\n" + "{cherries.value} {cherries.value} {cherries.value} +800\n" + "{two.value} {six.value} Bet * 4\n" + "{cherries.value} {cherries.value} Bet * 3\n\n" + "Three symbols: +500\n" + "Two symbols: Bet * 2".format(**SMReel.__dict__)) + + +def guild_only_check(): + async def pred(ctx: commands.Context): + if bank.is_global(): + return True + elif not bank.is_global() and ctx.guild is not None: + return True + else: + return False + return commands.check(pred) + + +class SetParser: + def __init__(self, argument): + allowed = ("+", "-") + self.sum = int(argument) + if argument and argument[0] in allowed: + if self.sum < 0: + self.operation = "withdraw" + elif self.sum > 0: + self.operation = "deposit" + else: + raise RuntimeError + self.sum = abs(self.sum) + elif argument.isdigit(): + self.operation = "set" + else: + raise RuntimeError + + +class Economy: + """Economy + + Get rich and have fun with imaginary currency!""" + + default_guild_settings = { + "PAYDAY_TIME": 300, + "PAYDAY_CREDITS": 120, + "SLOT_MIN": 5, + "SLOT_MAX": 100, + "SLOT_TIME": 0, + "REGISTER_CREDITS": 0 + } + + default_global_settings = default_guild_settings + + default_member_settings = { + "next_payday": 0, + "last_slot": 0 + } + + default_user_settings = default_member_settings + + def __init__(self, bot: Red): + self.bot = bot + self.file_path = "data/economy/settings.json" + self.config = Config.get_conf(self, 1256844281) + self.config.register_guild(**self.default_guild_settings) + self.config.register_global(**self.default_global_settings) + self.config.register_member(**self.default_member_settings) + self.config.register_user(**self.default_user_settings) + self.slot_register = defaultdict(dict) + + @commands.group(name="bank") + async def _bank(self, ctx: commands.Context): + """Bank operations""" + if ctx.invoked_subcommand is None: + await self.bot.send_cmd_help(ctx) + + @_bank.command() + async def balance(self, ctx: commands.Context, user: discord.Member = None): + """Shows balance of user. + + Defaults to yours.""" + if user is None: + user = ctx.author + + bal = bank.get_balance(user) + currency = bank.get_currency_name(ctx.guild) + + await ctx.send("{}'s balance is {} {}".format( + user.display_name, bal, currency)) + + @_bank.command() + async def transfer(self, ctx: commands.Context, to: discord.Member, amount: int): + """Transfer currency to other users""" + from_ = ctx.author + currency = bank.get_currency_name(ctx.guild) + + try: + await bank.transfer_credits(from_, to, amount) + except ValueError as e: + await ctx.send(str(e)) + + await ctx.send("{} transferred {} {} to {}".format( + from_.display_name, amount, currency, to.display_name + )) + + @_bank.command(name="set") + @check_global_setting_admin() + async def _set(self, ctx: commands.Context, to: discord.Member, creds: SetParser): + """Sets balance of user's bank account. See help for more operations + + Passing positive and negative values will add/remove currency instead + + Examples: + bank set @Twentysix 26 - Sets balance to 26 + bank set @Twentysix +2 - Increases balance by 2 + bank set @Twentysix -6 - Decreases balance by 6""" + author = ctx.author + currency = bank.get_currency_name(ctx.guild) + + if creds.operation == "deposit": + await bank.deposit_credits(to, creds.sum) + await ctx.send("{} added {} {} to {}'s account.".format( + author.display_name, creds.sum, currency, to.display_name + )) + elif creds.operation == "withdraw": + await bank.withdraw_credits(to, creds.sum) + await ctx.send("{} removed {} {} from {}'s account.".format( + author.display_name, creds.sum, currency, to.display_name + )) + else: + await bank.set_balance(to, creds.sum) + await ctx.send("{} set {}'s account to {} {}.".format( + author.display_name, to.display_name, creds.sum, currency + )) + + @_bank.command() + @guild_only_check() + @check_global_setting_guildowner() + async def reset(self, ctx, confirmation: bool = False): + """Deletes all guild's bank accounts""" + if confirmation is False: + await ctx.send( + "This will delete all bank accounts for {}.\nIf you're sure, type " + "{}bank reset yes".format( + self.bot.user.name if bank.is_global() else "this guild", + ctx.prefix + ) + ) + else: + if bank.is_global(): + # Bank being global means that the check would cause only + # the owner and any co-owners to be able to run the command + # so if we're in the function, it's safe to assume that the + # author is authorized to use owner-only commands + user = ctx.author + else: + user = ctx.guild.owner + success = await bank.wipe_bank(user) + if success: + await ctx.send("All bank accounts of this guild have been " + "deleted.") + + @commands.command() + @guild_only_check() + async def payday(self, ctx: commands.Context): + """Get some free currency""" + author = ctx.author + guild = ctx.guild + + cur_time = calendar.timegm(ctx.message.created_at.utctimetuple()) + credits_name = bank.get_currency_name(ctx.guild) + if bank.is_global(): + next_payday = self.config.user(author).next_payday() + if cur_time >= next_payday: + await bank.deposit_credits(author, self.config.PAYDAY_CREDITS()) + next_payday = cur_time + self.config.PAYDAY_TIME() + await self.config.user(author).next_payday.set(next_payday) + await ctx.send( + "{} Here, take some {}. Enjoy! (+{}" + " {}!)".format( + author.mention, credits_name, + str(self.config.PAYDAY_CREDITS()), + credits_name + ) + ) + else: + dtime = self.display_time(next_payday - cur_time) + await ctx.send( + "{} Too soon. For your next payday you have to" + " wait {}.".format(author.mention, dtime) + ) + else: + next_payday = self.config.member(author).next_payday() + if cur_time >= next_payday: + await bank.deposit_credits(author, self.config.guild(guild).PAYDAY_CREDITS()) + next_payday = cur_time + self.config.guild(guild).PAYDAY_TIME() + await self.config.member(author).next_payday.set(next_payday) + await ctx.send( + "{} Here, take some {}. Enjoy! (+{}" + " {}!)".format( + author.mention, credits_name, + str(self.config.guild(guild).PAYDAY_CREDITS()), + credits_name)) + else: + dtime = self.display_time(next_payday - cur_time) + await ctx.send( + "{} Too soon. For your next payday you have to" + " wait {}.".format(author.mention, dtime)) + + @commands.command() + @guild_only_check() + async def leaderboard(self, ctx: commands.Context, top: int = 10): + """Prints out the leaderboard + + Defaults to top 10""" + # Originally coded by Airenkun - edited by irdumb, rewritten by Palm__ for v3 + guild = ctx.guild + if top < 1: + top = 10 + if bank.is_global(): + bank_sorted = sorted(bank.get_global_accounts(ctx.author), + key=lambda x: x.balance, reverse=True) + else: + bank_sorted = sorted(bank.get_guild_accounts(guild), + key=lambda x: x.balance, reverse=True) + if len(bank_sorted) < top: + top = len(bank_sorted) + topten = bank_sorted[:top] + highscore = "" + place = 1 + for acc in topten: + dname = str(acc.name) + if len(dname) >= 23 - len(str(acc.balance)): + dname = dname[:(23 - len(str(acc.balance))) - 3] + dname += "... " + highscore += str(place).ljust(len(str(top)) + 1) + highscore += dname.ljust(23 - len(str(acc.balance))) + highscore += str(acc.balance) + "\n" + place += 1 + if highscore != "": + for page in pagify(highscore, shorten_by=12): + await ctx.send(box(page, lang="py")) + else: + await ctx.send("There are no accounts in the bank.") + + @commands.command() + @guild_only_check() + async def payouts(self, ctx: commands.Context): + """Shows slot machine payouts""" + await ctx.author.send(SLOT_PAYOUTS_MSG) + + @commands.command() + @guild_only_check() + async def slot(self, ctx: commands.Context, bid: int): + """Play the slot machine""" + author = ctx.author + guild = ctx.guild + channel = ctx.channel + if bank.is_global(): + valid_bid = self.config.SLOT_MIN() <= bid <= self.config.SLOT_MAX() + slot_time = self.config.SLOT_TIME() + last_slot = self.config.user(author).last_slot() + else: + valid_bid = self.config.guild(guild).SLOT_MIN() <= bid <= self.config.guild(guild).SLOT_MAX() + slot_time = self.config.guild(guild).SLOT_TIME() + last_slot = self.config.member(author).last_slot() + now = calendar.timegm(ctx.message.created_at.utctimetuple()) + + if (now - last_slot) < slot_time: + await ctx.send("You're on cooldown, try again in a bit.") + return + if not valid_bid: + await ctx.send("That's an invalid bid amount, sorry :/") + return + if not bank.can_spend(author, bid): + await ctx.send("You ain't got enough money, friend.") + return + if bank.is_global(): + await self.config.user(author).last_slot.set(now) + else: + await self.config.member(author).last_slot.set(now) + await self.slot_machine(author, channel, bid) + + async def slot_machine(self, author, channel, bid): + default_reel = deque(SMReel) + reels = [] + for i in range(3): + default_reel.rotate(random.randint(-999, 999)) # weeeeee + new_reel = deque(default_reel, maxlen=3) # we need only 3 symbols + reels.append(new_reel) # for each reel + rows = ((reels[0][0], reels[1][0], reels[2][0]), + (reels[0][1], reels[1][1], reels[2][1]), + (reels[0][2], reels[1][2], reels[2][2])) + + slot = "~~\n~~" # Mobile friendly + for i, row in enumerate(rows): # Let's build the slot to show + sign = " " + if i == 1: + sign = ">" + slot += "{}{} {} {}\n".format(sign, *[c.value for c in row]) + + payout = PAYOUTS.get(rows[1]) + if not payout: + # Checks for two-consecutive-symbols special rewards + payout = PAYOUTS.get((rows[1][0], rows[1][1]), + PAYOUTS.get((rows[1][1], rows[1][2]))) + if not payout: + # Still nothing. Let's check for 3 generic same symbols + # or 2 consecutive symbols + has_three = rows[1][0] == rows[1][1] == rows[1][2] + has_two = (rows[1][0] == rows[1][1]) or (rows[1][1] == rows[1][2]) + if has_three: + payout = PAYOUTS["3 symbols"] + elif has_two: + payout = PAYOUTS["2 symbols"] + + if payout: + then = bank.get_balance(author) + pay = payout["payout"](bid) + now = then - bid + pay + await bank.set_balance(author, now) + await channel.send("{}\n{} {}\n\nYour bid: {}\n{} → {}!" + "".format(slot, author.mention, + payout["phrase"], bid, then, now)) + else: + then = bank.get_balance(author) + await bank.withdraw_credits(author, bid) + now = then - bid + await channel.send("{}\n{} Nothing!\nYour bid: {}\n{} → {}!" + "".format(slot, author.mention, bid, then, now)) + + @commands.group() + @guild_only_check() + @check_global_setting_admin() + async def economyset(self, ctx: commands.Context): + """Changes economy module settings""" + guild = ctx.guild + if ctx.invoked_subcommand is None: + await self.bot.send_cmd_help(ctx) + if bank.is_global(): + slot_min = self.config.SLOT_MIN() + slot_max = self.config.SLOT_MAX() + slot_time = self.config.SLOT_TIME() + payday_time = self.config.PAYDAY_TIME() + payday_amount = self.config.PAYDAY_CREDITS() + else: + slot_min = self.config.guild(guild).SLOT_MIN() + slot_max = self.config.guild(guild).SLOT_MAX() + slot_time = self.config.guild(guild).SLOT_TIME() + payday_time = self.config.guild(guild).PAYDAY_TIME() + payday_amount = self.config.guild(guild).PAYDAY_CREDITS() + register_amount = bank.get_default_balance(guild) + msg = box( + "Minimum slot bid: {}\n" + "Maximum slot bid: {}\n" + "Slot cooldown: {}\n" + "Payday amount: {}\n" + "Payday cooldown: {}\n" + "Amount given at account registration: {}" + "".format( + slot_min, slot_max, slot_time, + payday_amount, payday_time, register_amount + ), + "Current Economy settings:" + ) + await ctx.send(msg) + + @economyset.command() + async def slotmin(self, ctx: commands.Context, bid: int): + """Minimum slot machine bid""" + if bid < 1: + await ctx.send('Invalid min bid amount.') + return + guild = ctx.guild + if bank.is_global(): + await self.config.SLOT_MIN.set(bid) + else: + await self.config.guild(guild).SLOT_MIN.set(bid) + credits_name = bank.get_currency_name(guild) + await ctx.send("Minimum bid is now {} {}.".format(bid, credits_name)) + + @economyset.command() + async def slotmax(self, ctx: commands.Context, bid: int): + """Maximum slot machine bid""" + slot_min = self.config.SLOT_MIN() + if bid < 1 or bid < slot_min: + await ctx.send('Invalid slotmax bid amount. Must be greater' + ' than slotmin.') + return + guild = ctx.guild + credits_name = bank.get_currency_name(guild) + if bank.is_global(): + await self.config.SLOT_MAX.set(bid) + else: + await self.config.guild(guild).SLOT_MAX.set(bid) + await ctx.send("Maximum bid is now {} {}.".format(bid, credits_name)) + + @economyset.command() + async def slottime(self, ctx: commands.Context, seconds: int): + """Seconds between each slots use""" + guild = ctx.guild + if bank.is_global(): + await self.config.SLOT_TIME.set(seconds) + else: + await self.config.guild(guild).SLOT_TIME.set(seconds) + await ctx.send("Cooldown is now {} seconds.".format(seconds)) + + @economyset.command() + async def paydaytime(self, ctx: commands.Context, seconds: int): + """Seconds between each payday""" + guild = ctx.guild + if bank.is_global(): + await self.config.PAYDAY_TIME.set(seconds) + else: + await self.config.guild(guild).PAYDAY_TIME.set(seconds) + await ctx.send("Value modified. At least {} seconds must pass " + "between each payday.".format(seconds)) + + @economyset.command() + async def paydayamount(self, ctx: commands.Context, creds: int): + """Amount earned each payday""" + guild = ctx.guild + credits_name = bank.get_currency_name(guild) + if creds <= 0: + await ctx.send("Har har so funny.") + return + if bank.is_global(): + await self.config.PAYDAY_CREDITS.set(creds) + else: + await self.config.guild(guild).PAYDAY_CREDITS.set(creds) + await ctx.send("Every payday will now give {} {}." + "".format(creds, credits_name)) + + @economyset.command() + async def registeramount(self, ctx: commands.Context, creds: int): + """Amount given on registering an account""" + guild = ctx.guild + if creds < 0: + creds = 0 + credits_name = bank.get_currency_name(guild) + await bank.set_default_balance(creds, guild) + await ctx.send("Registering an account will now give {} {}." + "".format(creds, credits_name)) + + # What would I ever do without stackoverflow? + def display_time(self, seconds, granularity=2): + intervals = ( # Source: http://stackoverflow.com/a/24542445 + ('weeks', 604800), # 60 * 60 * 24 * 7 + ('days', 86400), # 60 * 60 * 24 + ('hours', 3600), # 60 * 60 + ('minutes', 60), + ('seconds', 1), + ) + + result = [] + + for name, count in intervals: + value = seconds // count + if value: + seconds -= value * count + if value == 1: + name = name.rstrip('s') + result.append("{} {}".format(value, name)) + return ', '.join(result[:granularity]) diff --git a/core/bank.py b/core/bank.py new file mode 100644 index 000000000..3648f67bc --- /dev/null +++ b/core/bank.py @@ -0,0 +1,401 @@ +import datetime +from collections import namedtuple +from typing import Tuple, Generator, Union + +import discord +from copy import deepcopy + +from core import Config + +__all__ = ["get_balance", "set_balance", "withdraw_credits", "deposit_credits", + "can_spend", "transfer_credits", "wipe_bank", "get_guild_accounts", + "get_global_accounts", "get_account", "is_global", "get_bank_name", + "set_bank_name", "get_currency_name", "set_currency_name", + "get_default_balance", "set_default_balance"] + +_DEFAULT_GLOBAL = { + "is_global": False, + "bank_name": "Twentysix bank", + "currency": "credits", + "default_balance": 100 +} + +_DEFAULT_GUILD = { + "bank_name": "Twentysix bank", + "currency": "credits", + "default_balance": 100 +} + +_DEFAULT_MEMBER = { + "name": "", + "balance": 0, + "created_at": 0 +} + +_DEFAULT_USER = _DEFAULT_MEMBER + +_bank_type = type("Bank", (object,), {}) +Account = namedtuple("Account", "name balance created_at") + +_conf = Config.get_conf(_bank_type(), 384734293238749, force_registration=True) + + +def _register_defaults(): + _conf.register_global(**_DEFAULT_GLOBAL) + _conf.register_guild(**_DEFAULT_GUILD) + _conf.register_member(**_DEFAULT_MEMBER) + _conf.register_user(**_DEFAULT_USER) + + +_register_defaults() + + +def _encoded_current_time() -> int: + """ + Encoded current timestamp in UTC. + :return: + """ + now = datetime.datetime.utcnow() + return _encode_time(now) + + +def _encode_time(time: datetime.datetime) -> int: + """ + Goes from datetime object to serializable int. + :param time: + :return: + """ + ret = int(time.timestamp()) + return ret + + +def _decode_time(time: int) -> datetime.datetime: + """ + Returns decoded timestamp in UTC. + :param time: + :return: + """ + return datetime.datetime.utcfromtimestamp(time) + + +def get_balance(member: discord.Member) -> int: + """ + Gets the current balance of a member. + :param member: + :return: + """ + acc = get_account(member) + return acc.balance + + +def can_spend(member: discord.Member, amount: int) -> bool: + """ + Determines if a member can spend the given amount. + :param member: + :param amount: + :return: + """ + if _invalid_amount(amount): + return False + return get_balance(member) > amount + + +async def set_balance(member: discord.Member, amount: int) -> int: + """ + Sets an account balance. + + May raise ValueError if amount is invalid. + :param member: + :param amount: + :return: New account balance. + """ + if amount < 0: + raise ValueError("Not allowed to have negative balance.") + if is_global(): + group = _conf.user(member) + else: + group = _conf.member(member) + await group.balance.set(amount) + + if group.created_at() == 0: + time = _encoded_current_time() + await group.created_at.set(time) + + if group.name() == "": + await group.name.set(member.display_name) + + return amount + + +def _invalid_amount(amount: int) -> bool: + return amount <= 0 + + +async def withdraw_credits(member: discord.Member, amount: int) -> int: + """ + Removes a certain amount of credits from an account. + + May raise ValueError if the amount is invalid or if the account has + insufficient funds. + :param member: + :param amount: + :return: New account balance. + """ + if _invalid_amount(amount): + raise ValueError("Invalid withdrawal amount {} <= 0".format(amount)) + + bal = get_balance(member) + if amount > bal: + raise ValueError("Insufficient funds {} > {}".format(amount, bal)) + + return await set_balance(member, bal - amount) + + +async def deposit_credits(member: discord.Member, amount: int) -> int: + """ + Adds a given amount of credits to an account. + + May raise ValueError if the amount is invalid. + :param member: + :param amount: + :return: + """ + if _invalid_amount(amount): + raise ValueError("Invalid withdrawal amount {} <= 0".format(amount)) + + bal = get_balance(member) + return await set_balance(member, amount + bal) + + +async def transfer_credits(from_: discord.Member, to: discord.Member, amount: int): + """ + Transfers a given amount of credits from one account to another. + + May raise ValueError if the amount is invalid or if the from_ + account has insufficient funds. + :param from_: + :param to: + :param amount: + :return: + """ + if _invalid_amount(amount): + raise ValueError("Invalid transfer amount {} <= 0".format(amount)) + + await withdraw_credits(from_, amount) + return await deposit_credits(to, amount) + + +async def wipe_bank(user: Union[discord.User, discord.Member]): + """ + Deletes all accounts from the bank. + :return: + """ + if is_global(): + await _conf.user(user).clear() + else: + await _conf.member(user).clear() + + +def get_guild_accounts(guild: discord.Guild) -> Generator[Account, None, None]: + """ + Gets all account data for the given guild. + + May raise RuntimeError if the bank is currently global. + :param guild: + :return: + """ + if is_global(): + raise RuntimeError("The bank is currently global.") + + accs = _conf.member(guild.owner).all() + for user_id, acc in accs.items(): + acc_data = acc.copy() # There ya go kowlin + acc_data['created_at'] = _decode_time(acc_data['created_at']) + yield Account(**acc_data) + + +def get_global_accounts(user: discord.User) -> Generator[Account, None, None]: + """ + Gets all global account data. + + May raise RuntimeError if the bank is currently guild specific. + :param user: + :return: + """ + if not is_global(): + raise RuntimeError("The bank is not currently global.") + + accs = _conf.user(user).all() # this is a dict of user -> acc + for user_id, acc in accs.items(): + acc_data = acc.copy() + acc_data['created_at'] = _decode_time(acc_data['created_at']) + yield Account(**acc_data) + + +def get_account(member: Union[discord.Member, discord.User]) -> Account: + """ + Gets the appropriate account for the given member. + :param member: + :return: + """ + if is_global(): + acc_data = _conf.user(member)().copy() + default = _DEFAULT_USER.copy() + else: + acc_data = _conf.member(member)().copy() + default = _DEFAULT_MEMBER.copy() + + if acc_data == {}: + acc_data = default + acc_data['name'] = member.display_name + try: + acc_data['balance'] = get_default_balance(member.guild) + except AttributeError: + acc_data['balance'] = get_default_balance() + + acc_data['created_at'] = _decode_time(acc_data['created_at']) + return Account(**acc_data) + + +def is_global() -> bool: + """ + Determines if the bank is currently global. + :return: + """ + return _conf.is_global() + + +async def set_global(global_: bool, user: Union[discord.User, discord.Member]) -> bool: + """ + Sets global status of the bank, all accounts are reset when you switch! + :param global_: True will set bank to global mode. + :param user: Must be a Member object if changing TO global mode. + :return: New bank mode, True is global. + """ + if is_global() is global_: + return global_ + + if is_global(): + await _conf.user(user).clear_all() + elif isinstance(user, discord.Member): + await _conf.member(user).clear_all() + else: + raise RuntimeError("You must provide a member if you're changing to global" + " bank mode.") + + await _conf.is_global.set(global_) + return global_ + + +def get_bank_name(guild: discord.Guild=None) -> str: + """ + Gets the current bank name. If the bank is guild-specific the + guild parameter is required. + + May raise RuntimeError if guild is missing and required. + :param guild: + :return: + """ + if is_global(): + return _conf.bank_name() + elif guild is not None: + return _conf.guild(guild).bank_name() + else: + raise RuntimeError("Guild parameter is required and missing.") + + +async def set_bank_name(name: str, guild: discord.Guild=None) -> str: + """ + Sets the bank name, if bank is server specific the guild parameter is + required. + + May throw RuntimeError if guild is required and missing. + :param name: + :param guild: + :return: + """ + if is_global(): + await _conf.bank_name.set(name) + elif guild is not None: + await _conf.guild(guild).bank_name.set(name) + else: + raise RuntimeError("Guild must be provided if setting the name of a guild" + "-specific bank.") + return name + + +def get_currency_name(guild: discord.Guild=None) -> str: + """ + Gets the currency name of the bank. The guild parameter is required if + the bank is guild-specific. + + May raise RuntimeError if the guild is missing and required. + :param guild: + :return: + """ + if is_global(): + return _conf.currency() + elif guild is not None: + return _conf.guild(guild).currency() + else: + raise RuntimeError("Guild must be provided.") + + +async def set_currency_name(name: str, guild: discord.Guild=None) -> str: + """ + Sets the currency name for the bank, if bank is guild specific the + guild parameter is required. + + May raise RuntimeError if guild is missing and required. + :param name: + :param guild: + :return: + """ + if is_global(): + await _conf.currency.set(name) + elif guild is not None: + await _conf.guild(guild).currency.set(name) + else: + raise RuntimeError("Guild must be provided if setting the currency" + " name of a guild-specific bank.") + return name + + +def get_default_balance(guild: discord.Guild=None) -> int: + """ + Gets the current default balance amount. If the bank is guild-specific + you must pass guild. + + May raise RuntimeError if guild is missing and required. + :param guild: + :return: + """ + if is_global(): + return _conf.default_balance() + elif guild is not None: + return _conf.guild(guild).default_balance() + else: + raise RuntimeError("Guild is missing and required!") + + +async def set_default_balance(amount: int, guild: discord.Guild=None) -> int: + """ + Sets the default balance amount. Guild is required if the bank is + guild-specific. + + May raise RuntimeError if guild is missing and required. + May raise ValueError if amount is invalid. + :param guild: + :param amount: + :return: + """ + amount = int(amount) + if amount < 0: + raise ValueError("Amount must be greater than zero.") + + if is_global(): + await _conf.default_balance.set(amount) + elif guild is not None: + await _conf.guild(guild).default_balance.set(amount) + else: + raise RuntimeError("Guild is missing and required.") diff --git a/tests/cogs/test_economy.py b/tests/cogs/test_economy.py new file mode 100644 index 000000000..70bd6e98a --- /dev/null +++ b/tests/cogs/test_economy.py @@ -0,0 +1,55 @@ +import pytest + + +@pytest.fixture() +def bank(config): + from core import Config + Config.get_conf = lambda *args, **kwargs: config + + from core import bank + bank._register_defaults() + return bank + + +def test_bank_register(bank, ctx): + default_bal = bank.get_default_balance(ctx.guild) + assert default_bal == bank.get_account(ctx.author).balance + + +async def has_account(member, bank): + balance = bank.get_balance(member) + if balance == 0: + balance = 1 + await bank.set_balance(member, balance) + + +@pytest.mark.asyncio +async def test_bank_transfer(bank, member_factory): + mbr1 = member_factory.get() + mbr2 = member_factory.get() + bal1 = bank.get_account(mbr1).balance + bal2 = bank.get_account(mbr2).balance + await bank.transfer_credits(mbr1, mbr2, 50) + newbal1 = bank.get_account(mbr1).balance + newbal2 = bank.get_account(mbr2).balance + assert bal1 - 50 == newbal1 + assert bal2 + 50 == newbal2 + + +@pytest.mark.asyncio +async def test_bank_set(bank, member_factory): + mbr = member_factory.get() + await bank.set_balance(mbr, 250) + acc = bank.get_account(mbr) + assert acc.balance == 250 + + +@pytest.mark.asyncio +async def test_bank_can_spend(bank, member_factory): + mbr = member_factory.get() + canspend = bank.can_spend(mbr, 50) + assert canspend == (50 < bank.get_default_balance(mbr.guild)) + await bank.set_balance(mbr, 200) + acc = bank.get_account(mbr) + canspendnow = bank.can_spend(mbr, 100) + assert canspendnow diff --git a/tests/conftest.py b/tests/conftest.py index 5ecc9b178..f0ab65b9c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -87,13 +87,14 @@ def empty_role(): @pytest.fixture() def member_factory(guild_factory): - mock_member = namedtuple("Member", "id guild") + mock_member = namedtuple("Member", "id guild display_name") class MemberFactory: def get(self): return mock_member( random.randint(1, 999999999), - guild_factory.get()) + guild_factory.get(), + 'Testing_Name') return MemberFactory()