Refactor shared gacha logic

commands/commands/rollcoin.py

@@ -479,6 +479,14 @@ async def appraise(
     return make_response(f"That post is very interesting, {sender_name}!\nI will purchase it for {value} RollCoins!\nThat brings your balance to {sender_wallet.balance}")
 
 
+def validate_pulls(pulls, balance):
+    if not (1 <= pulls <= 10):
+        RollbotFailure.INVALID_ARGUMENTS.raise_exc("Number of pulls must be between 1 and 10")
+
+    if balance < pulls:
+        return f"You only have {sender_wallet.balance} RollCoins available, and gacha pulls cost one each!"
+
+
 def pull_gacha_color():
     pull = random.randint(1, 100)
     if pull == 100: # 1%
@@ -515,6 +523,23 @@ async def pull_gacha(req, name_api_key, logger):
     return f"{name} ({color_info}) ({rarity_info})", img
 
 
+async def run_gacha_pulls(receiver_name, req, name_api_key, sleep_time, logger, pulls):
+    info, img = await pull_gacha(req, name_api_key, logger)
+    if pulls == 1:
+        yield f"{receiver_name} received...\n\t{info}", img
+        return
+    yield f"{receiver_name} received (1/{pulls})...\n\t{info}", img
+    for i in range(pulls - 1):
+        await asyncio.sleep(sleep_time)
+        try:
+            info, img = await pull_gacha(req, name_api_key, logger)
+            yield f"{receiver_name} received ({i + 2}/{pulls})...\n\t{info}", img
+        except:
+            logger.exception("Failed to pull")
+            yield f"Failed to pull from gachapon! You will not be charged the coin for this pull.",
+
+
+
 @as_command
 async def gacha(
     sender_wallet: SenderWallet,
@@ -534,33 +559,20 @@ async def gacha(
     Spend one RollCoin to pull from the gachapon! It dispenses the hottest new commodity - broken virtual pets from the 90s!
     These "Non-Functional Tamagotchis" or "NFTs" have varying rarities - good luck!
     """
-    if not (1 <= pulls <= 10):
-        RollbotFailure.INVALID_ARGUMENTS.raise_exc("Number of pulls must be between 1 and 10")
-
-    if sender_wallet.balance < pulls:
-        yield f"You only have {sender_wallet.balance} RollCoins available, and gacha pulls cost one each!", reply
+    if (invalid := validate_pulls(pulls, sender_wallet.balance)) is not None:
+        yield invalid, reply
         return
 
-    info, img = await pull_gacha(req, name_api_key, logger)
-    pulled = [info]
-    if pulls == 1:
-        yield f"You received...\n\t{info}", img, reply
-    else:
-        yield f"You received (1/{pulls})...\n\t{info}", img, reply
-        for i in range(pulls - 1):
-            await asyncio.sleep(sleep_time)
-            try:
-                info, img = await pull_gacha(req, name_api_key, logger)
-                yield f"You received ({i + 2}/{pulls})...\n\t{info}", img, reply
-                pulled.append(info)
-            except:
-                logger.exception("Failed to pull")
-                yield f"Failed to pull from gachapon! You will not be charged the coin for this pull.", reply
+    pulled = []
+    async for result in run_gacha_pulls("You", req, name_api_key, sleep_time, logger, pulls):
+        pulled.append(result[0])
+        yield *result, reply
     
     sender_wallet = await get_sender_wallet()
     sender_wallet.balance -= len(pulled)
     sender_wallet.nfts += pulled
     await wallet_store.save(sender_id, sender_wallet)
+
     state = await get_state()
     state.treasury += len(pulled)
     await state_store.save(GLOBAL_STATE_KEY, state)
@@ -586,40 +598,26 @@ async def giftcha(
     """
     Spend one RollCoin to pull from the gachapon, but selflessly gift the resulting NFT to someone else!
     """
-    if not (1 <= pulls <= 10):
-        RollbotFailure.INVALID_ARGUMENTS.raise_exc("Number of pulls must be between 1 and 10")
-
-    if sender_wallet.balance < pulls:
-        yield f"You only have {sender_wallet.balance} RollCoins available, and gacha pulls cost one each!", reply
+    if (invalid := validate_pulls(pulls, sender_wallet.balance)) is not None:
+        yield invalid, reply
         return
 
     if (target_id := wallet_lookup.get(target_name.lower(), None)) is None:
         RollbotFailure.INVALID_ARGUMENTS.raise_exc(f"Could not find a wallet for {target_name}")
 
-    target_name = target_name.title()
-
-    info, img = await pull_gacha(req, name_api_key, logger)
-    pulled = [info]
-    if pulls == 1:
-        yield f"{target_name} received...\n\t{info}", img, reply
-    else:
-        yield f"{target_name} received (1/{pulls})...\n\t{info}", img, reply
-        for i in range(pulls - 1):
-            await asyncio.sleep(sleep_time)
-            try:
-                info, img = await pull_gacha(req, name_api_key, logger)
-                yield f"{target_name} received ({i + 2}/{pulls})...\n\t{info}", img, reply
-                pulled.append(info)
-            except:
-                logger.exception("Failed to pull")
-                yield f"Failed to pull from gachapon! You will not be charged the coin for this pull.", reply
+    pulled = []
+    async for result in run_gacha_pulls(target_name.title(), req, name_api_key, sleep_time, logger, pulls):
+        pulled.append(result[0])
+        yield result
     
     sender_wallet = await get_sender_wallet()
     sender_wallet.balance -= len(pulled)
     await wallet_store.save(sender_id, sender_wallet)
+
     target_wallet = await wallet_store.load_or(target_id)
     target_wallet.nfts += pulled
     await wallet_store.save(target_id, target_wallet)
+
     state = await get_state()
     state.treasury += len(pulled)
     await state_store.save(GLOBAL_STATE_KEY, state)