Resolve mypy type errors.

yellowstone/core.py

@@ -8,13 +8,20 @@
 import json
 import logging
 import time
-from typing import NoReturn, TypedDict
+from typing import NoReturn, TypedDict, cast
 
 import pugsql
 
 from .config import Config, getenv
 from .exceptions import UnknownJobError
-from .jobs import JobType, get_site, get_user, get_user_avatar, index_site_members
+from .jobs import (
+    JobType,
+    add_index_site_members_job,
+    get_site,
+    get_user,
+    get_user_avatar,
+    index_site_members,
+)
 from .jobs.index_site_members import START_OFFSET as START_MEMBER_OFFSET
 from .s3 import S3
 from .types import Json
@@ -73,8 +80,8 @@ def insert_all_sites(self) -> None:
     def queue_all_sites(self) -> None:
         for site_slug in self.config.site_slugs:
             logger.info("Queueing site start jobs for '%s'", site_slug)
-            # add_index_site_pages_job(site_slug)
-            # add_index_site_forums_job(site_slug)
+            # XXX add_index_site_pages_job(site_slug)
+            # XXX add_index_site_forums_job(site_slug)
             add_index_site_members_job(
                 self.database,
                 {
@@ -109,14 +116,17 @@ def process_job(self, job: JobDict) -> None:
                 case JobType.INDEX_SITE_FORUMS:
                     raise NotImplementedError
                 case JobType.INDEX_SITE_MEMBERS:
-                    assert isinstance(data, index_site_members.SiteMemberJob), "INDEX_SITE_MEMBERS"
-                    index_site_members.run(self, data)
+                    index_site_members.run(
+                        self,
+                        cast(index_site_members.SiteMemberJob, data),
+                    )
                 case JobType.FETCH_USER:
-                    assert isinstance(data, get_user.GetUserJob), "GET_USER"
-                    get_user.run(self, data)
+                    get_user.run(self, cast(get_user.GetUserJob, data))
                 case JobType.FETCH_USER_AVATAR:
-                    assert isinstance(data, get_user_avatar.GetUserAvatarJob), "GET_USER_AVATAR"
-                    get_user_avatar.run(self, data)
+                    get_user_avatar.run(
+                        self,
+                        cast(get_user_avatar.GetUserAvatarJob, data),
+                    )
                 case _:
                     raise UnknownJobError(f"Unknown job type: {job_type}")
         except UnknownJobError:

yellowstone/jobs/__init__.py

@@ -9,14 +9,13 @@
 
 import json
 from enum import Enum, unique
-from typing import TYPE_CHECKING
+from typing import cast
 
 from ..types import Json
+from .get_user import GetUserJob
+from .get_user_avatar import GetUserAvatarJob
 from .index_site_members import SiteMemberJob
 
-if TYPE_CHECKING:
-    pass
-
 
 @unique
 class JobType(Enum):
@@ -29,6 +28,7 @@ class JobType(Enum):
 
 # TODO change to JobManager class?
 
+
 def add_raw_job(database, type: JobType, data: Json) -> None:
     database.add_job(
         job_type=type.value,
@@ -45,4 +45,12 @@ def add_index_site_forums_job(database, data: None) -> None:
 
 
 def add_index_site_members_job(database, data: SiteMemberJob) -> None:
-    add_raw_job(database, JobType.INDEX_SITE_MEMBERS, data)
+    add_raw_job(database, JobType.INDEX_SITE_MEMBERS, cast(Json, data))
+
+
+def add_fetch_user_job(database, data: GetUserJob) -> None:
+    add_raw_job(database, JobType.FETCH_USER, cast(Json, data))
+
+
+def add_fetch_user_avatar_job(database, data: GetUserAvatarJob) -> None:
+    add_raw_job(database, JobType.FETCH_USER_AVATAR, cast(Json, data))

yellowstone/jobs/get_user.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING, TypedDict
 
 from ..requests import user as user_data
-from . import JobType
+from . import add_fetch_user_avatar_job
 
 if TYPE_CHECKING:
     from ..core import BackupDispatcher
@@ -39,4 +39,4 @@ def run(core: "BackupDispatcher", data: GetUserJob) -> None:
     )
 
     # Enqueue avatar job
-    core.add_job(JobType.FETCH_USER_AVATAR, data)
+    add_fetch_user_avatar_job(core.database, data)

yellowstone/jobs/index_site_members.py

@@ -9,7 +9,7 @@
 from typing import TYPE_CHECKING, TypedDict
 
 from ..requests import site_members
-from . import JobType
+from . import add_fetch_user_job, add_index_site_members_job
 
 if TYPE_CHECKING:
     from ..core import BackupDispatcher
@@ -48,7 +48,10 @@ def run(core: "BackupDispatcher", data: SiteMemberJob) -> None:
     if members:
         with core.database.transaction():
             # Queue the next offset, for iterating over pages using the job queue
-            core.add_job(JobType.INDEX_SITE_MEMBERS, site_slug, offset + 1)
+            add_index_site_members_job(
+                core.database,
+                {"site_slug": site_slug, "offset": offset + 1},
+            )
 
             # Add all site members, and queue their users for update
             for member in members:
@@ -57,4 +60,4 @@ def run(core: "BackupDispatcher", data: SiteMemberJob) -> None:
                     site_id=site_id,
                     joined_at=member.joined_at,
                 )
-                core.add_job(JobType.FETCH_USER, member.slug, member.id)
+                add_fetch_user_job(core.database, {"user_id": member.id})