diff --git a/.gitignore b/.gitignore index 3fc6b49..d4cd205 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,9 @@ # Config files config.yaml +# test files +*.sql + # Byte-compiled / optimized / DLL files *__pycache__* *.pyc diff --git a/movienightbot/commands/end_vote.py b/movienightbot/commands/end_vote.py index ec734b2..d5c11b0 100644 --- a/movienightbot/commands/end_vote.py +++ b/movienightbot/commands/end_vote.py @@ -21,7 +21,8 @@ async def end_vote_task(interaction: discord.Interaction): server_id = interaction.guild.id with vote_controller.transaction(): try: - vote_msg_id = vote_controller.get_by_id(server_id).message_id + vote_obj = vote_controller.get_by_id(server_id) + vote_msg_id = vote_obj.message_id except DoesNotExist: await interaction.response.send_message("No vote started!") return @@ -29,11 +30,16 @@ async def end_vote_task(interaction: discord.Interaction): # TODO: Make more robust so we don't assume the end message and vote message are in same channel # probably safe for now, only happens if admin changes bot channel in the middle of a vote vote_msg = await get_message(interaction.channel, vote_msg_id) + logger.debug("vote_msg: {}", vote_msg) if vote_msg: await vote_msg.clear_reactions() else: # Vote message was deleted or is unavailable, so make a new one vote_msg = await interaction.channel.send("replacement vote message") + vote_obj.message_id = vote_msg.id + vote_obj.channel_id = vote_msg.channel.id + vote_obj.save() + if len(winning_movies) == 1: winning_movie = winning_movies[0].movie_name imdb_info = winning_movies[0].imdb_id diff --git a/movienightbot/commands/suggest.py b/movienightbot/commands/suggest.py index 2922a38..5d2ea43 100644 --- a/movienightbot/commands/suggest.py +++ b/movienightbot/commands/suggest.py @@ -2,7 +2,7 @@ from typing import Union import discord -import imdb +import imdbinfo.services from discord import app_commands from peewee import DoesNotExist, IntegrityError @@ -13,7 +13,8 @@ MoviesController, ServerController, ) -from movienightbot.util import capitalize_movie_name, get_imdb_info +from movienightbot.util import capitalize_movie_name +from movienightbot.imdb import get_imdb_info logger = logging.getLogger("movienightbot") @@ -23,32 +24,23 @@ genre_controller = GenreController() -def imdb_data(movie: str, kind: str) -> tuple[Union[None, IMDBInfo], Union[None, imdb.Movie.Movie]]: +def get_or_create_imdb_data(movie: str, kind: str) -> tuple[Union[None, IMDBInfo], Union[None, imdbinfo.services.MovieDetail]]: suggestion = capitalize_movie_name(movie) imdb_info = get_imdb_info(suggestion, kind=kind) if not imdb_info: return None, None # see if the row already exists try: - imdb_row = imdb_controller.get_by_id(imdb_info.movieID) + imdb_row = imdb_controller.get_by_id(imdb_info.imdb_id) except DoesNotExist: pass else: return imdb_row, imdb_info - # row doesn't exist, so add it - imdb_row_data = { - "imdb_id": imdb_info.movieID, - "title": imdb_info["title"], - "canonical_title": imdb_info.get("canonical title", imdb_info["title"]), - "year": imdb_info.get("year", 0), - "thumbnail_poster_url": imdb_info.get("cover url", ""), - "full_size_poster_url": imdb_info.get("full-size cover url", ""), - } try: - imdb_row = imdb_controller.create(imdb_row_data) + imdb_row = imdb_controller.create_by_imdb_id(imdb_info.imdb_id) except IntegrityError as e: - logger.error(f"IMDB entry insert error: {imdb_data}\n{e!s}") + logger.error(f"IMDB entry insert error: {imdb_info.imdb_id}\n{e!s}") return None, None return imdb_row, imdb_info @@ -71,7 +63,7 @@ async def suggest(interaction: discord.Interaction, movie: str): if server_row.check_movie_names: allow_tv_shows = server_row.allow_tv_shows kind = None if allow_tv_shows else "movie" - imdb_row, imdb_info = imdb_data(movie=movie, kind=kind) + imdb_row, imdb_info = get_or_create_imdb_data(movie=movie, kind=kind) suggestion = capitalize_movie_name(imdb_row.title) if imdb_row else capitalize_movie_name(movie) if imdb_row is None: await interaction.followup.send("Could not find the movie title you suggested in IMDb.", ephemeral=True) @@ -105,9 +97,9 @@ async def suggest(interaction: discord.Interaction, movie: str): if imdb_info: try: - add_genre_info(server_id, suggestion, imdb_info["genres"]) + add_genre_info(server_id, suggestion, imdb_info.genres) except IntegrityError as e: - logger.error(f"Genre insert error: {server_id} {imdb_info['genres']} {suggestion}\n{e}") + logger.error(f"Genre insert error: {server_id} {imdb_info.genres} {suggestion}\n{e}") await interaction.followup.send(f"Error adding suggestion {suggestion}") return diff --git a/movienightbot/db/controllers.py b/movienightbot/db/controllers.py index 8549823..ef2445d 100644 --- a/movienightbot/db/controllers.py +++ b/movienightbot/db/controllers.py @@ -6,11 +6,10 @@ import discord import peewee as pw -from imdb import IMDb -from imdb._exceptions import IMDbDataAccessError from ..exc import VoteError from . import BaseController +from ..imdb import get_imdb_info_by_id from .models import ( IMDBInfo, Movie, @@ -31,7 +30,23 @@ class ServerController(BaseController): class IMDBInfoController(BaseController): model = IMDBInfo - def get_by_id(self, imdb_id: int) -> Union[Vote, None]: + def create_by_imdb_id(self, imdb_id: str) -> Union[None, IMDBInfo]: + imdb_info = get_imdb_info_by_id(imdb_id) + if imdb_info is None: + return 0 + found_imdb_id = imdb_info.imdb_id + imdb_data = { + "imdb_id": found_imdb_id, + "title": imdb_info.title, + "canonical_title": imdb_info.title_localized if imdb_info.title_localized else imdb_info.title, + "year": imdb_info.year if imdb_info.year else 0, + "thumbnail_poster_url": imdb_info.cover_url if imdb_info.cover_url else "", + "full_size_poster_url": imdb_info.cover_url if imdb_info.cover_url else "", + } + imdb_controller = IMDBInfoController() + return imdb_controller.create(imdb_data) + + def get_by_id(self, imdb_id: str) -> Union[Vote, None]: return super().get_by_id(id=imdb_id, primary_key="imdb_id") def get_by_name(self, movie_name: str) -> Union[IMDBInfo, None]: @@ -88,42 +103,19 @@ def get_suggested_for_server(self, server_id: int) -> list[Movie]: ) return Movie.select().order_by(obc).where((Movie.server == server_id) & Movie.watched_on.is_null()).execute() - def get_imdb_info_by_id(self, imdb_id: Union[int, str]): - if not imdb_id: - return None - - im_db = IMDb() - try: - result = im_db.get_movie(imdb_id) - except IMDbDataAccessError: - return None - return result - - def update_imdb_id(self, server_id: int, movie_name: str, imdb_id: str): - imdb_info = self.get_imdb_info_by_id(imdb_id) - if imdb_info is None: - return 0 - - imdb_data = { - "imdb_id": imdb_info.movieID, - "title": imdb_info["title"], - "canonical_title": imdb_info.get("canonical title", imdb_info["title"]), - "year": imdb_info.get("year", 0), - "thumbnail_poster_url": imdb_info.get("cover url", ""), - "full_size_poster_url": imdb_info.get("full-size cover url", ""), - } + def update_imdb_id(self, server_id: int, movie_name: str, imdb_id: str) -> int: imdb_controller = IMDBInfoController() try: - imdb_controller.create(imdb_data) + imdb_row = imdb_controller.create_by_imdb_id(imdb_id) except pw.IntegrityError as e: # IMDB entry already added, so ignore error - logger.debug(f"IMDB entry insert error: {imdb_data}\n{e!s}") - try: - imdb_row = imdb_controller.get_by_id(imdb_info.movieID) - except Exception as e: - logger.debug(f"IMDB entry get error: {imdb_info.movieID}\n{e!s}") - return 0 - logger.debug("IMDB row: " + str(imdb_row)) + logger.debug(f"IMDB entry insert error: {imdb_id}\n{e!s}") + try: + imdb_row = imdb_controller.get_by_id(imdb_id) + except Exception as e: + logger.debug(f"IMDB entry get error: {imdb_id}\n{e!s}") + return 0 + logger.debug("IMDB row: {}", str(imdb_row)) return ( Movie.update({Movie.imdb_id: imdb_row}) .where((Movie.movie_name == movie_name) & (Movie.server == server_id)) diff --git a/movienightbot/db/models.py b/movienightbot/db/models.py index a66b625..f4b0b44 100644 --- a/movienightbot/db/models.py +++ b/movienightbot/db/models.py @@ -22,7 +22,7 @@ class Meta: class IMDBInfo(BaseModel): - imdb_id = pw.TextField(primary_key=True) + imdb_id = pw.TextField(primary_key=True) # WITHOUT the tt in front (need text because can start with 0) title = pw.TextField(null=False) canonical_title = pw.TextField() year = pw.IntegerField() diff --git a/movienightbot/imdb.py b/movienightbot/imdb.py new file mode 100644 index 0000000..f511c51 --- /dev/null +++ b/movienightbot/imdb.py @@ -0,0 +1,49 @@ +import logging +import re +from typing import Optional, Union + +import imdbinfo + + +logger = logging.getLogger("movienightbot") + +imdb_url_regex = re.compile(r"title/tt([0-9]+)") + +def get_imdb_info_by_id(imdb_id: Union[int, str]) -> Union[None, imdbinfo.services.MovieDetail]: + if not imdb_id: + return None + + return imdbinfo.get_movie(str(imdb_id)) + + +def get_imdb_info(movie_name: str, kind: Optional[str] = None, year: Optional[int] = None) -> Union[None, imdbinfo.services.MovieDetail]: + if not movie_name: + return None + + if movie_name.lower().startswith("http"): + movie_id = imdb_url_regex.findall(movie_name) + logger.debug("movie regex: `{}` >> {}", movie_name, movie_id) + if len(movie_id) == 1: + imdb_id = movie_id[0] + else: + return None + else: + logger.debug(f"searching for `{movie_name}`") + results = imdbinfo.search_title(movie_name) + logger.debug("IMDB RESULTS: {}", str(results)) + for r in results.titles: + if kind and kind != r.kind: + continue + if year and year != r.year: + continue + if r.title.lower() == movie_name.lower(): + logger.debug("{} Matched {}", movie_name, r) + # Cant use r directly because it is a "MovieBriefInfo" object + imdb_id = r.imdb_id + break + # for/else hell yeah! + else: + logger.debug("{} Unmatched", movie_name) + return None + + return get_imdb_info_by_id(imdb_id) \ No newline at end of file diff --git a/movienightbot/util.py b/movienightbot/util.py index a3b5630..5f83cf5 100644 --- a/movienightbot/util.py +++ b/movienightbot/util.py @@ -1,11 +1,9 @@ import asyncio import datetime import logging -import re -from typing import Optional, Union +from typing import Union import discord -import imdb import peewee as pw from .db.controllers import MovieVote, MovieVoteController, ServerController @@ -147,46 +145,12 @@ def build_vote_embed(server_id: int): emojis_unicode = {v: k for k, v in emojis_text.items()} -imdb_url_regex = re.compile(r"title/tt([0-9]+)") - - async def add_vote_emojis(vote_msg: discord.Message, movie_votes: MovieVote): for movie_vote in movie_votes: await vote_msg.add_reaction(emojis_text[movie_vote.emoji]) await vote_msg.add_reaction(emojis_text[":arrows_counterclockwise:"]) -def get_imdb_info(movie_name: str, kind: Optional[str] = None) -> Union[None, imdb.Movie.Movie]: - if not movie_name: - return None - - im_db = imdb.IMDb() - if movie_name.lower().startswith("http"): - movie_id = imdb_url_regex.findall(movie_name) - logger.debug(f"movie regex: `{movie_name}` >> {movie_id}") - if len(movie_id) == 1: - imdb_id = movie_id[0] - else: - return None - else: - logger.debug(f"searching for `{movie_name}`") - results = im_db.search_movie(movie_name) - logger.debug("IMDB RESULTS: " + str(results)) - for r in results: - if kind and kind not in r.get("kind", ""): - continue - if r["title"].lower() == movie_name.lower(): - logger.debug(f"{movie_name} Matched {r}") - imdb_id = r.movieID - break - # for/else hell yeah! - else: - logger.debug(movie_name + " Unmatched") - return None - - return im_db.get_movie(imdb_id) - - def capitalize_movie_name(movie_name: str) -> str: clean_name = [] for word in movie_name.strip().split(" "): diff --git a/setup.py b/setup.py index 95726cd..1aa8ee0 100644 --- a/setup.py +++ b/setup.py @@ -3,11 +3,11 @@ from movienightbot import __version__ as bot_version requirements = [ - "discord-py @ git+https://github.com/Rapptz/discord.py.git", + "discord-py>=2.7.1", "peewee", "marshmallow", "pyyaml", - "cinemagoer>=2022.12.27", + "imdbinfo>=0.8.2", ] test_requirements = [