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# --------------------------------------------------------------------------26import asyncio27from collections.abc import AsyncIterator28import functools29import logging30from typing import Any, Callable, Optional, AsyncIterator as AsyncIteratorType31 32from oauthlib import oauth233import requests34from requests.models import CONTENT_CHUNK_SIZE35 36from ..exceptions import (37    TokenExpiredError,38    ClientRequestError,39    raise_with_traceback)40from . import AsyncHTTPSender, ClientRequest, AsyncClientResponse41from .requests import (42    BasicRequestsHTTPSender,43    RequestsHTTPSender,44    HTTPRequestsClientResponse45)46 47 48_LOGGER = logging.getLogger(__name__)49 50 51class AsyncBasicRequestsHTTPSender(BasicRequestsHTTPSender, AsyncHTTPSender):  # type: ignore52 53    async def __aenter__(self):54        return super(AsyncBasicRequestsHTTPSender, self).__enter__()55 56    async def __aexit__(self, *exc_details):  # pylint: disable=arguments-differ57        return super(AsyncBasicRequestsHTTPSender, self).__exit__()58 59    async def send(self, request: ClientRequest, **kwargs: Any) -> AsyncClientResponse:  # type: ignore60        """Send the request using this HTTP sender.61        """62        # It's not recommended to provide its own session, and is mostly63        # to enable some legacy code to plug correctly64        session = kwargs.pop('session', self.session)65 66        loop = kwargs.get("loop", asyncio.get_event_loop())67        future = loop.run_in_executor(68            None,69            functools.partial(70                session.request,71                request.method,72                request.url,73                **kwargs74            )75        )76        try:77            return AsyncRequestsClientResponse(78                request,79                await future80            )81        except requests.RequestException as err:82            msg = "Error occurred in request."83            raise_with_traceback(ClientRequestError, msg, err)84 85class AsyncRequestsHTTPSender(AsyncBasicRequestsHTTPSender, RequestsHTTPSender):  # type: ignore86 87    async def send(self, request: ClientRequest, **kwargs: Any) -> AsyncClientResponse:  # type: ignore88        """Send the request using this HTTP sender.89        """90        requests_kwargs = self._configure_send(request, **kwargs)91        return await super(AsyncRequestsHTTPSender, self).send(request, **requests_kwargs)92 93 94class _MsrestStopIteration(Exception):95    pass96 97def _msrest_next(iterator):98    """"To avoid:99    TypeError: StopIteration interacts badly with generators and cannot be raised into a Future100    """101    try:102        return next(iterator)103    except StopIteration:104        raise _MsrestStopIteration()105 106class StreamDownloadGenerator(AsyncIterator):107 108    def __init__(self, response: requests.Response, user_callback: Optional[Callable] = None, block: Optional[int] = None) -> None:109        self.response = response110        self.block = block or CONTENT_CHUNK_SIZE111        self.user_callback = user_callback112        self.iter_content_func = self.response.iter_content(self.block)113 114    async def __anext__(self):115        loop = asyncio.get_event_loop()116        try:117            chunk = await loop.run_in_executor(118                None,119                _msrest_next,120                self.iter_content_func,121            )122            if not chunk:123                raise _MsrestStopIteration()124            if self.user_callback and callable(self.user_callback):125                self.user_callback(chunk, self.response)126            return chunk127        except _MsrestStopIteration:128            self.response.close()129            raise StopAsyncIteration()130        except Exception as err:131            _LOGGER.warning("Unable to stream download: %s", err)132            self.response.close()133            raise134 135class AsyncRequestsClientResponse(AsyncClientResponse, HTTPRequestsClientResponse):136 137    def stream_download(self, chunk_size: Optional[int] = None, callback: Optional[Callable] = None) -> AsyncIteratorType[bytes]:138        """Generator for streaming request body data.139 140        :param callback: Custom callback for monitoring progress.141        :param int chunk_size:142        """143        return StreamDownloadGenerator(144            self.internal_response,145            callback,146            chunk_size147        )148 149 150# Trio support151try:152    import trio153 154    class TrioStreamDownloadGenerator(AsyncIterator):155 156        def __init__(self, response: requests.Response, user_callback: Optional[Callable] = None, block: Optional[int] = None) -> None:157            self.response = response158            self.block = block or CONTENT_CHUNK_SIZE159            self.user_callback = user_callback160            self.iter_content_func = self.response.iter_content(self.block)161 162        async def __anext__(self):163            try:164                chunk = await trio.to_thread.run_sync(165                    _msrest_next,166                    self.iter_content_func,167                )168                if not chunk:169                    raise _MsrestStopIteration()170                if self.user_callback and callable(self.user_callback):171                    self.user_callback(chunk, self.response)172                return chunk173            except _MsrestStopIteration:174                self.response.close()175                raise StopAsyncIteration()176            except Exception as err:177                _LOGGER.warning("Unable to stream download: %s", err)178                self.response.close()179                raise180 181    class TrioAsyncRequestsClientResponse(AsyncClientResponse, HTTPRequestsClientResponse):182 183        def stream_download(self, chunk_size: Optional[int] = None, callback: Optional[Callable] = None) -> AsyncIteratorType[bytes]:184            """Generator for streaming request body data.185 186            :param callback: Custom callback for monitoring progress.187            :param int chunk_size:188            """189            return TrioStreamDownloadGenerator(190                self.internal_response,191                callback,192                chunk_size193            )194 195 196    class AsyncTrioBasicRequestsHTTPSender(BasicRequestsHTTPSender, AsyncHTTPSender):  # type: ignore197 198        async def __aenter__(self):199            return super(AsyncTrioBasicRequestsHTTPSender, self).__enter__()200 201        async def __aexit__(self, *exc_details):  # pylint: disable=arguments-differ202            return super(AsyncTrioBasicRequestsHTTPSender, self).__exit__()203 204        async def send(self, request: ClientRequest, **kwargs: Any) -> AsyncClientResponse:  # type: ignore205            """Send the request using this HTTP sender.206            """207            # It's not recommended to provide its own session, and is mostly208            # to enable some legacy code to plug correctly209            session = kwargs.pop('session', self.session)210 211            trio_limiter = kwargs.get("trio_limiter", None)212            future = trio.to_thread.run_sync(213                functools.partial(214                    session.request,215                    request.method,216                    request.url,217                    **kwargs218                ),219                limiter=trio_limiter220            )221            try:222                return TrioAsyncRequestsClientResponse(223                    request,224                    await future225                )226            except requests.RequestException as err:227                msg = "Error occurred in request."228                raise_with_traceback(ClientRequestError, msg, err)229 230    class AsyncTrioRequestsHTTPSender(AsyncTrioBasicRequestsHTTPSender, RequestsHTTPSender):  # type: ignore231 232        async def send(self, request: ClientRequest, **kwargs: Any) -> AsyncClientResponse:  # type: ignore233            """Send the request using this HTTP sender.234            """235            requests_kwargs = self._configure_send(request, **kwargs)236            return await super(AsyncTrioRequestsHTTPSender, self).send(request, **requests_kwargs)237 238except ImportError:239    # trio not installed240    pass