ayaka.driver.ayakabot.model

  1import json
  2import dataclasses
  3from enum import Enum
  4from http.cookiejar import Cookie, CookieJar
  5from typing import (
  6    IO,
  7    Any,
  8    Dict,
  9    List,
 10    Tuple,
 11    Union,
 12    Mapping,
 13    Callable,
 14    Iterator,
 15    Optional,
 16    MutableMapping,
 17)
 18
 19from yarl import URL
 20from multidict import CIMultiDict
 21
 22RawURL = Tuple[bytes, bytes, Optional[int], bytes]
 23
 24SimpleQuery = Union[str, int, float]
 25QueryVariable = Union[SimpleQuery, List[SimpleQuery]]
 26QueryTypes = Union[
 27    None, str, Mapping[str, QueryVariable], List[Tuple[str, QueryVariable]]
 28]
 29
 30HeaderTypes = Union[
 31    None,
 32    CIMultiDict[str],
 33    Dict[str, str],
 34    List[Tuple[str, str]],
 35]
 36
 37CookieTypes = Union[None, "Cookies", CookieJar, Dict[str, str], List[Tuple[str, str]]]
 38
 39ContentTypes = Union[str, bytes, None]
 40DataTypes = Union[dict, None]
 41FileContent = Union[IO[bytes], bytes]
 42FileType = Tuple[Optional[str], FileContent, Optional[str]]
 43FileTypes = Union[
 44    # file (or bytes)
 45    FileContent,
 46    # (filename, file (or bytes))
 47    Tuple[Optional[str], FileContent],
 48    # (filename, file (or bytes), content_type)
 49    FileType,
 50]
 51FilesTypes = Union[Dict[str, FileTypes], List[Tuple[str, FileTypes]], None]
 52
 53
 54class HTTPVersion(Enum):
 55    H10 = "1.0"
 56    H11 = "1.1"
 57    H2 = "2"
 58
 59
 60class Request:
 61    def __init__(
 62        self,
 63        method: Union[str, bytes],
 64        url: Union["URL", str, RawURL],
 65        *,
 66        params: QueryTypes = None,
 67        headers: HeaderTypes = None,
 68        cookies: CookieTypes = None,
 69        content: ContentTypes = None,
 70        data: DataTypes = None,
 71        json: Any = None,
 72        files: FilesTypes = None,
 73        version: Union[str, HTTPVersion] = HTTPVersion.H11,
 74        timeout: Optional[float] = None,
 75        proxy: Optional[str] = None,
 76    ):
 77        # method
 78        self.method: str = (
 79            method.decode("ascii").upper()
 80            if isinstance(method, bytes)
 81            else method.upper()
 82        )
 83        # http version
 84        self.version: HTTPVersion = HTTPVersion(version)
 85        # timeout
 86        self.timeout: Optional[float] = timeout
 87        # proxy
 88        self.proxy: Optional[str] = proxy
 89
 90        # url
 91        if isinstance(url, tuple):
 92            scheme, host, port, path = url
 93            url = URL.build(
 94                scheme=scheme.decode("ascii"),
 95                host=host.decode("ascii"),
 96                port=port,
 97                path=path.decode("ascii"),
 98            )
 99        else:
100            url = URL(url)
101
102        if params is not None:
103            url = url.update_query(params)
104        self.url: URL = url
105
106        # headers
107        self.headers: CIMultiDict[str]
108        if headers is not None:
109            self.headers = CIMultiDict(headers)
110        else:
111            self.headers = CIMultiDict()
112
113        # cookies
114        self.cookies = Cookies(cookies)
115
116        # body
117        self.content: ContentTypes = content
118        self.data: DataTypes = data
119        self.json: Any = json
120        self.files: Optional[List[Tuple[str, FileType]]] = None
121        if files:
122            self.files = []
123            files_ = files.items() if isinstance(files, dict) else files
124            for name, file_info in files_:
125                if not isinstance(file_info, tuple):
126                    self.files.append((name, (None, file_info, None)))
127                elif len(file_info) == 2:
128                    self.files.append((name, (file_info[0], file_info[1], None)))
129                else:
130                    self.files.append((name, file_info))  # type: ignore
131
132    def __repr__(self) -> str:
133        class_name = self.__class__.__name__
134        url = str(self.url)
135        return f"<{class_name}({self.method!r}, {url!r})>"
136
137
138class Cookies(MutableMapping):
139    def __init__(self, cookies: CookieTypes = None) -> None:
140        self.jar: CookieJar = cookies if isinstance(cookies, CookieJar) else CookieJar()
141        if cookies is not None and not isinstance(cookies, CookieJar):
142            if isinstance(cookies, dict):
143                for key, value in cookies.items():
144                    self.set(key, value)
145            elif isinstance(cookies, list):
146                for key, value in cookies:
147                    self.set(key, value)
148            elif isinstance(cookies, Cookies):
149                for cookie in cookies.jar:
150                    self.jar.set_cookie(cookie)
151            else:
152                raise TypeError(f"Cookies must be dict or list, not {type(cookies)}")
153
154    def set(self, name: str, value: str, domain: str = "", path: str = "/") -> None:
155        cookie = Cookie(
156            version=0,
157            name=name,
158            value=value,
159            port=None,
160            port_specified=False,
161            domain=domain,
162            domain_specified=bool(domain),
163            domain_initial_dot=domain.startswith("."),
164            path=path,
165            path_specified=bool(path),
166            secure=False,
167            expires=None,
168            discard=True,
169            comment=None,
170            comment_url=None,
171            rest={},
172            rfc2109=False,
173        )
174        self.jar.set_cookie(cookie)
175
176    def get(
177        self,
178        name: str,
179        default: Optional[str] = None,
180        domain: str = None,
181        path: str = None,
182    ) -> Optional[str]:
183        value: Optional[str] = None
184        for cookie in self.jar:
185            if (
186                cookie.name == name
187                and (domain is None or cookie.domain == domain)
188                and (path is None or cookie.path == path)
189            ):
190                if value is not None:
191                    message = f"Multiple cookies exist with name={name}"
192                    raise ValueError(message)
193                value = cookie.value
194
195        return default if value is None else value
196
197    def delete(
198        self, name: str, domain: Optional[str] = None, path: Optional[str] = None
199    ) -> None:
200        if domain is not None and path is not None:
201            return self.jar.clear(domain, path, name)
202
203        remove = [
204            cookie
205            for cookie in self.jar
206            if cookie.name == name
207            and (domain is None or cookie.domain == domain)
208            and (path is None or cookie.path == path)
209        ]
210
211        for cookie in remove:
212            self.jar.clear(cookie.domain, cookie.path, cookie.name)
213
214    def clear(self, domain: Optional[str] = None, path: Optional[str] = None) -> None:
215        self.jar.clear(domain, path)
216
217    def update(self, cookies: CookieTypes = None) -> None:
218        cookies = Cookies(cookies)
219        for cookie in cookies.jar:
220            self.jar.set_cookie(cookie)
221
222    def __setitem__(self, name: str, value: str) -> None:
223        return self.set(name, value)
224
225    def __getitem__(self, name: str) -> str:
226        value = self.get(name)
227        if value is None:
228            raise KeyError(name)
229        return value
230
231    def __delitem__(self, name: str) -> None:
232        return self.delete(name)
233
234    def __len__(self) -> int:
235        return len(self.jar)
236
237    def __iter__(self) -> Iterator[Cookie]:
238        return (cookie for cookie in self.jar)
239
240    def __repr__(self) -> str:
241        cookies_repr = ", ".join(
242            [
243                f"<Cookie {cookie.name}={cookie.value} for {cookie.domain} />"
244                for cookie in self.jar
245            ]
246        )
247
248        return f"<Cookies [{cookies_repr}]>"
249
250
251@dataclasses.dataclass
252class WebSocketServerSetup:
253    """WebSocket 服务器路由配置。"""
254
255    path: URL  # path should not be absolute, check it by URL.is_absolute() == False
256    name: str
257    handle_func: Callable
258
259class DataclassEncoder(json.JSONEncoder):
260    """在JSON序列化 `Message` (List[Dataclass]) 时使用的 `JSONEncoder`"""
261
262    def default(self, o):
263        if dataclasses.is_dataclass(o):
264            return dataclasses.asdict(o)
265        return super().default(o)