feat: Pagination for bubble tea count (#126)

* feat: pagination for bubble tea count - list * generic paginator, pagination to stats and leaderboard * add grouped embed back, refactored embed code to consolidate more logic in there * refactor bbt stats code
placeTW · Jan 12, 2025 · f91d41c · f91d41c
1 parent 28499f4
commit f91d41c
Show file tree

Hide file tree

Showing 4 changed files with 250 additions and 82 deletions.
diff --git a/commands/bbt_count/bbt_count.py b/commands/bbt_count/bbt_count.py
@@ -2,9 +2,9 @@
 import discord
 from discord import app_commands
 
-
 from bot import TWPlaceClient
 from commands.bbt_count.consts import BBT_LIST_GROUP_BY_CHOICES
+from components.paginator import GenericPaginator
 
 from .db_functions import (
     add_bbt_entry,
@@ -24,8 +24,8 @@
     user_transfer_embed,
     bbt_stats_embed,
 )
-from .helpers import bubble_tea_data
-from modules import logging, content_moderation
+from .helpers import bubble_tea_data, calculate_prices
+from modules import logging
 
 
 def register_commands(
@@ -311,23 +311,38 @@ async def bbt_count_list(
     ):
         await interaction.response.defer()
         entries = get_bbt_entries(user.id if user else interaction.user.id, year)
-        embed = (
-            bbt_list_default_embed(
+
+        if not entries:
+            embed = discord.Embed(
+                title="No bubble tea entries found",
+                color=discord.Color.red()
+            )
+            await interaction.followup.send(embed=embed)
+            return
+
+        total_prices = calculate_prices(entries, None)["default_group"]
+
+        if group_by:
+            embeds = bbt_list_grouped_embed(
                 user.id if user else interaction.user.id,
                 entries,
                 year,
                 interaction.created_at.astimezone().tzinfo,
+                group_by.value,
+                total_prices,
             )
-            if not group_by
-            else bbt_list_grouped_embed(
+        else:
+            embeds = bbt_list_default_embed(
                 user.id if user else interaction.user.id,
                 entries,
                 year,
                 interaction.created_at.astimezone().tzinfo,
-                group_by.value,
+                total_prices,
+                len(entries),
             )
-        )
-        await interaction.followup.send(embed=embed)
+
+        paginator = GenericPaginator(embeds)
+        await interaction.followup.send(embed=embeds[0], view=paginator)
 
     # leaderboard
     @bbt_count.command(
@@ -341,22 +356,40 @@ async def bbt_count_leaderboard(interaction: discord.Interaction, year: int = No
             interaction.guild.id,
             (datetime.datetime(year=year, month=1, day=1) if year else interaction.created_at),
         )
-        embed = discord.Embed(
-            title=f"Top bubble tea drinkers of {year if year else interaction.created_at.year} in {client.guilds_dict[interaction.guild.id]['server_name']}",
-            color=discord.Color.blue(),
-        )
-        embed.description = "\n".join(
-            [
-                f"{i+1}. <@{user_data.get('user_id')}>: {user_data.get('count')} 🧋"
-                + (
-                    f" *average given rating: {user_data.get('average_rating'):.3f}*"
-                    if user_data.get("average_rating")
-                    else ""
-                )
-                for i, user_data in enumerate(leaderboard)
-            ]
-        )
-        await interaction.followup.send(embed=embed)
+
+        if not leaderboard:
+            embed = discord.Embed(
+                title="No leaderboard data found",
+                color=discord.Color.red()
+            )
+            await interaction.followup.send(embed=embed)
+            return
+
+        # Split into chunks of 10 entries per page
+        chunks = [leaderboard[i:i + 10] for i in range(0, len(leaderboard), 10)]
+        embeds = []
+
+        for i, chunk in enumerate(chunks):
+            embed = discord.Embed(
+                title=f"Top bubble tea drinkers of {year if year else interaction.created_at.year} in {client.guilds_dict[interaction.guild.id]['server_name']}",
+                color=discord.Color.blue(),
+            )
+            embed.description = "\n".join(
+                [
+                    f"{j + (i*10) + 1}. <@{user_data.get('user_id')}>: {user_data.get('count')} 🧋"
+                    + (
+                        f" *average given rating: {user_data.get('average_rating'):.3f}*"
+                        if user_data.get("average_rating")
+                        else ""
+                    )
+                    for j, user_data in enumerate(chunk)
+                ]
+            )
+            embed.set_footer(text=f"Page {i+1} of {len(chunks)} • Total entries: {len(leaderboard)}")
+            embeds.append(embed)
+
+        paginator = GenericPaginator(embeds)
+        await interaction.followup.send(embed=embeds[0], view=paginator)
 
     @bbt_count.command(name="stats", description="Get the stats for a user for a given year")
     @app_commands.describe(user="User to get stats for (optional, default to self)")
@@ -375,25 +408,61 @@ async def bbt_count_stats(
             (datetime.datetime(year=year, month=1, day=1) if year else interaction.created_at),
             group_by_location,
         )
+
+        if not stats:
+            embed = discord.Embed(
+                title="No stats found",
+                description="No bubble tea entries found for the specified period.",
+                color=discord.Color.red(),
+            )
+            await interaction.followup.send(embed=embed)
+            return
+
         monthly_counts = get_bubble_tea_monthly_counts(
             user_id,
             (datetime.datetime(year=year, month=1, day=1) if year else interaction.created_at),
-        )
+        ) or []
 
         latest = (
             get_latest_bubble_tea_entry(user_id)
             if not year or (year and year == datetime.datetime.now().year)
             else None
         )
-        embed = bbt_stats_embed(
-            user_id,
-            stats,
-            year,
-            group_by_location,
-            monthly_counts,
-            latest,
-            interaction.created_at.astimezone().tzinfo,
-        )
-        await interaction.followup.send(embed=embed)
+
+        total_entries = sum([len(entry.get("prices_list", [])) for entry in stats])
+
+        if group_by_location:
+            chunks = [stats[i:i + 10] for i in range(0, len(stats), 10)]
+            embeds = []
+
+            for i, chunk in enumerate(chunks):
+                embed = bbt_stats_embed(
+                    user_id=user_id,
+                    entries=chunk,
+                    year=year,
+                    group_by_location=group_by_location,
+                    monthly_counts=None,
+                    latest=None,
+                    timezone=interaction.created_at.astimezone().tzinfo,
+                    total_entries=total_entries,
+                )
+                embed.set_footer(text=f"Page {i+1} of {len(chunks)} • Total locations: {len(stats)}")
+                embeds.append(embed)
+
+            paginator = GenericPaginator(embeds)
+            await interaction.followup.send(embed=embeds[0], view=paginator)
+        else:
+            embed = bbt_stats_embed(
+                user_id,
+                stats,
+                year,
+                group_by_location,
+                monthly_counts,
+                latest,
+                interaction.created_at.astimezone().tzinfo,
+                total_entries=total_entries,
+            )
+            await interaction.followup.send(embed=embed)
+
 
     tree.add_command(bbt_count, guilds=guilds)
diff --git a/commands/bbt_count/embeds.py b/commands/bbt_count/embeds.py
@@ -6,6 +6,7 @@
     bubble_tea_string,
     calculate_prices,
     entry_string,
+    organize_entries_by_group,
     price_string,
     cost_string,
     average_year_string,
@@ -75,48 +76,85 @@ def bbt_list_default_embed(
     entries: list[dict],
     year: int | None,
     timezone: datetime.tzinfo,
+    total_prices: dict = None,
+    total_entries: int = 0,
+    items_per_page: int = 10,
 ):
-    prices = calculate_prices(entries, None).get("default_group", {})
-    embed = discord.Embed(
-        title=f"Bubble tea entries {f'for {year}' if year else 'for the past year'} 🧋",
-        color=discord.Color.blue(),
-    )
-    embed.description = (
-        f"For <@{user_id}>: **{len(entries)} total entries**\n{average_year_string(year, len(entries))}"
-        + (
-            "\n\n__Total costs__:\n"
-            + "\n".join([cost_string(prices[currency]["prices"], currency) for currency in prices])
-            + "\n\n"
-            + "\n".join([entry_string(entry, timezone) for entry in entries])
+    chunks = [entries[i:i + items_per_page] for i in range(0, len(entries), items_per_page)]
+    embeds = []
+
+    for i, chunk in enumerate(chunks):
+        prices = calculate_prices(chunk, None).get("default_group", {})
+
+        embed = discord.Embed(
+            title=f"Bubble tea entries {f'for {year}' if year else 'for the past year'} 🧋",
+            color=discord.Color.blue(),
         )
-        if len(entries)
-        else ""
-    )
-    return embed
-
+        embed.description = (
+            f"For <@{user_id}>: **{total_entries} total entries**\n{average_year_string(year, total_entries)}"
+            + (
+                "\n\n__Total costs (all entries)__:\n"
+                + "\n".join([cost_string(total_prices[currency]["prices"], currency) for currency in total_prices])
+                + "\n\n__Current page entries:__\n"
+                + "\n".join([entry_string(entry, timezone) for entry in chunk])
+            )
+            if len(chunk)
+            else ""
+        )
+        embed.set_footer(text=f"Page {i+1}/{len(chunks)} • Total entries: {total_entries}")
+        embeds.append(embed)
+
+    return embeds
 
 def bbt_list_grouped_embed(
     user_id: int,
     entries: list[dict],
-    year: int,
+    year: int | None,
     timezone: datetime.tzinfo,
     group_by: str,
+    total_prices: dict = None,
+    items_per_page: int = 10,
 ):
-    prices = calculate_prices(entries, group_by)
-    embed = discord.Embed(
-        title=f"Bubble tea entries {f'for {year}' if year else 'for the past year'} grouped by {group_by} 🧋",
-        color=discord.Color.blue(),
-    )
-    embed.description = f"For <@{user_id}>: **{len(entries)} total entries**"
+    organized_data = organize_entries_by_group(entries, group_by)
+    embeds = []
+    total_entries = len(entries)
 
-    for group in prices:
-        group_entries = [entry for entry in entries if entry[group_by] == group]
-        embed.description += f"\n\n---\n**{group}: {len(group_entries)} entries**"
-        for currency in prices[group]:
-            embed.description += f"\n{currency}: {cost_string(prices[group][currency]['prices'], currency)}"
-        embed.description += "\n\n"
-        embed.description += "\n".join([entry_string(entry, timezone) for entry in group_entries])
-    return embed
+    for group_name, group_data in organized_data.items():
+        group_entries = group_data["entries"]
+        chunks = [group_entries[i:i + items_per_page] for i in range(0, len(group_entries), items_per_page)]
+
+        for i, chunk in enumerate(chunks):
+            embed = discord.Embed(
+                title=f"Bubble tea entries {f'for {year}' if year else 'for the past year'} grouped by {group_by} 🧋",
+                color=discord.Color.blue(),
+            )
+
+            # Add total information
+            embed.description = f"For <@{user_id}>: **{total_entries} total entries**"
+
+            if i == 0:  # Only for first page of each group
+                embed.description += "\n\n__Total costs (all entries)__:\n"
+                embed.description += "\n".join([
+                    cost_string(total_prices[currency]["prices"], currency) 
+                    for currency in total_prices
+                ])
+
+            # Add group information
+            embed.description += f"\n\n---\n**{group_name}: {len(group_entries)} entries**"
+            for currency in group_data["prices"]:
+                embed.description += f"\n{currency}: {cost_string(group_data['prices'][currency]['prices'], currency)}"
+
+            # Add entries for this chunk
+            embed.description += "\n\n"
+            embed.description += "\n".join([
+                entry_string(entry, timezone) 
+                for entry in chunk
+            ])
+
+            embed.set_footer(text=f"Group: {group_name} • Page {i+1}/{len(chunks)} • Total entries: {total_entries}")
+            embeds.append(embed)
+
+    return embeds
 
 
 def bbt_stats_embed(
@@ -127,16 +165,17 @@ def bbt_stats_embed(
     monthly_counts: list[dict],
     latest: dict,
     timezone: datetime.tzinfo,
+    total_entries: int = None,
 ):
-    total_count = sum([len(entry.get("prices_list", 0)) for entry in entries])
+    current_count = sum([len(entry.get("prices_list", [])) for entry in entries])
     embed = discord.Embed(
         title=f"Bubble tea stats {f'for {year}' if year else 'for the past year'} {'grouped by location ' if group_by_location else ''}🧋",
         color=discord.Color.green(),
     )
     embed.description = (
-        f"For <@{user_id}>: **{total_count} total entries**\n{average_year_string(year, total_count)}\n\n"
+        f"For <@{user_id}>: **{total_entries or current_count} total entries**\n{average_year_string(year, total_entries or current_count)}\n\n"
     )
-    if total_count > 0:
+    if current_count > 0:
         embed.description += f"{'__Total costs__' if not group_by_location else '__Costs by location__'}:\n"
         embed.description += "\n".join(
             [
@@ -152,16 +191,18 @@ def bbt_stats_embed(
                 for entry in entries
             ]
         )
-        current_year = year if year else datetime.datetime.now().year
-        embed.description += "\n\n__Monthly counts__:\n"
-        embed.description += "\n".join(
-            [
-                f"**{calendar.month_name[monthly_count.get('month', 0)]}**: {monthly_count.get('entry_count')} entries ({average_month_string(current_year, monthly_count.get('month', 0), monthly_count.get('entry_count'))})"
-                + ('\n- ' + rating_string(monthly_count) if monthly_count.get("average_rating") else "")
-                for monthly_count in monthly_counts
-            ]
-        )
-        if latest:
+
+        if monthly_counts and not group_by_location:
+            current_year = year if year else datetime.datetime.now().year
+            embed.description += "\n\n__Monthly counts__:\n"
+            embed.description += "\n".join(
+                [
+                    f"**{calendar.month_name[monthly_count.get('month', 0)]}**: {monthly_count.get('entry_count')} entries ({average_month_string(current_year, monthly_count.get('month', 0), monthly_count.get('entry_count'))})"
+                    + ('\n- ' + rating_string(monthly_count) if monthly_count.get("average_rating") else "")
+                    for monthly_count in monthly_counts
+                ]
+            )
+        if latest and not group_by_location:
             embed.description += "\n\n**Latest entry**:\n"
             embed.description += entry_string(latest, timezone)
     return embed