codekingpro/portable-devtools

1# --------------------------------------------------------------------------2#3# Copyright (c) Microsoft Corporation. All rights reserved.4#5# The MIT License (MIT)6#7# Permission is hereby granted, free of charge, to any person obtaining a copy8# of this software and associated documentation files (the ""Software""), to9# deal in the Software without restriction, including without limitation the10# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or11# sell copies of the Software, and to permit persons to whom the Software is12# furnished to do so, subject to the following conditions:13#14# The above copyright notice and this permission notice shall be included in15# all copies or substantial portions of the Software.16#17# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING22# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS23# IN THE SOFTWARE.24#25# --------------------------------------------------------------------------26from typing import Any, Callable, AsyncIterator, Optional27 28import aiohttp29from multidict import CIMultiDict30 31from . import AsyncHTTPSender, ClientRequest, AsyncClientResponse32 33# Matching requests, because why not?34CONTENT_CHUNK_SIZE = 10 * 102435 36 37class AioHTTPSender(AsyncHTTPSender):38    """AioHttp HTTP sender implementation.39    """40 41    def __init__(self, *, loop=None):42        self._session = aiohttp.ClientSession(loop=loop)43 44    async def __aenter__(self):45        await self._session.__aenter__()46        return self47 48    async def __aexit__(self, *exc_details):  # pylint: disable=arguments-differ49        await self._session.__aexit__(*exc_details)50 51    async def send(self, request: ClientRequest, **config: Any) -> AsyncClientResponse:52        """Send the request using this HTTP sender.53 54        Will pre-load the body into memory to be available with a sync method.55        pass stream=True to avoid this behavior.56        """57        result = await self._session.request(58            request.method,59            request.url,60            **config61        )62        response = AioHttpClientResponse(request, result)63        if not config.get("stream", False):64            await response.load_body()65        return response66 67 68class AioHttpClientResponse(AsyncClientResponse):69    def __init__(self, request: ClientRequest, aiohttp_response: aiohttp.ClientResponse) -> None:70        super(AioHttpClientResponse, self).__init__(request, aiohttp_response)71        # https://aiohttp.readthedocs.io/en/stable/client_reference.html#aiohttp.ClientResponse72        self.status_code = aiohttp_response.status73        self.headers = CIMultiDict(aiohttp_response.headers)74        self.reason = aiohttp_response.reason75        self._body = None76 77    def body(self) -> bytes:78        """Return the whole body as bytes in memory.79        """80        if not self._body:81            raise ValueError("Body is not available. Call async method load_body, or do your call with stream=False.")82        return self._body83 84    async def load_body(self) -> None:85        """Load in memory the body, so it could be accessible from sync methods."""86        self._body = await self.internal_response.read()87 88    def raise_for_status(self):89        self.internal_response.raise_for_status()90 91    def stream_download(self, chunk_size: Optional[int] = None, callback: Optional[Callable] = None) -> AsyncIterator[bytes]:92        """Generator for streaming request body data.93        """94        chunk_size = chunk_size or CONTENT_CHUNK_SIZE95        async def async_gen(resp):96            while True:97                chunk = await resp.content.read(chunk_size)98                if not chunk:99                    break100                callback(chunk, resp)101        return async_gen(self.internal_response)102