diff --git a/main.py b/main.py index 7f40d50..6697e61 100644 --- a/main.py +++ b/main.py @@ -5,6 +5,8 @@ import discord from d4dj_utils.manager.asset_manager import AssetManager from discord.ext import commands +from miyu_bot.commands.common.master_asset_manager import MasterFilterManager + logging.basicConfig(level=logging.INFO) with open('config.json') as f: @@ -14,6 +16,8 @@ bot = commands.Bot(command_prefix='!', case_insensitive=True) asset_manager = AssetManager('assets') +masters = MasterFilterManager(asset_manager) + bot.load_extension('miyu_bot.commands.cogs.card') bot.load_extension('miyu_bot.commands.cogs.event') bot.load_extension('miyu_bot.commands.cogs.music') diff --git a/miyu_bot/commands/cogs/event.py b/miyu_bot/commands/cogs/event.py index a12cea4..d1e70bc 100644 --- a/miyu_bot/commands/cogs/event.py +++ b/miyu_bot/commands/cogs/event.py @@ -7,25 +7,19 @@ from d4dj_utils.master.event_master import EventMaster, EventState from discord.ext import commands from pytz import UnknownTimeZoneError -from main import asset_manager +from main import asset_manager, masters from miyu_bot.commands.common.emoji import attribute_emoji_ids_by_attribute_id, unit_emoji_ids_by_unit_id, \ parameter_bonus_emoji_ids_by_parameter_id, \ event_point_emoji_id from miyu_bot.commands.common.formatting import format_info from miyu_bot.commands.common.fuzzy_matching import romanize -from miyu_bot.commands.common.master_asset_manager import MasterAssetManager +from miyu_bot.commands.common.master_asset_manager import MasterFilter class Event(commands.Cog): def __init__(self, bot: commands.Bot): self.bot = bot self.logger = logging.getLogger(__name__) - self.events = MasterAssetManager( - asset_manager.event_master, - naming_function=lambda e: e.name, - filter_function=lambda e: e.start_datetime < dt.datetime.now( - dt.timezone.utc) + dt.timedelta(hours=12), - ) @commands.command(name='event', aliases=['ev'], @@ -36,7 +30,7 @@ class Event(commands.Cog): event: EventMaster if arg: - event = self.events.get(arg, ctx) + event = masters.events.get(arg, ctx) else: event = self.get_latest_event(ctx) @@ -51,7 +45,7 @@ class Event(commands.Cog): logo = discord.File(event.logo_path, filename='logo.png') except FileNotFoundError: # Just a fallback - logo = discord.File(self.events.get('1', ctx).logo_path, filename='logo.png') + logo = discord.File(masters.events.get('1', ctx).logo_path, filename='logo.png') embed = discord.Embed(title=event.name) embed.set_thumbnail(url=f'attachment://logo.png') @@ -191,14 +185,14 @@ class Event(commands.Cog): """Returns the oldest event that has not ended or the newest event otherwise.""" try: # NY event overlapped with previous event - return min((v for v in self.events.values(ctx) if v.state() == EventState.Open), + return min((v for v in masters.events.values(ctx) if v.state() == EventState.Open), key=lambda e: e.start_datetime) except ValueError: try: - return min((v for v in self.events.values(ctx) if v.state() < EventState.Ended), + return min((v for v in masters.events.values(ctx) if v.state() < EventState.Ended), key=lambda e: e.start_datetime) except ValueError: - return max(self.events.values(ctx), key=lambda v: v.start_datetime) + return max(masters.events.values(ctx), key=lambda v: v.start_datetime) def setup(bot): diff --git a/miyu_bot/commands/cogs/music.py b/miyu_bot/commands/cogs/music.py index 2aee918..f354827 100644 --- a/miyu_bot/commands/cogs/music.py +++ b/miyu_bot/commands/cogs/music.py @@ -2,7 +2,7 @@ import asyncio import contextlib import logging import wave -from typing import Optional, Tuple +from typing import Tuple import discord from d4dj_utils.master.chart_master import ChartDifficulty, ChartMaster @@ -10,11 +10,10 @@ from d4dj_utils.master.common_enums import ChartSectionType from d4dj_utils.master.music_master import MusicMaster from discord.ext import commands -from main import asset_manager +from main import asset_manager, masters from miyu_bot.commands.common.emoji import difficulty_emoji_ids from miyu_bot.commands.common.formatting import format_info -from miyu_bot.commands.common.fuzzy_matching import romanize, FuzzyFilteredMap -from miyu_bot.commands.common.master_asset_manager import MasterAssetManager +from miyu_bot.commands.common.fuzzy_matching import romanize from miyu_bot.commands.common.reaction_message import run_tabbed_message, run_paged_message @@ -22,12 +21,6 @@ class Music(commands.Cog): def __init__(self, bot: commands.Bot): self.bot = bot self.logger = logging.getLogger(__name__) - self.music = MasterAssetManager( - asset_manager.music_master, - naming_function=lambda m: f'{m.name} {m.special_unit_name}', - filter_function=lambda m: m.is_released, - fallback_naming_function=lambda m: m.id, - ) @property def reaction_emojis(self): @@ -55,7 +48,7 @@ class Music(commands.Cog): async def song(self, ctx: commands.Context, *, arg: str): self.logger.info(f'Searching for song "{arg}".') - song = self.music.get(arg, ctx) + song = masters.music.get(arg, ctx) if not song: msg = f'Failed to find song "{arg}".' @@ -68,7 +61,7 @@ class Music(commands.Cog): thumb = discord.File(song.jacket_path, filename='jacket.png') except FileNotFoundError: # dig delight is just a fallback - thumb = discord.File(self.music.get('110001', ctx).jacket_path, filename='jacket.png') + thumb = discord.File(masters.music.get('110001', ctx).jacket_path, filename='jacket.png') embed = discord.Embed(title=song.name) embed.set_thumbnail(url=f'attachment://jacket.png') @@ -108,7 +101,7 @@ class Music(commands.Cog): self.logger.info(f'Searching for chart "{arg}".') name, difficulty = self.parse_chart_args(arg) - song = self.music.get(name, ctx) + song = masters.music.get(name, ctx) if not song: msg = f'Failed to find chart "{name}".' @@ -131,7 +124,7 @@ class Music(commands.Cog): self.logger.info(f'Searching for chart sections "{arg}".') name, difficulty = self.parse_chart_args(arg) - song = self.music.get(name, ctx) + song = masters.music.get(name, ctx) if not song: msg = f'Failed to find chart "{name}".' @@ -163,7 +156,7 @@ class Music(commands.Cog): elif arg == 'duration': sort = 'duration' arg = '' - songs = self.music.get_sorted(arg, ctx) + songs = masters.music.get_sorted(arg, ctx) if sort == 'relevance': listing = [f'{song.name}{" (" + song.special_unit_name + ")" if song.special_unit_name else ""}' for song in songs] @@ -187,7 +180,7 @@ class Music(commands.Cog): thumb = discord.File(song.jacket_path, filename='jacket.png') except FileNotFoundError: # dig delight is just a fallback - thumb = discord.File(self.music.get('110001', None).jacket_path, filename='jacket.png') + thumb = discord.File(masters.music.get('110001', None).jacket_path, filename='jacket.png') files = [thumb] diff --git a/miyu_bot/commands/common/master_asset_manager.py b/miyu_bot/commands/common/master_asset_manager.py index ecc3544..4747a33 100644 --- a/miyu_bot/commands/common/master_asset_manager.py +++ b/miyu_bot/commands/common/master_asset_manager.py @@ -1,12 +1,39 @@ +from functools import cached_property from typing import Callable, Any, Optional +from d4dj_utils.manager.asset_manager import AssetManager from d4dj_utils.master.master_asset import MasterDict from discord.ext import commands from miyu_bot.commands.common.fuzzy_matching import FuzzyFilteredMap +import datetime as dt -class MasterAssetManager: + +class MasterFilterManager: + def __init__(self, manager: AssetManager): + self.manager = manager + + @cached_property + def music(self): + return MasterFilter( + self.manager.music_master, + naming_function=lambda m: f'{m.name} {m.special_unit_name}', + filter_function=lambda m: m.is_released, + fallback_naming_function=lambda m: m.id, + ) + + @cached_property + def events(self): + return MasterFilter( + self.manager.event_master, + naming_function=lambda e: e.name, + filter_function=lambda e: e.start_datetime < dt.datetime.now( + dt.timezone.utc) + dt.timedelta(hours=12), + ) + + +class MasterFilter: def __init__(self, masters: MasterDict, naming_function: Callable[[Any], str], filter_function=lambda _: True, fallback_naming_function: Optional[Callable[[Any], str]] = None): self.masters = masters