colab_tunnel/_registry.py

"""
Реестр провайдеров туннелей.

Регистрация нового провайдера:

@tunnel_provider('myprovider')
def get_myprovider_url(port: int) -> str: ...

Отключить провайдера (сервис временно недоступен):

@tunnel_provider('tunwg', enabled=False)
# или в рантайме:
disable_provider('tunwg')

Включить обратно:

enable_provider('tunwg')

Посмотреть статус всех:

list_providers()
"""
from typing import Callable

# Активные провайдеры — используются get_share_link(), try_all(), benchmark()
proxies_functions: dict[str, Callable[[int], str]] = {}

# Все зарегистрированные провайдеры (активные + отключённые)
_all_providers: dict[str, Callable[[int], str]] = {}

# Имена отключённых провайдеров
_disabled_providers: set[str] = set()


def tunnel_provider(name: str, *, enabled: bool = True) -> Callable:
    """
    Декоратор регистрации провайдера туннеля.

    Args:
        name:    Уникальное имя провайдера.
        enabled: False — провайдер зарегистрирован, но не активен.
                 Включить позднее через enable_provider(name).
    """
    def decorator(func: Callable[[int], str]) -> Callable[[int], str]:
        _all_providers[name] = func
        if enabled:
            proxies_functions[name] = func
        else:
            _disabled_providers.add(name)
        return func
    return decorator


def disable_provider(name: str) -> None:
    """
    Деактивирует провайдера: убирает из активного реестра.
    Функция сохраняется — включить обратно через enable_provider().

    Raises:
        KeyError: Провайдер с таким именем не зарегистрирован.
    """
    if name not in _all_providers:
        raise KeyError(f'Провайдер "{name}" не зарегистрирован. Доступные: {list(_all_providers)}')
    proxies_functions.pop(name, None)
    _disabled_providers.add(name)


def enable_provider(name: str) -> None:
    """
    Активирует ранее отключённого провайдера.

    Raises:
        KeyError: Провайдер с таким именем не зарегистрирован.
    """
    if name not in _all_providers:
        raise KeyError(f'Провайдер "{name}" не зарегистрирован. Доступные: {list(_all_providers)}')
    _disabled_providers.discard(name)
    proxies_functions[name] = _all_providers[name]


def list_providers() -> dict[str, bool]:
    """
    Возвращает статус всех зарегистрированных провайдеров.

    Returns:
        Словарь {имя: активен}, в порядке регистрации.

    Пример вывода:
        {'cloudflared': True, 'tunwg': False, 'bore': True, ...}
    """
    return {name: (name not in _disabled_providers) for name in _all_providers}