From 80258fe2d44e39118bdd82571c52a279569d241e Mon Sep 17 00:00:00 2001 From: yanyongyu Date: Fri, 24 Dec 2021 23:15:28 +0800 Subject: [PATCH] :alembic: add support for data json file --- nonebot/drivers/_model.py | 40 +++++++++++++++++++++++++++++++++----- nonebot/drivers/aiohttp.py | 8 +++++++- nonebot/drivers/fastapi.py | 30 +++++++++++++++++++++++++--- nonebot/drivers/httpx.py | 3 +++ nonebot/drivers/quart.py | 22 +++++++++++++++++---- 5 files changed, 90 insertions(+), 13 deletions(-) diff --git a/nonebot/drivers/_model.py b/nonebot/drivers/_model.py index d217ab33..a5998373 100644 --- a/nonebot/drivers/_model.py +++ b/nonebot/drivers/_model.py @@ -2,6 +2,8 @@ import abc from enum import Enum from http.cookiejar import Cookie, CookieJar from typing import ( + IO, + Any, Dict, List, Tuple, @@ -9,7 +11,6 @@ from typing import ( Mapping, Iterator, Optional, - Sequence, MutableMapping, ) @@ -19,21 +20,34 @@ from multidict import CIMultiDict RawURL = Tuple[bytes, bytes, Optional[int], bytes] SimpleQuery = Union[str, int, float] -QueryVariable = Union[SimpleQuery, Sequence[SimpleQuery]] +QueryVariable = Union[SimpleQuery, List[SimpleQuery]] QueryTypes = Union[ - None, str, Mapping[str, QueryVariable], Sequence[Tuple[str, QueryVariable]] + None, str, Mapping[str, QueryVariable], List[Tuple[str, QueryVariable]] ] HeaderTypes = Union[ None, CIMultiDict[str], Dict[str, str], - Sequence[Tuple[str, str]], + List[Tuple[str, str]], ] -ContentTypes = Union[str, bytes, None] CookieTypes = Union[None, "Cookies", CookieJar, Dict[str, str], List[Tuple[str, str]]] +ContentTypes = Union[str, bytes, None] +DataTypes = Union[dict, None] +FileContent = Union[IO[bytes], bytes] +FileType = Tuple[Optional[str], FileContent, Optional[str]] +FileTypes = Union[ + # file (or bytes) + FileContent, + # (filename, file (or bytes)) + Tuple[Optional[str], FileContent], + # (filename, file (or bytes), content_type) + FileType, +] +FilesTypes = Union[Dict[str, FileTypes], List[Tuple[str, FileTypes]], None] + class HTTPVersion(Enum): H10 = "1.0" @@ -51,6 +65,9 @@ class Request: headers: HeaderTypes = None, cookies: CookieTypes = None, content: ContentTypes = None, + data: DataTypes = None, + json: Any = None, + files: FilesTypes = None, version: Union[str, HTTPVersion] = HTTPVersion.H11, timeout: Optional[float] = None, ): @@ -93,6 +110,19 @@ class Request: # body self.content: ContentTypes = content + self.data: DataTypes = data + self.json: Any = json + self.files: Optional[List[Tuple[str, FileType]]] = None + if files: + self.files = [] + files_ = files.items() if isinstance(files, dict) else files + for name, file_info in files_: + if not isinstance(file_info, tuple): + self.files.append((name, (None, file_info, None))) + elif len(file_info) == 2: + self.files.append((name, (file_info[0], file_info[1], None))) + else: + self.files.append((name, file_info)) # type: ignore def __repr__(self) -> str: class_name = self.__class__.__name__ diff --git a/nonebot/drivers/aiohttp.py b/nonebot/drivers/aiohttp.py index f935342c..3f0784d9 100644 --- a/nonebot/drivers/aiohttp.py +++ b/nonebot/drivers/aiohttp.py @@ -35,11 +35,17 @@ class Mixin(ForwardMixin): raise RuntimeError(f"Unsupported HTTP version: {setup.version}") timeout = aiohttp.ClientTimeout(setup.timeout) + files = None + if setup.files: + files = aiohttp.FormData() + for name, file in setup.files: + files.add_field(name, file[1], content_type=file[2], filename=file[0]) async with aiohttp.ClientSession(version=version) as session: async with session.request( setup.method, setup.url, - data=setup.content, + data=setup.content or setup.data or files, + json=setup.json, headers=setup.headers, timeout=timeout, ) as response: diff --git a/nonebot/drivers/fastapi.py b/nonebot/drivers/fastapi.py index 41bd6c71..57170a59 100644 --- a/nonebot/drivers/fastapi.py +++ b/nonebot/drivers/fastapi.py @@ -11,17 +11,17 @@ FastAPI 驱动适配 """ import logging -from typing import List, Callable, Optional +from typing import Any, List, Tuple, Callable, Optional import uvicorn from pydantic import BaseSettings from fastapi.responses import Response -from fastapi import FastAPI, Request, status +from fastapi import FastAPI, Request, UploadFile, status from starlette.websockets import WebSocket, WebSocketState +from ._model import FileTypes from nonebot.config import Env from nonebot.typing import overrides -from nonebot.utils import escape_tag from nonebot.config import Config as NoneBotConfig from nonebot.drivers import Request as BaseRequest from nonebot.drivers import WebSocket as BaseWebSocket @@ -238,12 +238,36 @@ class Driver(ReverseDriver): request: Request, setup: HTTPServerSetup, ) -> Response: + json: Any = None + try: + json = await request.json() + except Exception: + pass + + data: Optional[dict] = None + files: Optional[List[Tuple[str, FileTypes]]] = None + try: + form = await request.form() + data = {} + files = [] + for key, value in form.multi_items(): + if isinstance(value, UploadFile): + files.append( + (key, (value.filename, value.file, value.content_type)) + ) + else: + data[key] = value + except Exception: + pass http_request = BaseRequest( request.method, str(request.url), headers=request.headers.items(), cookies=request.cookies, content=await request.body(), + data=data, + json=json, + files=files, version=request.scope["http_version"], ) diff --git a/nonebot/drivers/httpx.py b/nonebot/drivers/httpx.py index 09234710..f6789ab4 100644 --- a/nonebot/drivers/httpx.py +++ b/nonebot/drivers/httpx.py @@ -30,6 +30,9 @@ class Mixin(ForwardMixin): setup.method, str(setup.url), content=setup.content, + data=setup.data, + json=setup.json, + files=setup.files, headers=tuple(setup.headers.items()), timeout=setup.timeout, ) diff --git a/nonebot/drivers/quart.py b/nonebot/drivers/quart.py index 89822ee6..68eedd67 100644 --- a/nonebot/drivers/quart.py +++ b/nonebot/drivers/quart.py @@ -8,15 +8,14 @@ Quart 驱动适配 https://pgjones.gitlab.io/quart/index.html """ -from typing import List, TypeVar, Callable, Optional, Coroutine +from typing import List, Tuple, TypeVar, Callable, Optional, Coroutine import uvicorn from pydantic import BaseSettings +from ._model import FileTypes from nonebot.config import Env -from nonebot.log import logger from nonebot.typing import overrides -from nonebot.utils import escape_tag from nonebot.config import Config as NoneBotConfig from nonebot.drivers import Request as BaseRequest from nonebot.drivers import WebSocket as BaseWebSocket @@ -24,9 +23,9 @@ from nonebot.drivers import ReverseDriver, HTTPServerSetup, WebSocketServerSetup try: from quart import request as _request - import werkzeug.exceptions as exceptions from quart import websocket as _websocket from quart import Quart, Request, Response + from quart.datastructures import FileStorage from quart import Websocket as QuartWebSocket except ImportError: raise ValueError("Please install Quart by using `pip install nonebot2[quart]`") @@ -213,6 +212,18 @@ class Driver(ReverseDriver): async def _handle_http(self, setup: HTTPServerSetup) -> Response: request: Request = _request + json = None + if request.is_json: + json = await request.get_json() + + data = await request.form + files_dict = await request.files + files: List[Tuple[str, FileTypes]] = [] + key: str + value: FileStorage + for key, value in files_dict.items(): + files.append((key, (value.filename, value.stream, value.content_type))) + http_request = BaseRequest( request.method, request.url, @@ -221,6 +232,9 @@ class Driver(ReverseDriver): content=await request.get_data( cache=False, as_text=False, parse_form_data=False ), + data=data or None, + json=json, + files=files or None, version=request.http_version, )