Coverage for polar/oauth2/requests.py: 59%

1import typing 

2from collections import defaultdict 

3 

4from authlib.oauth2.rfc6749 import ( 

5 JsonPayload, 

6 JsonRequest, 

7 OAuth2Payload, 

8 OAuth2Request, 

9) 

10from starlette.datastructures import URL, ImmutableMultiDict, UploadFile 

11from starlette.requests import Request 

12 

13 

14class RequestPathParamsMixin: 

15 _request: Request 

16 

17 @property 

18 def path_params(self) -> dict[str, typing.Any]: 

19 return self._request.path_params 

20 

21 

22class StarletteOAuth2Payload(OAuth2Payload): 

23 def __init__(self, request: Request) -> None: 

24 # Merge query parameters and form data into a single dictionary 

25 datalist = defaultdict(list) 

26 sources: list[ImmutableMultiDict[str, str | UploadFile]] = [ 

27 request.query_params 

28 ] 

29 if request._form is not None: 

30 sources.append(request._form) 

31 for source in sources: 

32 for key, value in source.multi_items(): 

33 if not isinstance(value, UploadFile): 

34 datalist[key].append(value) 

35 self._datalist: dict[str, list[str]] = dict(datalist) 

36 self._data = {k: v[0] for k, v in self._datalist.items()} 

37 

38 @property 

39 def data(self) -> dict[str, str]: 

40 return self._data 

41 

42 @property 

43 def datalist(self) -> dict[str, list[str]]: 

44 return self._datalist 

45 

46 

47class StarletteOAuth2Request(RequestPathParamsMixin, OAuth2Request): 

48 def __init__(self, request: Request): 

49 super().__init__(request.method, str(request.url), headers=request.headers) 

50 self.user = getattr(request.state, "user", None) 

51 self.payload = StarletteOAuth2Payload(request) 

52 self._args = dict(request.query_params) 

53 self._form = dict(request._form) if request._form else {} 

54 

55 @property 

56 def args(self) -> dict[str, str | None]: 

57 return typing.cast(dict[str, str | None], self._args) 

58 

59 @property 

60 def form(self) -> dict[str, str]: 

61 return typing.cast(dict[str, str], self._form) 

62 

63 

64class StarletteJsonPayload(JsonPayload): 

65 def __init__(self, request: Request) -> None: 

66 self._data = getattr(request.state, "parsed_data", None) 

67 

68 @property 

69 def data(self) -> dict[str, str]: 

70 return self._data or {} 

71 

72 

73class StarletteJsonRequest(RequestPathParamsMixin, JsonRequest): 

74 credential: str | None = None 

75 

76 def __init__(self, request: Request): 

77 super().__init__(request.method, str(request.url), request.headers) 

78 self.user = getattr(request.state, "user", None) 

79 self.payload = StarletteJsonPayload(request) 

80 self._request = request 

81 

82 def url_for(self, name: str, /, **path_params: typing.Any) -> URL: 

83 return self._request.url_for(name, **path_params)