yt-dlp/yt_dlp/extractor/youtube/jsc/provider.py

"""PUBLIC API"""

from __future__ import annotations

import abc
import dataclasses
import enum
import typing

from yt_dlp.extractor.youtube.jsc._registry import _jsc_preferences, _jsc_providers
from yt_dlp.extractor.youtube.pot._provider import (
    IEContentProvider,
    IEContentProviderError,
    register_preference_generic,
    register_provider_generic,
)
from yt_dlp.utils import ExtractorError

__all__ = [
    'JsChallengeProvider',
    'JsChallengeProviderError',
    'JsChallengeProviderRejectedRequest',
    'JsChallengeProviderResponse',
    'JsChallengeRequest',
    'JsChallengeResponse',
    'JsChallengeType',
    'NChallengeInput',
    'NChallengeOutput',
    'SigChallengeInput',
    'SigChallengeOutput',
    'register_preference',
    'register_provider',
]


class JsChallengeType(enum.Enum):
    N = 'n'
    SIG = 'sig'


@dataclasses.dataclass(frozen=True)
class JsChallengeRequest:
    type: JsChallengeType
    input: NChallengeInput | SigChallengeInput
    video_id: str | None = None


@dataclasses.dataclass(frozen=True)
class NChallengeInput:
    player_url: str
    challenges: list[str] = dataclasses.field(default_factory=list)


@dataclasses.dataclass(frozen=True)
class SigChallengeInput:
    player_url: str
    challenges: list[str] = dataclasses.field(default_factory=list)


@dataclasses.dataclass(frozen=True)
class NChallengeOutput:
    results: dict[str, str] = dataclasses.field(default_factory=dict)


@dataclasses.dataclass(frozen=True)
class SigChallengeOutput:
    results: dict[str, str] = dataclasses.field(default_factory=dict)


@dataclasses.dataclass
class JsChallengeProviderResponse:
    request: JsChallengeRequest
    response: JsChallengeResponse | None = None
    error: Exception | None = None


@dataclasses.dataclass
class JsChallengeResponse:
    type: JsChallengeType
    output: NChallengeOutput | SigChallengeOutput


class JsChallengeProviderRejectedRequest(IEContentProviderError):
    """Reject the JsChallengeRequest (cannot handle the request)"""

    def __init__(self, msg=None, expected: bool = False, *, _skipped_components=None):
        super().__init__(msg, expected)
        self._skipped_components = _skipped_components


class JsChallengeProviderError(IEContentProviderError):
    """An error occurred while solving the challenge"""


class JsChallengeProvider(IEContentProvider, abc.ABC, suffix='JCP'):

    # Set to None to disable the check
    _SUPPORTED_TYPES: tuple[JsChallengeType] | None = ()

    def __validate_request(self, request: JsChallengeRequest):
        if not self.is_available():
            raise JsChallengeProviderRejectedRequest(f'{self.PROVIDER_NAME} is not available')

        # Validate request using built-in settings
        if (
            self._SUPPORTED_TYPES is not None
            and request.type not in self._SUPPORTED_TYPES
        ):
            raise JsChallengeProviderRejectedRequest(
                f'JS Challenge type "{request.type}" is not supported by {self.PROVIDER_NAME}')

    def bulk_solve(self, requests: list[JsChallengeRequest]) -> typing.Generator[JsChallengeProviderResponse, None, None]:
        """Solve multiple JS challenges and return the results"""
        validated_requests = []
        for request in requests:
            try:
                self.__validate_request(request)
                validated_requests.append(request)
            except JsChallengeProviderRejectedRequest as e:
                yield JsChallengeProviderResponse(request=request, error=e)
                continue
        yield from self._real_bulk_solve(validated_requests)

    @abc.abstractmethod
    def _real_bulk_solve(self, requests: list[JsChallengeRequest]) -> typing.Generator[JsChallengeProviderResponse, None, None]:
        """Subclasses can override this method to handle bulk solving"""
        raise NotImplementedError(f'{self.PROVIDER_NAME} does not implement bulk solving')

    def _get_player(self, video_id, player_url):
        try:
            return self.ie._load_player(
                video_id=video_id,
                player_url=player_url,
                fatal=True,
            )
        except ExtractorError as e:
            raise JsChallengeProviderError(
                f'Failed to load player for JS challenge: {e}') from e


def register_provider(provider: type[JsChallengeProvider]):
    """Register a JsChallengeProvider class"""
    return register_provider_generic(
        provider=provider,
        base_class=JsChallengeProvider,
        registry=_jsc_providers.value,
    )


def register_preference(*providers: type[JsChallengeProvider]) -> typing.Callable[[Preference], Preference]:
    """Register a preference for a JsChallengeProvider class."""
    return register_preference_generic(
        JsChallengeProvider,
        _jsc_preferences.value,
        *providers,
    )


if typing.TYPE_CHECKING:
    Preference = typing.Callable[[JsChallengeProvider, list[JsChallengeRequest]], int]
    __all__.append('Preference')