102 lines
3.9 KiB
Python
102 lines
3.9 KiB
Python
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()
|