py_client_generator/base.py - master - semak/py-client-generator

py_client_generator/base.py
280 строк9 KB
semak
Stream в ответе
22 июн 2026, 17:51
22 июн 2026, 17:51c993176
import json
from abc import ABC, abstractmethod
from typing import Type, Dict, Any, Union, TypeVar, AsyncGenerator
import httpx
from httpx import Request, Response
from pydantic import BaseModel
import base64


class NothingToParse(Exception):
    def __init__(self, code, message):
        self.code = code
        self.message = message


class ClientParseError(Exception):
    def __init__(self, code, message):
        self.code = code
        self.message = message


class NeedAuthException(Exception):
    def __init__(self, auth_class):
        self.auth_class = auth_class

    def __str__(self):
        return f"Need to init authorization with {self.auth_class}"


T = TypeVar("T", bound=BaseModel)
DefaultT = TypeVar("DefaultT", bound=BaseModel)


class ClientHttpError(httpx.HTTPStatusError):
    def __init__(
            self,
            message: str,
            *,
            request: Request,
            response: Response,
            status_model_map: Dict[int, Type[T] | None],
            default_model: Type[DefaultT] | None = None,
    ) -> None:
        super().__init__(message, request=request, response=response)
        self.raw_data = response.text
        self.parsing_error = None

        try:
            self.model_to_parse = status_model_map.get(
                response.status_code, default_model
            )
            self.parsed_data = (
                self.model_to_parse.model_validate(response.json())
                if self.model_to_parse
                else None
            )
        except Exception as ex:
            self.parsed_data = None
            self.parsing_error = str(ex)


class BaseRawClient:
    def __init__(
            self,
            client: httpx.AsyncClient,
    ):
        """
        Args:
            client: Настроенный экземпляр httpx.AsyncClient
        """
        self._client = client


TClient = TypeVar("TClient", bound=BaseRawClient)


class BaseClient:
    def __init__(self, raw_client: "TClient"):
        self.raw_client: TClient = raw_client
        self.ignore_auth_settings = False

    @staticmethod
    def parse(
            response: httpx.Response,
            status_model_map: Dict[int, Type[T] | None],
            default_model: Type[DefaultT] | None,
    ) -> Union[T, DefaultT]:
        try:
            response.raise_for_status()
        except httpx.HTTPStatusError as http_err:
            raise ClientHttpError(
                "HTTPError",
                request=http_err.request,
                response=http_err.response,
                status_model_map=status_model_map,
                default_model=default_model,
            )

        model_to_parse = status_model_map.get(response.status_code) or default_model
        if model_to_parse is None:
            raise ClientParseError(
                code=str(response.status_code), message=response.text
            )
        return model_to_parse.model_validate(response.json())

    @staticmethod
    async def parse_stream(
            response: httpx.Response,
            status_model_map: Dict[int, Type[T] | None],
            default_model: Type[DefaultT] | None,
    ) -> AsyncGenerator[Union[T, DefaultT], None]:
        model_to_parse = status_model_map.get(response.status_code) or default_model
        if model_to_parse is None:
            raise ClientParseError(
                code=str(response.status_code), message=response.text
            )
        try:
            if not response.is_success:
                error_body = ""
                try:
                    async for line in response.aiter_lines():
                        error_body += line
                except Exception:
                    pass

                response.raise_for_status()

            async for line in response.aiter_lines():
                if line.startswith("data: "):
                    data = line[6:].strip()
                    if data == "[DONE]":
                        break
                    if data:
                        js = json.loads(data)
                        yield model_to_parse.model_validate(js)
            response.raise_for_status()
        except httpx.HTTPStatusError as http_err:
            raise ClientHttpError(
                "HTTPError",
                request=http_err.request,
                response=http_err.response,
                status_model_map=status_model_map,
                default_model=default_model,
            )
        '''
        model_to_parse = status_model_map.get(response.status_code) or default_model
        if model_to_parse is None:
            raise ClientParseError(
                code=str(response.status_code), message=response.text
            )
        return model_to_parse.model_validate(response.json())
        '''

class BaseSecurityAuth(ABC):
    @property
    @abstractmethod
    def headers(self) -> Dict[str, str]:
        raise NotImplementedError()


class BearerAuth(BaseSecurityAuth):
    """Аутентификация через Bearer токен."""

    def __init__(self, bearer_token: str):
        """
        Args:
            bearer_token: Bearer токен для аутентификации
        """
        self.bearer_token = bearer_token

    @property
    def headers(self) -> Dict[str, str]:
        """
        Возвращает заголовки для Bearer аутентификации.

        Returns:
            Dict[str, str]: Заголовок Authorization с Bearer токеном
        """
        return {"Authorization": f"Bearer {self.bearer_token}"}

    def __repr__(self) -> str:
        token_preview = (
            self.bearer_token[:10] + "..."
            if len(self.bearer_token) > 10
            else self.bearer_token
        )
        return f"BearerAuth(token='{token_preview}')"


class BasicAuth(BaseSecurityAuth):

    def __init__(self, username: str | None = None, password: str | None = None,
                 encoded_credentials: str | None = None):
        self.username = username
        self.password = password
        self.encoded_credentials = encoded_credentials

    @property
    def headers(self) -> Dict[str, str]:
        """
        Возвращает заголовки для Basic аутентификации.

        Returns:
            Dict[str, str]: Заголовок Authorization с Basic credentials
        """

        credentials = f"{self.username}:{self.password}"
        self.encoded_credentials = self.encoded_credentials or base64.b64encode(credentials.encode()).decode()
        return {"Authorization": f"Basic {self.encoded_credentials}"}

    def __repr__(self) -> str:
        return f"BasicAuth(username='{self.username}')"


class ApiKeyAuth(BaseSecurityAuth):
    """Аутентификация через API Key."""

    def __init__(
            self, api_key: str, header_name: str = "X-API-Key", location: str = "header"
    ):
        """
        Args:
            api_key: API ключ
            header_name: Имя заголовка (по умолчанию "X-API-Key")
            location: Место размещения - "header", "query" или "cookie"
        """
        self.api_key = api_key
        self.header_name = header_name
        self.location = location

        if location not in ["header", "query", "cookie"]:
            raise ValueError(
                f"Неподдерживаемое расположение API ключа: {location}. "
                f"Поддерживаются: 'header', 'query', 'cookie'"
            )

    @property
    def headers(self) -> Dict[str, str]:
        """
        Возвращает заголовки для API Key аутентификации.

        Returns:
            Dict[str, str]: Заголовок с API ключом (если location == "header")
                            или пустой dict для других расположений
        """
        if self.location == "header":
            return {self.header_name: self.api_key}
        return {}

    @property
    def query_params(self) -> Dict[str, str]:
        """
        Возвращает query параметры для API Key аутентификации.

        Returns:
            Dict[str, str]: Query параметр с API ключом (если location == "query")
                            или пустой dict для других расположений
        """
        if self.location == "query":
            return {self.header_name: self.api_key}
        return {}

    @property
    def cookies(self) -> Dict[str, str]:
        """
        Возвращает cookies для API Key аутентификации.

        Returns:
            Dict[str, str]: Cookie с API ключом (если location == "cookie")
                            или пустой dict для других расположений
        """
        if self.location == "cookie":
            return {self.header_name: self.api_key}
        return {}

    def __repr__(self) -> str:
        key_preview = (
            self.api_key[:10] + "..." if len(self.api_key) > 10 else self.api_key
        )
        return f"ApiKeyAuth(key='{key_preview}', header='{self.header_name}', location='{self.location}')"