TheLovinator/ttvdrops
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Improve performance and add type hints

Browse source
This commit is contained in:
Joakim Hellsén 2026-04-11 00:44:16 +02:00
commit b7e10e766e
Signed by: Joakim Hellsén
SSH key fingerprint: SHA256:/9h/CsExpFp+PRhsfA0xznFx2CGfTT5R/kpuFfUgEQk
23 changed files with 745 additions and 178 deletions
Download patch file Download diff file Expand all files Collapse all files

14
twitch/management/commands/backup_db.py
View file

@ -7,6 +7,7 @@ from compression import zstd
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING
from typing import Protocol
from django.conf import settings
from django.core.management.base import BaseCommand
@ -19,6 +20,15 @@ if TYPE_CHECKING:
from argparse import ArgumentParser
class SupportsStr(Protocol):
"""Protocol for values that provide a string representation."""
def __str__(self) -> str: ...
type SqlSerializable = bool | int | float | bytes | SupportsStr | None
class Command(BaseCommand):
"""Create a compressed SQL dump of the Twitch and Kick dataset tables."""
@ -285,7 +295,7 @@ def _write_postgres_dump(output_path: Path, tables: list[str]) -> None:
raise CommandError(msg)
def _sql_literal(value: object) -> str:
def _sql_literal(value: SqlSerializable) -> str:
"""Convert a Python value to a SQL literal.
Args:
@ -305,7 +315,7 @@ def _sql_literal(value: object) -> str:
return "'" + str(value).replace("'", "''") + "'"
def _json_default(value: object) -> str:
def _json_default(value: bytes | SupportsStr) -> str:
"""Convert non-serializable values to JSON-compatible strings.
Args:

116
twitch/management/commands/better_import_drops.py
View file

@ -583,16 +583,31 @@ class Command(BaseCommand):
Returns:
Organization instance.
"""
org_obj, created = Organization.objects.get_or_create(
cache: dict[str, Organization] = getattr(self, "_org_cache", {})
if not hasattr(self, "_org_cache"):
self._org_cache = cache
cached_org: Organization | None = cache.get(org_data.twitch_id)
if cached_org is not None:
self._save_if_changed(cached_org, {"name": org_data.name})
return cached_org
org_obj: Organization | None = Organization.objects.filter(
twitch_id=org_data.twitch_id,
defaults={"name": org_data.name},
)
if not created:
self._save_if_changed(org_obj, {"name": org_data.name})
else:
).first()
_created: bool = org_obj is None
if org_obj is None:
org_obj = Organization.objects.create(
twitch_id=org_data.twitch_id,
name=org_data.name,
)
tqdm.write(
f"{Fore.GREEN}✓{Style.RESET_ALL} Created new organization: {org_data.name}",
)
else:
self._save_if_changed(org_obj, {"name": org_data.name})
cache[org_data.twitch_id] = org_obj
return org_obj
@ -621,6 +636,10 @@ class Command(BaseCommand):
if campaign_org_obj:
owner_orgs.add(campaign_org_obj)
cache: dict[str, Game] = getattr(self, "_game_cache", {})
if not hasattr(self, "_game_cache"):
self._game_cache = cache
defaults: dict[str, object] = {
"display_name": game_data.display_name or (game_data.name or ""),
"name": game_data.name or "",
@ -628,10 +647,22 @@ class Command(BaseCommand):
"box_art": game_data.box_art_url or "",
}
game_obj, created = Game.objects.get_or_create(
cached_game: Game | None = cache.get(game_data.twitch_id)
if cached_game is not None:
if owner_orgs:
cached_game.owners.add(*owner_orgs)
self._save_if_changed(cached_game, defaults)
return cached_game
game_obj: Game | None = Game.objects.filter(
twitch_id=game_data.twitch_id,
defaults=defaults,
)
).first()
created: bool = game_obj is None
if game_obj is None:
game_obj = Game.objects.create(
twitch_id=game_data.twitch_id,
**defaults,
)
# Set owners (ManyToMany)
if created or owner_orgs:
game_obj.owners.add(*owner_orgs)
@ -642,6 +673,7 @@ class Command(BaseCommand):
f"{Fore.GREEN}✓{Style.RESET_ALL} Created new game: {game_data.display_name}",
)
self._download_game_box_art(game_obj, game_obj.box_art)
cache[game_data.twitch_id] = game_obj
return game_obj
def _download_game_box_art(self, game_obj: Game, box_art_url: str | None) -> None:
@ -701,7 +733,7 @@ class Command(BaseCommand):
return channel_obj
def process_responses(
def process_responses( # noqa: PLR0915
self,
responses: list[dict[str, Any]],
file_path: Path,
@ -792,13 +824,18 @@ class Command(BaseCommand):
"account_link_url": drop_campaign.account_link_url,
}
campaign_obj, created = DropCampaign.objects.get_or_create(
campaign_obj: DropCampaign | None = DropCampaign.objects.filter(
twitch_id=drop_campaign.twitch_id,
defaults=defaults,
)
if not created:
self._save_if_changed(campaign_obj, defaults)
).first()
created: bool = campaign_obj is None
if campaign_obj is None:
campaign_obj = DropCampaign.objects.create(
twitch_id=drop_campaign.twitch_id,
**defaults,
)
else:
self._save_if_changed(campaign_obj, defaults)
if created:
tqdm.write(
f"{Fore.GREEN}✓{Style.RESET_ALL} Created new campaign: {drop_campaign.name}",
)
@ -882,13 +919,18 @@ class Command(BaseCommand):
if end_at_dt is not None:
drop_defaults["end_at"] = end_at_dt
drop_obj, created = TimeBasedDrop.objects.get_or_create(
drop_obj: TimeBasedDrop | None = TimeBasedDrop.objects.filter(
twitch_id=drop_schema.twitch_id,
defaults=drop_defaults,
)
if not created:
self._save_if_changed(drop_obj, drop_defaults)
).first()
created: bool = drop_obj is None
if drop_obj is None:
drop_obj = TimeBasedDrop.objects.create(
twitch_id=drop_schema.twitch_id,
**drop_defaults,
)
else:
self._save_if_changed(drop_obj, drop_defaults)
if created:
tqdm.write(
f"{Fore.GREEN}✓{Style.RESET_ALL} Created TimeBasedDrop: {drop_schema.name}",
)
@ -900,6 +942,10 @@ class Command(BaseCommand):
def _get_or_update_benefit(self, benefit_schema: DropBenefitSchema) -> DropBenefit:
"""Return a DropBenefit, creating or updating as needed."""
cache: dict[str, DropBenefit] = getattr(self, "_benefit_cache", {})
if not hasattr(self, "_benefit_cache"):
self._benefit_cache = cache
distribution_type: str = (benefit_schema.distribution_type or "").strip()
benefit_defaults: dict[str, str | int | datetime | bool | None] = {
"name": benefit_schema.name,
@ -914,10 +960,20 @@ class Command(BaseCommand):
if created_at_dt:
benefit_defaults["created_at"] = created_at_dt
benefit_obj, created = DropBenefit.objects.get_or_create(
cached_benefit: DropBenefit | None = cache.get(benefit_schema.twitch_id)
if cached_benefit is not None:
self._save_if_changed(cached_benefit, benefit_defaults)
return cached_benefit
benefit_obj: DropBenefit | None = DropBenefit.objects.filter(
twitch_id=benefit_schema.twitch_id,
defaults=benefit_defaults,
)
).first()
created: bool = benefit_obj is None
if benefit_obj is None:
benefit_obj = DropBenefit.objects.create(
twitch_id=benefit_schema.twitch_id,
**benefit_defaults,
)
if not created:
self._save_if_changed(benefit_obj, benefit_defaults)
else:
@ -925,6 +981,8 @@ class Command(BaseCommand):
f"{Fore.GREEN}✓{Style.RESET_ALL} Created DropBenefit: {benefit_schema.name}",
)
cache[benefit_schema.twitch_id] = benefit_obj
return benefit_obj
def _process_benefit_edges(
@ -946,11 +1004,17 @@ class Command(BaseCommand):
)
defaults = {"entitlement_limit": edge_schema.entitlement_limit}
edge_obj, created = DropBenefitEdge.objects.get_or_create(
edge_obj: DropBenefitEdge | None = DropBenefitEdge.objects.filter(
drop=drop_obj,
benefit=benefit_obj,
defaults=defaults,
)
).first()
created: bool = edge_obj is None
if edge_obj is None:
edge_obj = DropBenefitEdge.objects.create(
drop=drop_obj,
benefit=benefit_obj,
**defaults,
)
if not created:
self._save_if_changed(edge_obj, defaults)
else:

8
twitch/management/commands/download_box_art.py
View file

@ -39,9 +39,13 @@ class Command(BaseCommand):
help="Re-download even if a local box art file already exists.",
)
def handle(self, *_args: object, **options: object) -> None: # noqa: PLR0914, PLR0915
def handle( # noqa: PLR0914, PLR0915
self,
*_args: str,
**options: str | bool | int | None,
) -> None:
"""Download Twitch box art images for all games."""
limit_value: object | None = options.get("limit")
limit_value: str | bool | int | None = options.get("limit")
limit: int | None = limit_value if isinstance(limit_value, int) else None
force: bool = bool(options.get("force"))

8
twitch/management/commands/download_campaign_images.py
View file

@ -50,10 +50,14 @@ class Command(BaseCommand):
help="Re-download even if a local image file already exists.",
)
def handle(self, *_args: object, **options: object) -> None:
def handle(
self,
*_args: str,
**options: str | bool | int | None,
) -> None:
"""Download images for campaigns, benefits, and/or rewards."""
model_choice: str = str(options.get("model", "all"))
limit_value: object | None = options.get("limit")
limit_value: str | bool | int | None = options.get("limit")
limit: int | None = limit_value if isinstance(limit_value, int) else None
force: bool = bool(options.get("force"))

27
twitch/management/commands/import_chat_badges.py
View file

@ -196,9 +196,12 @@ class Command(BaseCommand):
Returns:
Tuple of (ChatBadgeSet instance, created flag)
"""
badge_set_obj, created = ChatBadgeSet.objects.get_or_create(
badge_set_obj: ChatBadgeSet | None = ChatBadgeSet.objects.filter(
set_id=badge_set_schema.set_id,
)
).first()
created: bool = badge_set_obj is None
if badge_set_obj is None:
badge_set_obj = ChatBadgeSet.objects.create(set_id=badge_set_schema.set_id)
if created:
self.stdout.write(
@ -258,11 +261,25 @@ class Command(BaseCommand):
"click_url": version_schema.click_url,
}
_badge_obj, created = ChatBadge.objects.update_or_create(
badge_obj: ChatBadge | None = ChatBadge.objects.filter(
badge_set=badge_set_obj,
badge_id=version_schema.badge_id,
defaults=defaults,
)
).first()
created: bool = badge_obj is None
if badge_obj is None:
badge_obj = ChatBadge.objects.create(
badge_set=badge_set_obj,
badge_id=version_schema.badge_id,
**defaults,
)
else:
changed_fields: list[str] = []
for field, value in defaults.items():
if getattr(badge_obj, field) != value:
setattr(badge_obj, field, value)
changed_fields.append(field)
if changed_fields:
badge_obj.save(update_fields=changed_fields)
if created:
msg: str = (