change used IMDB library, update to centralize code

.gitignore

@@ -7,6 +7,9 @@
 # Config files
 config.yaml
 
+# test files
+*.sql
+
 # Byte-compiled / optimized / DLL files
 *__pycache__*
 *.pyc

movienightbot/commands/end_vote.py

@@ -21,19 +21,25 @@ 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
         winning_movies = vote_controller.end_vote(server_id)
     # 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

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
 

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))

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()

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)

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(" "):

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 = [