request_forwarder_classes.py

import datetime
from typing import Any, Dict, Tuple, Literal

import requests
from pydantic import BaseModel

RequestMethod = Literal['GET', 'POST', 'PUT', 'DELETE', 'HEAD', 'OPTIONS', 'PATCH']
ResponseType = Literal['raw', 'text', 'json']


class ForwardedRequest(BaseModel):
    method: RequestMethod
    url: str
    headers: Dict[str, str] | None = None
    data: Dict[str, Any] | Tuple[Tuple[str, Any]] | str | bytes | None = None
    params: Dict[str, Any] | Tuple[Tuple[str, Any]] | str | None = None
    auth: Tuple[str, str] | None = None
    cookies: Dict[str, str] | None = None
    json: Any | None = None
    response_content_type: ResponseType = 'text'

    @staticmethod
    def from_request(request: requests.Request, response_content_type: ResponseType = 'text') -> 'ForwardedRequest':
        return ForwardedRequest(
            method=request.method,
            url=request.url,
            headers=request.headers,
            data=request.data,
            params=request.params,
            auth=request.auth,
            cookies=request.cookies,
            json=request.json,
            response_content_type=response_content_type
        )

    @staticmethod
    def from_json(json_dict: Dict[str, Any]) -> 'ForwardedRequest':
        return ForwardedRequest(
            method=json_dict['method'],
            url=json_dict['url'],
            headers=json_dict['headers'],
            data=json_dict['data'],
            params=json_dict['params'],
            auth=json_dict['auth'],
            cookies=json_dict['cookies'],
            json=json_dict['json'],
            response_content_type=json_dict['response_content_type'],
        )

    def to_json_dict(self):
        return {
            "method": self.method,
            "url": self.url,
            "headers": self.headers,
            "data": self.data,
            "params": self.params,
            "auth": self.auth,
            "cookies": self.cookies,
            "json": self.json,
            "response_content_type": self.response_content_type
        }


class ForwardedResponse(BaseModel):
    status_code: int | None = None
    headers: Dict[str, Any] = dict[str, Any]()
    url: str | None = None
    encoding: str | None = None
    reason: str | None = None
    cookies: Dict[str, Any] = dict[str, Any]()
    elapsed: datetime.timedelta = datetime.timedelta(0)
    content: Any | None
    content_type: ResponseType = 'text'

    @staticmethod
    def from_response(response: requests.Response, response_content_type: ResponseType = 'text') -> 'ForwardedResponse':
        forwarded_response = ForwardedResponse(
            status_code=response.status_code,
            headers=response.headers,
            url=response.url,
            encoding=response.encoding,
            reason=response.reason,
            cookies=response.cookies,
            elapsed=response.elapsed,
            content=None,
            content_type=response_content_type
        )
        match forwarded_response.content_type:
            case 'raw':
                forwarded_response.content = response.content
            case 'text':
                forwarded_response.content = response.text
            case 'json':
                forwarded_response.content = response.json()
        return forwarded_response

    @staticmethod
    def from_json(json_dict: Dict[str, Any]) -> 'ForwardedResponse':
        return ForwardedResponse(
            status_code=json_dict['status_code'],
            headers=json_dict['headers'],
            url=json_dict['url'],
            encoding=json_dict['encoding'],
            reason=json_dict['reason'],
            cookies=json_dict['cookies'],
            elapsed=json_dict['elapsed'],
            content=json_dict['content'],
            content_type=json_dict['content_type'],
        )

    def to_json_dict(self):
        def serialize_timedelta(timedelta: datetime.timedelta):
            return {
                'days': timedelta.days,
                'seconds': timedelta.seconds,
                'microseconds': timedelta.microseconds
            }

        return {
            "status_code": self.status_code,
            "headers": self.headers,
            "url": self.url,
            "encoding": self.encoding,
            "reason": self.reason,
            "cookies": self.cookies,
            "elapsed": serialize_timedelta(self.elapsed),
            "content": self.content,
            "content_type": self.content_type,
        }