import requests from requests.adapters import HTTPAdapter from typing import List, Dict, Any, Tuple from rest_client_models import GetListResponse, GeneralResponse, MessageResponse from utils import get_size_from_dict, NetworkingSize class TrafficAdapter(HTTPAdapter): def __init__(self, *args, **kwargs): self.last_request_size = 0 self.last_response_size = 0 super().__init__(*args, **kwargs) def send(self, request, **kwargs): self.last_request_size = len(request.method.encode()) + len(request.url.encode()) + 12 # "HTTP/1.1\r\n" if request.headers: for name, value in request.headers.items(): self.last_request_size += len(name.encode()) + len(value.encode()) + 4 # ": " + "\r\n" self.last_request_size += 2 # "\r\n" if request.body: if isinstance(request.body, bytes): self.last_request_size += len(request.body) else: self.last_request_size += len(str(request.body).encode()) response = super().send(request, **kwargs) self.last_response_size = len(f"HTTP/1.1 {response.status_code}".encode()) + 2 if response.headers: for name, value in response.headers.items(): self.last_response_size += len(name.encode()) + len(value.encode()) + 4 self.last_response_size += 2 # "\r\n" self.last_response_size += len(response.content) return response class TrafficSession(requests.Session): def __init__(self): super().__init__() self.traffic_adapter = TrafficAdapter() self.mount('http://', self.traffic_adapter) self.mount('https://', self.traffic_adapter) def get_last_traffic(self): return { 'request_bytes': self.traffic_adapter.last_request_size, 'response_bytes': self.traffic_adapter.last_response_size } class RestClient: def __init__(self, base_url): self.base_url = base_url # self.session = requests.Session() self.session = TrafficSession() def get_list(self, pages: int = 1, per_page: int = 100, list_data_limit=30) -> Tuple[GetListResponse, NetworkingSize]: response = self.session.get( f"{self.base_url}/api/get_list", params={"pages": pages, "per_page": per_page, "list_data_limit": list_data_limit} ) response.raise_for_status() data = response.json() return GetListResponse(**data), get_size_from_dict(self.session.get_last_traffic()) def add_book(self, books: List[Dict[str, Any]], test_only=False) -> Tuple[GeneralResponse, NetworkingSize]: data = {"books": books, "test_only": test_only} response = self.session.post( f"{self.base_url}/api/add_book", json=data ) response.raise_for_status() return GeneralResponse(**response.json()), get_size_from_dict(self.session.get_last_traffic()) def delete_books(self, book_ids: List[int], delete_last_count=-1) -> Tuple[GeneralResponse, NetworkingSize]: data = {"book_ids": book_ids, "delete_last_count": delete_last_count} response = self.session.delete( f"{self.base_url}/api/delete_books", json=data ) response.raise_for_status() return GeneralResponse(**response.json()), get_size_from_dict(self.session.get_last_traffic()) def update_book(self, book_data: Dict[str, Any]) -> Tuple[GeneralResponse, NetworkingSize]: data = {"book": book_data} response = self.session.put( f"{self.base_url}/api/update_book", json=data ) response.raise_for_status() return GeneralResponse(**response.json()), get_size_from_dict(self.session.get_last_traffic()) def ping(self): self.session.get(f"{self.base_url}/api/ping") def close(self): self.session.close()