273 lines
9.4 KiB
Python
273 lines
9.4 KiB
Python
"""Management command to import Twitch global chat badges."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from typing import TYPE_CHECKING
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from colorama import Fore
|
|
from colorama import Style
|
|
from colorama import init as colorama_init
|
|
from django.core.management.base import BaseCommand
|
|
from django.core.management.base import CommandError
|
|
from django.core.management.base import CommandParser
|
|
from pydantic import ValidationError
|
|
|
|
from twitch.models import ChatBadge
|
|
from twitch.models import ChatBadgeSet
|
|
from twitch.schemas import ChatBadgeSetSchema
|
|
from twitch.schemas import GlobalChatBadgesResponse
|
|
|
|
if TYPE_CHECKING:
|
|
from twitch.schemas import ChatBadgeVersionSchema
|
|
|
|
logger: logging.Logger = logging.getLogger("ttvdrops")
|
|
|
|
|
|
class Command(BaseCommand):
|
|
"""Import Twitch global chat badges from the Twitch Helix API."""
|
|
|
|
help = "Import Twitch global chat badges from the Twitch Helix API"
|
|
requires_migrations_checks = True
|
|
|
|
def add_arguments(self, parser: CommandParser) -> None:
|
|
"""Add command arguments."""
|
|
parser.add_argument(
|
|
"--client-id",
|
|
type=str,
|
|
help="Twitch Client ID (or set TWITCH_CLIENT_ID environment variable)",
|
|
)
|
|
parser.add_argument(
|
|
"--client-secret",
|
|
type=str,
|
|
help="Twitch Client Secret (or set TWITCH_CLIENT_SECRET environment variable)",
|
|
)
|
|
parser.add_argument(
|
|
"--access-token",
|
|
type=str,
|
|
help="Twitch Access Token (optional - will be obtained automatically if not provided)",
|
|
)
|
|
|
|
def handle(self, *args, **options) -> None: # noqa: ARG002
|
|
"""Main entry point for the command.
|
|
|
|
Raises:
|
|
CommandError: If required arguments are missing or API calls fail
|
|
"""
|
|
colorama_init(autoreset=True)
|
|
|
|
# Get credentials from arguments or environment
|
|
client_id: str | None = options.get("client_id") or os.getenv("TWITCH_CLIENT_ID")
|
|
client_secret: str | None = options.get("client_secret") or os.getenv("TWITCH_CLIENT_SECRET")
|
|
access_token: str | None = options.get("access_token") or os.getenv("TWITCH_ACCESS_TOKEN")
|
|
|
|
if not client_id:
|
|
msg = (
|
|
"Twitch Client ID is required. "
|
|
"Provide it via --client-id argument or set TWITCH_CLIENT_ID environment variable."
|
|
)
|
|
raise CommandError(msg)
|
|
|
|
# If access token is not provided, obtain it automatically using client credentials
|
|
if not access_token:
|
|
if not client_secret:
|
|
msg = (
|
|
"Either --access-token or --client-secret must be provided. "
|
|
"Set TWITCH_ACCESS_TOKEN or TWITCH_CLIENT_SECRET environment variable, "
|
|
"or provide them via command arguments."
|
|
)
|
|
raise CommandError(msg)
|
|
|
|
self.stdout.write("Obtaining access token from Twitch...")
|
|
try:
|
|
access_token = self._get_app_access_token(client_id, client_secret)
|
|
self.stdout.write(self.style.SUCCESS("✓ Access token obtained successfully"))
|
|
except httpx.HTTPError as e:
|
|
msg = f"Failed to obtain access token: {e}"
|
|
raise CommandError(msg) from e
|
|
|
|
self.stdout.write("Fetching global chat badges from Twitch API...")
|
|
|
|
try:
|
|
badges_data: GlobalChatBadgesResponse = self._fetch_global_chat_badges(
|
|
client_id=client_id,
|
|
access_token=access_token,
|
|
)
|
|
except httpx.HTTPError as e:
|
|
msg: str = f"Failed to fetch chat badges from Twitch API: {e}"
|
|
raise CommandError(msg) from e
|
|
except ValidationError as e:
|
|
msg: str = f"Failed to validate chat badges response: {e}"
|
|
raise CommandError(msg) from e
|
|
|
|
self.stdout.write(f"Received {len(badges_data.data)} badge sets")
|
|
|
|
# Process and store badge data
|
|
created_sets = 0
|
|
updated_sets = 0
|
|
created_badges = 0
|
|
updated_badges = 0
|
|
|
|
for badge_set_schema in badges_data.data:
|
|
badge_set_obj, set_created = self._process_badge_set(badge_set_schema)
|
|
|
|
if set_created:
|
|
created_sets += 1
|
|
else:
|
|
updated_sets += 1
|
|
|
|
# Process each badge version in the set
|
|
for version_schema in badge_set_schema.versions:
|
|
badge_created: bool = self._process_badge_version(
|
|
badge_set_obj=badge_set_obj,
|
|
version_schema=version_schema,
|
|
)
|
|
|
|
if badge_created:
|
|
created_badges += 1
|
|
else:
|
|
updated_badges += 1
|
|
|
|
# Print summary
|
|
self.stdout.write("\n" + "=" * 50)
|
|
self.stdout.write(
|
|
self.style.SUCCESS(
|
|
f"✓ Created {created_sets} new badge sets, updated {updated_sets} existing",
|
|
),
|
|
)
|
|
self.stdout.write(
|
|
self.style.SUCCESS(
|
|
f"✓ Created {created_badges} new badges, updated {updated_badges} existing",
|
|
),
|
|
)
|
|
self.stdout.write(f"Total badge sets: {created_sets + updated_sets}")
|
|
self.stdout.write(f"Total badges: {created_badges + updated_badges}")
|
|
self.stdout.write("=" * 50)
|
|
|
|
def _fetch_global_chat_badges(
|
|
self,
|
|
client_id: str,
|
|
access_token: str,
|
|
) -> GlobalChatBadgesResponse:
|
|
"""Fetch global chat badges from Twitch Helix API.
|
|
|
|
Args:
|
|
client_id: Twitch Client ID
|
|
access_token: Twitch Access Token (app or user token)
|
|
|
|
Returns:
|
|
Validated GlobalChatBadgesResponse
|
|
"""
|
|
url = "https://api.twitch.tv/helix/chat/badges/global"
|
|
headers: dict[str, str] = {
|
|
"Authorization": f"Bearer {access_token}",
|
|
"Client-Id": client_id,
|
|
}
|
|
|
|
with httpx.Client() as client:
|
|
response: httpx.Response = client.get(url, headers=headers, timeout=30.0)
|
|
response.raise_for_status()
|
|
|
|
data: dict[str, Any] = response.json()
|
|
logger.debug("Received chat badges response: %s", data)
|
|
|
|
# Validate response with Pydantic
|
|
return GlobalChatBadgesResponse.model_validate(data)
|
|
|
|
def _process_badge_set(
|
|
self,
|
|
badge_set_schema: ChatBadgeSetSchema,
|
|
) -> tuple[ChatBadgeSet, bool]:
|
|
"""Get or create a ChatBadgeSet from the schema.
|
|
|
|
Args:
|
|
badge_set_schema: Validated badge set schema
|
|
|
|
Returns:
|
|
Tuple of (ChatBadgeSet instance, created flag)
|
|
"""
|
|
badge_set_obj, created = ChatBadgeSet.objects.get_or_create(
|
|
set_id=badge_set_schema.set_id,
|
|
)
|
|
|
|
if created:
|
|
self.stdout.write(
|
|
f"{Fore.GREEN}✓{Style.RESET_ALL} Created new badge set: {badge_set_schema.set_id}",
|
|
)
|
|
else:
|
|
self.stdout.write(
|
|
f"{Fore.YELLOW}→{Style.RESET_ALL} Badge set already exists: {badge_set_schema.set_id}",
|
|
)
|
|
|
|
return badge_set_obj, created
|
|
|
|
def _get_app_access_token(self, client_id: str, client_secret: str) -> str:
|
|
"""Get an app access token from Twitch.
|
|
|
|
Args:
|
|
client_id: Twitch Client ID
|
|
client_secret: Twitch Client Secret
|
|
|
|
Returns:
|
|
Access token string
|
|
"""
|
|
url = "https://id.twitch.tv/oauth2/token"
|
|
params: dict[str, str] = {
|
|
"client_id": client_id,
|
|
"client_secret": client_secret,
|
|
"grant_type": "client_credentials",
|
|
}
|
|
|
|
with httpx.Client() as client:
|
|
response: httpx.Response = client.post(url, params=params, timeout=30.0)
|
|
response.raise_for_status()
|
|
data: dict[str, str] = response.json()
|
|
return data["access_token"]
|
|
|
|
def _process_badge_version(
|
|
self,
|
|
badge_set_obj: ChatBadgeSet,
|
|
version_schema: ChatBadgeVersionSchema,
|
|
) -> bool:
|
|
"""Get or create a ChatBadge from the schema.
|
|
|
|
Args:
|
|
badge_set_obj: Parent ChatBadgeSet instance
|
|
version_schema: Validated badge version schema
|
|
|
|
Returns:
|
|
True if created, False if updated
|
|
"""
|
|
defaults: dict[str, str | None] = {
|
|
"image_url_1x": version_schema.image_url_1x,
|
|
"image_url_2x": version_schema.image_url_2x,
|
|
"image_url_4x": version_schema.image_url_4x,
|
|
"title": version_schema.title,
|
|
"description": version_schema.description,
|
|
"click_action": version_schema.click_action,
|
|
"click_url": version_schema.click_url,
|
|
}
|
|
|
|
_badge_obj, created = ChatBadge.objects.update_or_create(
|
|
badge_set=badge_set_obj,
|
|
badge_id=version_schema.badge_id,
|
|
defaults=defaults,
|
|
)
|
|
|
|
if created:
|
|
msg: str = (
|
|
f"{Fore.GREEN}✓{Style.RESET_ALL} Created badge: "
|
|
f"{badge_set_obj.set_id}/{version_schema.badge_id} - {version_schema.title}"
|
|
)
|
|
self.stdout.write(msg)
|
|
else:
|
|
msg: str = (
|
|
f"{Fore.YELLOW}→{Style.RESET_ALL} Updated badge: "
|
|
f"{badge_set_obj.set_id}/{version_schema.badge_id} - {version_schema.title}"
|
|
)
|
|
self.stdout.write(msg)
|
|
|
|
return created
|