-
Notifications
You must be signed in to change notification settings - Fork 643
/
requests_backend.py
165 lines (133 loc) · 5.3 KB
/
requests_backend.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
from __future__ import annotations
import dataclasses
from typing import Any, Dict, Optional, TYPE_CHECKING, Union
import requests
from requests import PreparedRequest
from requests.auth import AuthBase
from requests.structures import CaseInsensitiveDict
from requests_toolbelt.multipart.encoder import MultipartEncoder # type: ignore
from . import protocol
class TokenAuth:
def __init__(self, token: str):
self.token = token
class OAuthTokenAuth(TokenAuth, AuthBase):
def __call__(self, r: PreparedRequest) -> PreparedRequest:
r.headers["Authorization"] = f"Bearer {self.token}"
r.headers.pop("PRIVATE-TOKEN", None)
r.headers.pop("JOB-TOKEN", None)
return r
class PrivateTokenAuth(TokenAuth, AuthBase):
def __call__(self, r: PreparedRequest) -> PreparedRequest:
r.headers["PRIVATE-TOKEN"] = self.token
r.headers.pop("JOB-TOKEN", None)
r.headers.pop("Authorization", None)
return r
class JobTokenAuth(TokenAuth, AuthBase):
def __call__(self, r: PreparedRequest) -> PreparedRequest:
r.headers["JOB-TOKEN"] = self.token
r.headers.pop("PRIVATE-TOKEN", None)
r.headers.pop("Authorization", None)
return r
@dataclasses.dataclass
class SendData:
content_type: str
data: Optional[Union[Dict[str, Any], MultipartEncoder]] = None
json: Optional[Union[Dict[str, Any], bytes]] = None
def __post_init__(self) -> None:
if self.json is not None and self.data is not None:
raise ValueError(
f"`json` and `data` are mutually exclusive. Only one can be set. "
f"json={self.json!r} data={self.data!r}"
)
class RequestsResponse(protocol.BackendResponse):
def __init__(self, response: requests.Response) -> None:
self._response: requests.Response = response
@property
def response(self) -> requests.Response:
return self._response
@property
def status_code(self) -> int:
return self._response.status_code
@property
def headers(self) -> CaseInsensitiveDict[str]:
return self._response.headers
@property
def content(self) -> bytes:
return self._response.content
@property
def reason(self) -> str:
return self._response.reason
def json(self) -> Any:
return self._response.json()
class RequestsBackend(protocol.Backend):
def __init__(self, session: Optional[requests.Session] = None) -> None:
self._client: requests.Session = session or requests.Session()
@property
def client(self) -> requests.Session:
return self._client
@staticmethod
def prepare_send_data(
files: Optional[Dict[str, Any]] = None,
post_data: Optional[Union[Dict[str, Any], bytes]] = None,
raw: bool = False,
) -> SendData:
if files:
if post_data is None:
post_data = {}
else:
# When creating a `MultipartEncoder` instance with data-types
# which don't have an `encode` method it will cause an error:
# object has no attribute 'encode'
# So convert common non-string types into strings.
if TYPE_CHECKING:
assert isinstance(post_data, dict)
for k, v in post_data.items():
if isinstance(v, bool):
v = int(v)
if isinstance(v, (complex, float, int)):
post_data[k] = str(v)
post_data["file"] = files.get("file")
post_data["avatar"] = files.get("avatar")
data = MultipartEncoder(fields=post_data)
return SendData(data=data, content_type=data.content_type)
if raw and post_data:
return SendData(data=post_data, content_type="application/octet-stream")
return SendData(json=post_data, content_type="application/json")
def http_request(
self,
method: str,
url: str,
json: Optional[Union[Dict[str, Any], bytes]] = None,
data: Optional[Union[Dict[str, Any], MultipartEncoder]] = None,
params: Optional[Any] = None,
timeout: Optional[float] = None,
verify: Optional[Union[bool, str]] = True,
stream: Optional[bool] = False,
**kwargs: Any,
) -> RequestsResponse:
"""Make HTTP request
Args:
method: The HTTP method to call ('get', 'post', 'put', 'delete', etc.)
url: The full URL
data: The data to send to the server in the body of the request
json: Data to send in the body in json by default
timeout: The timeout, in seconds, for the request
verify: Whether SSL certificates should be validated. If
the value is a string, it is the path to a CA file used for
certificate validation.
stream: Whether the data should be streamed
Returns:
A requests Response object.
"""
response: requests.Response = self._client.request(
method=method,
url=url,
params=params,
data=data,
timeout=timeout,
stream=stream,
verify=verify,
json=json,
**kwargs,
)
return RequestsResponse(response=response)