wi1-bot 1.2.12__tar.gz → 1.3.1__tar.gz

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/.gitignore

@@ -1,5 +1,7 @@
 config.yaml
+/tests/files/
+
+build/
 dist/
 *.egg-info/
 _version.py
-/tests/files/

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/.pre-commit-config.yaml

@@ -24,6 +24,6 @@ repos:
       - id: mypy
         name: mypy
         language: system
-        entry: mypy -p wi1_bot
+        entry: mypy
         pass_filenames: false
         types: [python]

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: wi1-bot
-Version: 1.2.12
+Version: 1.3.1
 Summary: Discord bot for Radarr/Sonarr integration
 Home-page: https://github.com/wthueb/wi1-bot
 Author-email: William Huebner <wilhueb@gmail.com>
@@ -57,6 +57,8 @@ Requires Python >=3.10.
 
 ### TODO
 
+- Use `mypy --strict`
+- Use Discord slash commands instead of normal text commands
 - Web dashboard for seeing transcode queue, transcode progress, quotas
 - Enforce quotas
 - !linktmdb

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/README.md

@@ -19,6 +19,8 @@ Requires Python >=3.10.
 
 ### TODO
 
+- Use `mypy --strict`
+- Use Discord slash commands instead of normal text commands
 - Web dashboard for seeing transcode queue, transcode progress, quotas
 - Enforce quotas
 - !linktmdb

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/config.yaml.template

@@ -4,12 +4,16 @@ radarr:
   url: http://localhost:7878
   # radarr api key (Settings->General->Security->API Key)
   api_key: XXX
+  # radarr root folder (absolute path)
+  root_folder: /full/path/movies
 
 sonarr:
   # sonarr url you use to get to the dashboard
   url: http://localhost:8989
   # sonarr api key (Settings->General->Security->API Key)
   api_key: XXX
+  # sonarr root folder (absolute path)
+  root_folder: /full/path/shows
 
 discord:
   # discord bot token
@@ -45,6 +49,8 @@ transcoding:
   profiles:
     # name of profile must match name in radarr/sonarr
     good:
+      # -map 0 in ffmpeg, optional
+      copy_all_streams: true
       # -c:v in ffmpeg, optional
       video_codec: hevc_nvenc
       # -b:v in ffmpeg, optional

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/pyproject.toml

@@ -16,28 +16,26 @@ classifiers = [
 dynamic = ["version"]
 requires-python = ">=3.10"
 dependencies = [
-    "discord.py==1.7.3",
+    "discord.py==2.1.0",
     "Flask==2.1.2",
     "mongoengine==0.24.1",
     "pyarr==3.1.3",
-    "pymongo==4.1.1",
     "python-pushover==0.4",
-    "PyYAML==6.0",
-    "websockets==10.3"
+    "PyYAML==6.0"
 ]
 
 [project.optional-dependencies]
 dev = [
-    "black==22.3.0",
-    "discord.py-stubs==1.7.3",
-    "flake8==4.0.1",
-    "isort==5.10.1",
+    "black==22.12.0",
+    "flake8==6.0.0",
+    "isort==5.12.0",
     "mongo-types==0.15.1",
-    "mypy==0.950",
-    "pre-commit==2.19.0",
-    "requests==2.27.1",
-    "types-PyYAML==6.0.7",
-    "types-requests==2.27.29"
+    "mypy==0.991",
+    "pre-commit==3.0.2",
+    "requests==2.28.2",
+    "types-flask==1.1.6",
+    "types-PyYAML==6.0.12.3",
+    "types-requests==2.28.11.8"
 ]
 
 [project.urls]
@@ -46,7 +44,6 @@ Homepage = "https://github.com/wthueb/wi1-bot"
 [project.scripts]
 wi1-bot = "wi1_bot.scripts.start:main"
 transcode-item = "wi1_bot.scripts.transcode_item:main"
-add-tag = "wi1_bot.scripts.add_tag:main"
 
 [tool.setuptools]
 packages = ["wi1_bot"]
@@ -55,13 +52,13 @@ packages = ["wi1_bot"]
 write_to = "wi1_bot/_version.py"
 
 [tool.mypy]
-namespace_packages = true
+packages = "wi1_bot"
+strict = true
 
 [[tool.mypy.overrides]]
 module = [
     "pushover",
     "pyarr",
-    "setuptools_scm"
 ]
 ignore_missing_imports = true
 

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/tests/movie_downloaded.py

@@ -1,13 +1,13 @@
+import pathlib
 import shutil
 
 import requests
 
 from wi1_bot.config import config
 
-shutil.copy(
-    "./tests/files/jellyfish-10-mbps-hd-h264.mkv.bak",
-    "./tests/files/jellyfish-10-mbps-hd-h264.mkv",
-)
+path = pathlib.Path("./tests/files/jellyfish-10-mbps-hd-h264.mkv")
+
+shutil.copy(f"{path}.bak", path)
 
 header = {
     "X-Api-Key": config["radarr"]["api_key"],
@@ -19,7 +19,7 @@ data = {
         "id": 1,  # just to get quality profile
         "folderPath": "./tests/files",
     },
-    "movieFile": {"relativePath": "jellyfish-10-mbps-hd-h264.mkv"},
+    "movieFile": {"relativePath": path.name},
 }
 
 requests.post("http://localhost:9000/", json=data)

wi1-bot-1.3.1/wi1_bot/_version.py

@@ -0,0 +1,4 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+__version__ = version = '1.3.1'
+__version_tuple__ = version_tuple = (1, 3, 1)

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/wi1_bot/arr/download.py

@@ -1,9 +1,11 @@
+from typing import Any
+
 from .episode import Episode
 from .movie import Movie
 
 
 class Download:
-    def __init__(self, data: dict) -> None:
+    def __init__(self, data: dict[str, Any]) -> None:
         self.content: Movie | Episode | str
 
         if "movie" in data:

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/wi1_bot/arr/episode.py

@@ -1,7 +1,10 @@
+from typing import Any
+
+
 class Episode:
     def __init__(
         self,
-        ep_json: dict,
+        ep_json: dict[str, Any],
         *,
         series_title: str,
         series_tvdb_id: int,

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/wi1_bot/arr/movie.py

@@ -1,5 +1,8 @@
+from typing import Any
+
+
 class Movie:
-    def __init__(self, movie_json: dict) -> None:
+    def __init__(self, movie_json: dict[str, Any]) -> None:
         self._json = movie_json
 
         self.title: str = movie_json["title"]

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/wi1_bot/arr/radarr.py

@@ -1,34 +1,37 @@
 from shutil import rmtree
+from typing import Any, cast
 
 from pyarr import RadarrAPI
 
 from .download import Download
 from .movie import Movie
 
+json = dict[str, Any]
+
 
 class Radarr:
     def __init__(self, url: str, api_key: str) -> None:
         self._radarr = RadarrAPI(url, api_key)
 
     def lookup_movie(self, query: str) -> list[Movie]:
-        possible_movies = self._radarr.lookup_movie(query)
+        possible_movies = cast(list[json], self._radarr.lookup_movie(query))
 
         return [Movie(m) for m in possible_movies]
 
     def lookup_library(self, query: str) -> list[Movie]:
-        possible_movies = self._radarr.lookup_movie(query)
+        possible_movies = cast(list[json], self._radarr.lookup_movie(query))
 
         return [Movie(m) for m in possible_movies if "id" in m]
 
-    def lookup_user_movies(self, query: str, user_id: int) -> list[Movie]:
+    def lookup_user_library(self, query: str, user_id: int) -> list[Movie]:
         try:
             tag_id = self._get_tag_for_user_id(user_id)
         except ValueError:
             return []
 
-        tag_detail = self._radarr.get_tag_detail(tag_id)
+        tag_detail = cast(dict[str, Any], self._radarr.get_tag_detail(tag_id))
 
-        possible_movies = self._radarr.lookup_movie(query)
+        possible_movies = cast(list[json], self._radarr.lookup_movie(query))
 
         user_movie_ids = tag_detail["movieIds"]
 
@@ -42,7 +45,7 @@ class Radarr:
 
         quality_profile_id = self._get_quality_profile_id(profile)
 
-        root_folder = self._radarr.get_root_folder()[0]["path"]
+        root_folder = cast(list[json], self._radarr.get_root_folder())[0]["path"]
 
         self._radarr.add_movie(
             db_id=movie.tmdb_id,
@@ -53,7 +56,7 @@ class Radarr:
         return True
 
     def del_movie(self, movie: Movie) -> None:
-        potential = self._radarr.get_movie(movie.tmdb_id)
+        potential = cast(list[json], self._radarr.get_movie(movie.tmdb_id))
 
         if not potential:
             raise ValueError(f"{movie} is not in the library")
@@ -71,7 +74,7 @@ class Radarr:
             pass
 
     def movie_downloaded(self, movie: Movie) -> bool:
-        potential = self._radarr.get_movie(movie.tmdb_id)
+        potential = cast(list[json], self._radarr.get_movie(movie.tmdb_id))
 
         if not potential:
             return False
@@ -88,9 +91,12 @@ class Radarr:
 
     def add_tag(self, movie: Movie | list[Movie], user_id: int) -> bool:
         if isinstance(movie, Movie):
-            ids = [self._radarr.get_movie(movie.tmdb_id)[0]["id"]]
+            ids = [cast(list[json], self._radarr.get_movie(movie.tmdb_id))[0]["id"]]
         else:
-            ids = [self._radarr.get_movie(m.tmdb_id)[0]["id"] for m in movie]
+            ids = [
+                cast(list[json], self._radarr.get_movie(m.tmdb_id))[0]["id"]
+                for m in movie
+            ]
 
         try:
             tag_id = self._get_tag_for_user_id(user_id)
@@ -106,7 +112,7 @@ class Radarr:
         return True
 
     def get_downloads(self) -> list[Download]:
-        queue = self._radarr.get_queue_details()
+        queue = cast(list[json], self._radarr.get_queue_details())
 
         downloads = [Download(d) for d in queue]
 
@@ -120,46 +126,52 @@ class Radarr:
 
         total = 0
 
-        for movie in self._radarr.get_movie():
+        for movie in cast(list[json], self._radarr.get_movie()):
             if tag_id in movie["tags"]:
                 total += movie["sizeOnDisk"]
 
         return total
 
-    def get_quality_profile_name(self, profile_id: int):
-        profiles = self._radarr.get_quality_profile()
+    def get_quality_profile_name(self, profile_id: int) -> str:
+        profiles = cast(list[json], self._radarr.get_quality_profile())
 
         for profile in profiles:
             if profile["id"] == profile_id:
-                return profile["name"]
+                return cast(str, profile["name"])
 
         raise ValueError(f"no quality profile with the id {profile_id}")
 
     def rescan_movie(self, movie_id: int) -> None:
-        self._radarr.post_command("RescanMovie", movieIds=[movie_id])
+        self._radarr.post_command("RescanMovie", movieId=movie_id)
+
+    def refresh_movie(self, movie_id: int) -> None:
+        self._radarr.post_command("RefreshMovie", movieIds=[movie_id])
 
     def search_missing(self) -> None:
         self._radarr.post_command(name="MissingMoviesSearch")
 
     def _get_quality_profile_id(self, name: str) -> int:
-        profiles = self._radarr.get_quality_profile()
+        profiles = cast(list[json], self._radarr.get_quality_profile())
 
         for profile in profiles:
             if profile["name"].lower() == name.lower():
-                return profile["id"]
+                return cast(int, profile["id"])
 
         raise ValueError(f"no quality profile with the name {name}")
 
     def _get_tag_for_user_id(self, user_id: int) -> int:
-        tags = self._radarr.get_tag()
+        tags = cast(list[json], self._radarr.get_tag())
 
         for tag in tags:
             if str(user_id) in tag["label"]:
-                return tag["id"]
+                return cast(int, tag["id"])
 
         raise ValueError(f"no tag with the user id {user_id}")
 
 
+__all__ = ["Movie"]
+
+
 if __name__ == "__main__":
     from wi1_bot.config import config
 

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/wi1_bot/arr/sonarr.py

@@ -1,12 +1,15 @@
 from shutil import rmtree
+from typing import Any, cast
 
 from pyarr import SonarrAPI
 
 from .download import Download
 
+json = dict[str, Any]
+
 
 class Series:
-    def __init__(self, series_json: dict) -> None:
+    def __init__(self, series_json: dict[str, Any]) -> None:
         self.title: str = series_json["title"]
         self.year: int = series_json["year"]
         self.tvdb_id: int = series_json["tvdbId"]
@@ -32,33 +35,59 @@ class Series:
         return str(self.__dict__)
 
 
+class SonarrError(Exception):
+    pass
+
+
 class Sonarr:
     def __init__(self, url: str, api_key: str) -> None:
         self._sonarr = SonarrAPI(url, api_key)
 
     def lookup_series(self, query: str) -> list[Series]:
-        possible_series = self._sonarr.lookup_series(query)
+        possible_series = cast(
+            list[json] | dict[str, Any], self._sonarr.lookup_series(query)
+        )
+
+        if isinstance(possible_series, dict):
+            raise SonarrError(possible_series["message"])
 
         return [Series(s) for s in possible_series]
 
     def lookup_library(self, query: str) -> list[Series]:
-        possible_series = self._sonarr.lookup_series(query)
+        possible_series = cast(list[json], self._sonarr.lookup_series(query))
 
         return [Series(s) for s in possible_series if "id" in s]
 
+    def lookup_user_library(self, query: str, user_id: int) -> list[Series]:
+        try:
+            tag_id = self._get_tag_for_user_id(user_id)
+        except ValueError:
+            return []
+
+        possible_series = cast(list[json], self._sonarr.lookup_series(query))
+
+        # self._sonarr.get_tag_detail is broken/not supported in v1 sonarr API so have
+        # to filter the tmdb lookup (probably slower but saves an API call)
+        user_series = [s for s in possible_series if tag_id in s["tags"]]
+
+        return [Series(s) for s in user_series]
+
     def add_series(self, series: Series, profile: str = "good") -> bool:
         if series.db_id is not None:
             return False
 
         quality_profile_id = self._get_quality_profile_id(profile)
 
-        root_folder = self._sonarr.get_root_folder()[0]["path"]
+        root_folder = cast(list[json], self._sonarr.get_root_folder())[0]["path"]
 
-        series_json = self._sonarr.add_series(
-            tvdb_id=series.tvdb_id,
-            quality_profile_id=quality_profile_id,
-            root_dir=root_folder,
-            search_for_missing_episodes=True,
+        series_json = cast(
+            json,
+            self._sonarr.add_series(
+                tvdb_id=series.tvdb_id,
+                quality_profile_id=quality_profile_id,
+                root_dir=root_folder,
+                search_for_missing_episodes=True,
+            ),
         )
 
         series.db_id = series_json["id"]
@@ -69,7 +98,7 @@ class Sonarr:
         if series.db_id is None:
             raise ValueError(f"{series} is not in the library")
 
-        series_json = self._sonarr.get_series(series.db_id)
+        series_json = cast(json, self._sonarr.get_series(series.db_id))
 
         self._sonarr.del_series(series.db_id, delete_files=True)
 
@@ -100,7 +129,7 @@ class Sonarr:
 
             return False
 
-        series_json = self._sonarr.get_series(series.db_id)
+        series_json = cast(json, self._sonarr.get_series(series.db_id))
 
         series_json["tags"].append(tag_id)
 
@@ -109,7 +138,7 @@ class Sonarr:
         return True
 
     def get_downloads(self) -> list[Download]:
-        queue = self._sonarr.get_queue()
+        queue = cast(list[json], self._sonarr.get_queue())
 
         downloads = [Download(d) for d in queue]
 
@@ -123,18 +152,21 @@ class Sonarr:
 
         total = 0
 
-        for series in self._sonarr.get_series():
+        for series in cast(list[json], self._sonarr.get_series()):
             if tag_id in series["tags"]:
-                total += series["sizeOnDisk"]
+                try:
+                    total += series["sizeOnDisk"]
+                except KeyError:
+                    continue
 
         return total
 
-    def get_quality_profile_name(self, profile_id: int):
-        profiles = self._sonarr.get_quality_profile()
+    def get_quality_profile_name(self, profile_id: int) -> str:
+        profiles = cast(list[json], self._sonarr.get_quality_profile())
 
         for profile in profiles:
             if profile["id"] == profile_id:
-                return profile["name"]
+                return cast(str, profile["name"])
 
         raise ValueError(f"no quality profile with the id {profile_id}")
 
@@ -142,19 +174,19 @@ class Sonarr:
         self._sonarr.post_command("RescanSeries", seriesId=series_id)
 
     def _get_quality_profile_id(self, name: str) -> int:
-        profiles = self._sonarr.get_quality_profile()
+        profiles = cast(list[json], self._sonarr.get_quality_profile())
 
         for profile in profiles:
             if profile["name"].lower() == name.lower():
-                return profile["id"]
+                return cast(int, profile["id"])
 
         raise ValueError(f"no quality profile with the name {name}")
 
     def _get_tag_for_user_id(self, user_id: int) -> int:
-        tags = self._sonarr.get_tag()
+        tags = cast(list[json], self._sonarr.get_tag())
 
         for tag in tags:
             if str(user_id) in tag["label"]:
-                return tag["id"]
+                return cast(int, tag["id"])
 
         raise ValueError(f"no tag with the user id {user_id}")

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/wi1_bot/config.py

@@ -26,6 +26,7 @@ if _config_path is None:
 class ArrConfig(TypedDict):
     url: str
     api_key: str
+    root_folder: str
 
 
 class PushoverConfig(TypedDict):
@@ -79,6 +80,7 @@ class DiscordConfig(DiscordConfigOptional):
 
 
 class TranscodingProfile(TypedDict, total=False):
+    copy_all_streams: bool
     video_codec: str
     video_bitrate: int
     audio_codec: str

{wi1-bot-1.2.12 → wi1-bot-1.3.1}/wi1_bot/discord/bot.py

@@ -1,4 +1,7 @@
+import asyncio
 import logging
+import traceback
+from typing import Any
 
 import discord
 from discord.ext import commands
@@ -18,23 +21,32 @@ sonarr = Sonarr(config["sonarr"]["url"], config["sonarr"]["api_key"])
 
 
 @bot.check
-async def check_channel(ctx: commands.Context) -> bool:
+async def check_channel(ctx: commands.Context[Any]) -> bool:
     return ctx.channel.id == config["discord"]["channel_id"]
 
 
 @bot.event
-async def on_command_error(ctx: commands.Context, error: commands.CommandError) -> None:
+async def on_command_error(
+    ctx: commands.Context[Any], error: commands.CommandError
+) -> None:
     match error:
-        case commands.CommandNotFound():
+        case commands.CommandNotFound() | commands.CheckFailure():
             pass
         case commands.MissingRole():
             await reply(ctx.message, "you don't have permission to do that")
         case commands.MemberNotFound():
             await reply(ctx.message, "that user doesn't exist")
+        case commands.CommandOnCooldown():
+            await reply(ctx.message, str(error))
         case commands.MissingRequiredArgument():
             await reply(ctx.message, str(error))
         case _:
-            logger.error(repr(error))
+            logger.error(
+                "".join(
+                    traceback.format_exception(type(error), error, error.__traceback__)
+                )
+            )
+
             await reply(
                 ctx.message,
                 f"something went wrong (<@!{config['discord']['admin_id']}>)",
@@ -55,15 +67,15 @@ async def on_ready() -> None:
 
 
 @bot.before_invoke
-async def before_invoke(ctx: commands.Context) -> None:
+async def before_invoke(ctx: commands.Context[Any]) -> None:
     logger.debug(f"got command from {ctx.message.author}: {ctx.message.content}")
 
 
-@commands.cooldown(1, 10)  # one time every 10 seconds
+@commands.cooldown(1, 10)  # type: ignore
 @bot.command(
     name="downloads", aliases=["queue", "q"], help="see the status of movie downloads"
 )
-async def downloads_cmd(ctx: commands.Context) -> None:
+async def downloads_cmd(ctx: commands.Context[Any]) -> None:
     async with ctx.typing():
         queue = radarr.get_downloads() + sonarr.get_downloads()
 
@@ -76,9 +88,9 @@ async def downloads_cmd(ctx: commands.Context) -> None:
     await reply(ctx.message, "\n\n".join(map(str, queue)), title="download progress")
 
 
-@commands.cooldown(1, 60, commands.BucketType.user)
+@commands.cooldown(1, 60, commands.BucketType.user)  # type: ignore
 @bot.command(name="quota", help="see your used space on the plex")
-async def quota_cmd(ctx: commands.Context) -> None:
+async def quota_cmd(ctx: commands.Context[Any]) -> None:
     async with ctx.typing():
         used = (
             radarr.get_quota_amount(ctx.message.author.id)
@@ -102,9 +114,9 @@ async def quota_cmd(ctx: commands.Context) -> None:
     await reply(ctx.message, msg)
 
 
-@commands.cooldown(1, 60)
+@commands.cooldown(1, 60)  # type: ignore
 @bot.command(name="quotas", help="see everyone's used space on the plex")
-async def quotas_cmd(ctx: commands.Context) -> None:
+async def quotas_cmd(ctx: commands.Context[Any]) -> None:
     try:
         quotas = config["discord"]["quotas"]
     except ValueError:
@@ -135,17 +147,30 @@ async def quotas_cmd(ctx: commands.Context) -> None:
     )
 
 
-bot.add_cog(MovieCog(bot))
-bot.add_cog(SeriesCog(bot))
+@bot.command(name="addtag", help="add a user tag")  # type: ignore
+@commands.has_role("plex-admin")
+async def addtag_cmd(
+    ctx: commands.Context[Any], name: str, user: discord.Member
+) -> None:
+    tag = f"{name}: {user.id}"
 
+    radarr.create_tag(tag)
+    sonarr.create_tag(tag)
 
-def run() -> None:
+    await reply(ctx.message, f"tag `{tag}` added for {user.display_name}")
+
+
+async def run() -> None:
     logger.debug("starting bot")
 
-    bot.run(config["discord"]["bot_token"])
+    async with bot:
+        await bot.add_cog(MovieCog(bot))
+        await bot.add_cog(SeriesCog(bot))
+
+        await bot.start(config["discord"]["bot_token"])
 
 
 if __name__ == "__main__":
     logger.addHandler(logging.StreamHandler())
 
-    bot.run(config["discord"]["bot_token"])
+    asyncio.run(run())