CopyBot/Red-DiscordBot
1
0
Fork 0
mirror of https://github.com/Cog-Creators/Red-DiscordBot.git synced 2025-11-09 20:58:54 -05:00
Code Issues Packages Projects Releases Wiki Activity

[Economy] [WIP] rewrite (#781)

Browse Source 
* [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
This commit is contained in:
palmtree5 2017-08-09 17:23:41 -08:00 committed by GitHub
parent 99bfb2fc7a
commit 4923ffe98a
8 changed files with 1102 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
cogs/bank/__init__.py Normal file
View File

@ -0,0 +1,5 @@
from .bank import Bank, check_global_setting_guildowner, check_global_setting_admin
def setup(bot):
bot.add_cog(Bank(bot))

67
cogs/bank/bank.py Normal file
View File

@ -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

37
cogs/bank/errors.py Normal file
View File

@ -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

6
cogs/economy/__init__.py Normal file
View File

@ -0,0 +1,6 @@
from .economy import Economy
from core.bot import Red
def setup(bot: Red):
bot.add_cog(Economy(bot))

528
cogs/economy/economy.py Normal file
View File

@ -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])

401
core/bank.py Normal file
View File

@ -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.")

55
tests/cogs/test_economy.py Normal file
View File

@ -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

5
tests/conftest.py
View File

@ -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()