include PR Cog-Creators#5325 from Yami

Author: vertyco

docs/framework_events.rst

@@ -26,3 +26,54 @@ Economy
     :param payload.new_balance: The new balance of the user after the payday claim. (:class:`int`)
     :method payload.to_dict: Returns a serializable dictionary representation of the payload. (:class:`dict`)
     :method payload.to_json: Returns the payload as JSON. (:class:`str`)
+
+
+Bank
+^^^^
+
+.. py:method:: red_bank_set_balance(payload: redbot.core.bank.BankSetBalanceInformation)
+
+    Dispatched when a user's balance is changed.
+
+    :type payload: redbot.core.bank.BankSetBalanceInformation
+    :param payload.recipient: The member whose balance is being set. (:class:`discord.Member`, :class:`discord.User`)
+    :param payload.guild: The guild where the balance is being set. (:class:`discord.Guild`, :class:`NoneType`)
+    :param payload.recipient_old_balance: The old balance of the user before the change. (:class:`int`)
+    :param payload.recipient_new_balance: The new balance of the user after the change. (:class:`int`)
+    :method payload.to_dict: Returns a serializable dictionary representation of the payload. (:class:`dict`)
+    :method payload.to_json: Returns the payload as JSON. (:class:`str`)
+
+.. py:method:: red_bank_transfer_credits(payload: redbot.core.bank.BankTransferInformation)
+
+    Dispatched when a user transfers currency to another user.
+
+    :type payload: redbot.core.bank.BankTransferInformation
+    :param payload.sender: The member who is sending currency. (:class:`discord.Member`, :class:`discord.User`)
+    :param payload.recipient: The member who is receiving currency. (:class:`discord.Member`, :class:`discord.User`)
+    :param payload.guild: The guild where the transfer is taking place. (:class:`discord.Guild`, :class:`NoneType`)
+    :param payload.transfer_amount: The amount of currency being transferred. (:class:`int`)
+    :param payload.sender_new_balance: The new balance of the sender after the transfer. (:class:`int`)
+    :param payload.recipient_new_balance: The new balance of the recipient after the transfer. (:class:`int`)
+    :method payload.to_dict: Returns a serializable dictionary representation of the payload. (:class:`dict`)
+    :method payload.to_json: Returns the payload as JSON. (:class:`str`)
+
+
+.. py:method:: red_bank_prune(payload: redbot.core.bank.BankPruneInformation)
+
+    Dispatched when a user is pruned from the bank.
+
+    :type payload: redbot.core.bank.BankPruneInformation
+    :param payload.guild: The guild where the user is being pruned. (:class:`discord.Guild`, :class:`NoneType`)
+    :param payload.user_id: The ID of the user being pruned. (:class:`int`, :class:`NoneType`)
+    :param payload.pruned_users: Dict of pruned user accounts {user_id: {name: str, balance: int, created_at: int}}. (:class:`dict`)
+    :method payload.to_dict: Returns a serializable dictionary representation of the payload. (:class:`dict`)
+    :method payload.to_json: Returns the payload as JSON. (:class:`str`)
+
+
+.. py:method:: red_bank_wipe(guild_id: int)
+
+    Dispatched when a guild's bank is wiped. :code:`guild_id` will be the ID of the Guild that was wiped, -1 if all users were wiped (global bank), or None if all Guilds were wiped (local bank).
+
+.. py:method:: red_bank_set_global(global_state: bool)
+
+    Dispatched when the global bank is enabled or disabled. :code:`global_state` will be True if the Bank is being set to Global or False if the bank is being set to Local

redbot/core/bank.py

@@ -3,7 +3,7 @@
 import asyncio
 import logging
 from datetime import datetime, timezone
-from typing import Union, List, Optional, TYPE_CHECKING, Literal, NamedTuple
+from typing import Union, List, Optional, TYPE_CHECKING, Literal, NamedTuple, Dict
 from functools import wraps
 import json
 
@@ -47,6 +47,10 @@
     "set_default_balance",
     "AbortPurchase",
     "cost",
+    "PaydayClaimInformation",
+    "BankSetBalanceInformation",
+    "BankTransferInformation",
+    "BankPruneInformation",
 )
 
 _MAX_BALANCE = 2**63 - 1
@@ -82,9 +86,13 @@
 _cache_is_global = None
 _cache = {"bank_name": None, "currency": None, "default_balance": None, "max_balance": None}
 
+_bot_ref: Optional[Red] = None
 
-async def _init():
+
+async def _init(bot: Red):
     global _config
+    global _bot_ref
+    _bot_ref = bot
     _config = Config.get_conf(None, 384734293238749, cog_name="Bank", force_registration=True)
     _config.register_global(**_DEFAULT_GLOBAL)
     _config.register_guild(**_DEFAULT_GUILD)
@@ -337,6 +345,7 @@ async def set_balance(member: Union[discord.Member, discord.User], amount: int)
         group = _config.user(member)
     else:
         group = _config.member(member)
+    old_balance = await group.balance()
     await group.balance.set(amount)
 
     if await group.created_at() == 0:
@@ -345,7 +354,8 @@ async def set_balance(member: Union[discord.Member, discord.User], amount: int)
 
     if await group.name() == "":
         await group.name.set(member.display_name)
-
+    payload = BankSetBalanceInformation(member, guild, old_balance, amount)
+    _bot_ref.dispatch("red_bank_set_balance", payload)
     return amount
 
 
@@ -484,8 +494,11 @@ async def transfer_credits(
             user=to.display_name, max_balance=max_bal, currency_name=currency
         )
 
-    await withdraw_credits(from_, amount)
-    return await deposit_credits(to, amount)
+    sender_new = await withdraw_credits(from_, amount)
+    recipient_new = await deposit_credits(to, amount)
+    payload = BankTransferInformation(from_, to, guild, amount, sender_new, recipient_new)
+    _bot_ref.dispatch("red_bank_transfer_credits", payload)
+    return recipient_new
 
 
 async def wipe_bank(guild: Optional[discord.Guild] = None) -> None:
@@ -500,8 +513,10 @@ async def wipe_bank(guild: Optional[discord.Guild] = None) -> None:
     """
     if await is_global():
         await _config.clear_all_users()
+        _bot_ref.dispatch("red_bank_wipe", -1)
     else:
         await _config.clear_all_members(guild)
+        _bot_ref.dispatch("red_bank_wipe", getattr(guild, "id", None))
 
 
 async def bank_prune(bot: Red, guild: discord.Guild = None, user_id: int = None) -> None:
@@ -524,46 +539,49 @@ async def bank_prune(bot: Red, guild: discord.Guild = None, user_id: int = None)
         If guild is :code:`None` and the bank is Local.
 
     """
-
     global_bank = await is_global()
+    if not global_bank and guild is None:
+        raise BankPruneError("'guild' can't be None when pruning a local bank")
 
+    _guilds = set()
+    _uguilds = set()
     if global_bank:
-        _guilds = set()
-        _uguilds = set()
+        group = _config._get_base_group(_config.USER)
         if user_id is None:
             async for g in AsyncIter(bot.guilds, steps=100):
                 if not g.unavailable and g.large and not g.chunked:
                     _guilds.add(g)
                 elif g.unavailable:
                     _uguilds.add(g)
-        group = _config._get_base_group(_config.USER)
 
     else:
-        if guild is None:
-            raise BankPruneError("'guild' can't be None when pruning a local bank")
-        if user_id is None:
-            _guilds = {guild} if not guild.unavailable and guild.large else set()
-            _uguilds = {guild} if guild.unavailable else set()
         group = _config._get_base_group(_config.MEMBER, str(guild.id))
-
-    if user_id is None:
-        for _guild in _guilds:
-            await _guild.chunk()
-        accounts = await group.all()
-        tmp = accounts.copy()
-        members = bot.get_all_members() if global_bank else guild.members
-        user_list = {str(m.id) for m in members if m.guild not in _uguilds}
-
-    async with group.all() as bank_data:  # FIXME: use-config-bulk-update
         if user_id is None:
-            for acc in tmp:
-                if acc not in user_list:
-                    del bank_data[acc]
+            if guild.large and not guild.unavailable:
+                _guilds.add(guild)
+            if guild.unavailable:
+                _uguilds.add(guild)
+
+    pruned = {}
+    async with group.all() as bank_data:
+        if user_id in None:
+            for _guild in _guilds:
+                await _guild.chunk()
+            members = bot.get_all_members() if global_bank else guild.members
+            valid_users = {str(m.id) for m in members if m.guild not in _uguilds}
+            for account_user_id, account in bank_data.copy().items():
+                if account_user_id not in valid_users:
+                    pruned[account_user_id] = account
+                    del bank_data[account_user_id]
         else:
             user_id = str(user_id)
             if user_id in bank_data:
+                pruned[user_id] = bank_data[user_id]
                 del bank_data[user_id]
 
+    payload = BankPruneInformation(guild, user_id, pruned)
+    _bot_ref.dispatch("red_bank_prune", payload)
+
 
 async def get_leaderboard(positions: int = None, guild: discord.Guild = None) -> List[tuple]:
     """
@@ -725,11 +743,14 @@ async def set_global(global_: bool) -> bool:
 
     if await is_global():
         await _config.clear_all_users()
+        _bot_ref.dispatch("red_bank_wipe", -1)
     else:
         await _config.clear_all_members()
+        _bot_ref.dispatch("red_bank_wipe", None)
 
     await _config.is_global.set(global_)
     _cache_is_global = global_
+    _bot_ref.dispatch("red_bank_set_global", global_)
     return global_
 
 
@@ -1106,3 +1127,69 @@ def to_dict(self) -> dict:
 
     def to_json(self) -> str:
         return json.dumps(self.to_dict())
+
+
+class BankSetBalanceInformation(NamedTuple):
+    recipient: Union[discord.Member, discord.User]
+    guild: Union[discord.Guild, None]
+    recipient_old_balance: int
+    recipient_new_balance: int
+
+    def to_dict(self) -> dict:
+        return {
+            "recipient": self.recipient.id,
+            "guild": getattr(self.guild, "id", None),
+            "recipient_old_balance": self.recipient_old_balance,
+            "recipient_new_balance": self.recipient_new_balance,
+        }
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict())
+
+
+class BankTransferInformation(NamedTuple):
+    sender: Union[discord.Member, discord.User]
+    recipient: Union[discord.Member, discord.User]
+    guild: Union[discord.Guild, None]
+    transfer_amount: int
+    sender_new_balance: int
+    recipient_new_balance: int
+
+    def to_dict(self) -> dict:
+        return {
+            "sender": self.sender.id,
+            "recipient": self.recipient.id,
+            "guild": getattr(self.guild, "id", None),
+            "transfer_amount": self.transfer_amount,
+            "sender_new_balance": self.sender_new_balance,
+            "recipient_new_balance": self.recipient_new_balance,
+        }
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict())
+
+
+class BankPruneInformation(NamedTuple):
+    guild: Union[discord.Guild, None]
+    user_id: Union[int, None]
+    # {user_id: {name: str, balance: int, created_at: int}}
+    pruned_users: Dict[str, Dict[str, Union[int, str]]]
+
+    @property
+    def scope(self) -> str:
+        if self.guild is None and self.user_id is None:
+            return "global"
+        elif self.guild is not None and self.user_id is None:
+            return "guild"
+        return "user"
+
+    def to_dict(self) -> dict:
+        return {
+            "guild": getattr(self.guild, "id", None),
+            "user_id": self.user_id,
+            "scope": self.scope,
+            "pruned_users": self.pruned_users,
+        }
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict())

redbot/core/bot.py

@@ -1154,7 +1154,7 @@ async def _pre_connect(self) -> None:
             await self.add_cog(Dev())
 
         await modlog._init(self)
-        await bank._init()
+        await bank._init(self)
 
         packages = OrderedDict()
 

redbot/pytest/economy.py

@@ -5,11 +5,11 @@
 
 
 @pytest.fixture()
-async def bank(config, monkeypatch):
+async def bank(config, monkeypatch, red):
     from redbot.core import Config
 
     with monkeypatch.context() as m:
         m.setattr(Config, "get_conf", lambda *args, **kwargs: config)
         # noinspection PyProtectedMember
-        await bank_module._init()
+        await bank_module._init(red)
         return bank_module