Coverage for src/prisma/_async_http.py: 78%

1import json 

2from typing import Any 

3from typing_extensions import override 

4 

5import httpx 

6 

7from ._types import Method 

8from .http_abstract import AbstractHTTP, AbstractResponse 

9 

10__all__ = ('HTTP', 'AsyncHTTP', 'Response', 'client') 

11 

12 

13class AsyncHTTP(AbstractHTTP[httpx.AsyncClient, httpx.Response]): 

14 session: httpx.AsyncClient 

15 

16 @override 

17 async def download(self, url: str, dest: str) -> None: 

18 async with self.session.stream('GET', url, timeout=None) as resp: 

19 resp.raise_for_status() 

20 with open(dest, 'wb') as fd: 

21 async for chunk in resp.aiter_bytes(): 

22 fd.write(chunk) 

23 

24 @override 

25 async def request(self, method: Method, url: str, **kwargs: Any) -> 'Response': 

26 return Response(await self.session.request(method, url, **kwargs)) 

27 

28 @override 

29 def open(self) -> None: 

30 self.session = httpx.AsyncClient(**self.session_kwargs) 

31 

32 @override 

33 async def close(self) -> None: 

34 if self.should_close(): 

35 await self.session.aclose() 

36 

37 # mypy doesn't like us assigning None as the type of 

38 # session is not optional, however the argument that 

39 # the setter takes is optional, so this is fine 

40 self.session = None # type: ignore[assignment] 

41 

42 

43HTTP = AsyncHTTP 

44 

45 

46client: HTTP = HTTP() 

47 

48 

49class Response(AbstractResponse[httpx.Response]): 

50 __slots__ = () 

51 

52 @property 

53 @override 

54 def status(self) -> int: 

55 return self.original.status_code 

56 

57 @property 

58 @override 

59 def headers(self) -> httpx.Headers: 

60 return self.original.headers 

61 

62 @override 

63 async def json(self, **kwargs: Any) -> Any: 

64 return json.loads(await self.original.aread(), **kwargs) 

65 

66 @override 

67 async def text(self, **kwargs: Any) -> str: 

68 return ''.join([part async for part in self.original.aiter_text(**kwargs)])