adapter.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168
  1. # SPDX-FileCopyrightText: 2015 Eric Larson
  2. #
  3. # SPDX-License-Identifier: Apache-2.0
  4. from __future__ import annotations
  5. import functools
  6. import types
  7. import weakref
  8. import zlib
  9. from typing import TYPE_CHECKING, Any, Collection, Mapping
  10. from pip._vendor.requests.adapters import HTTPAdapter
  11. from pip._vendor.cachecontrol.cache import DictCache
  12. from pip._vendor.cachecontrol.controller import PERMANENT_REDIRECT_STATUSES, CacheController
  13. from pip._vendor.cachecontrol.filewrapper import CallbackFileWrapper
  14. if TYPE_CHECKING:
  15. from pip._vendor.requests import PreparedRequest, Response
  16. from pip._vendor.urllib3 import HTTPResponse
  17. from pip._vendor.cachecontrol.cache import BaseCache
  18. from pip._vendor.cachecontrol.heuristics import BaseHeuristic
  19. from pip._vendor.cachecontrol.serialize import Serializer
  20. class CacheControlAdapter(HTTPAdapter):
  21. invalidating_methods = {"PUT", "PATCH", "DELETE"}
  22. def __init__(
  23. self,
  24. cache: BaseCache | None = None,
  25. cache_etags: bool = True,
  26. controller_class: type[CacheController] | None = None,
  27. serializer: Serializer | None = None,
  28. heuristic: BaseHeuristic | None = None,
  29. cacheable_methods: Collection[str] | None = None,
  30. *args: Any,
  31. **kw: Any,
  32. ) -> None:
  33. super().__init__(*args, **kw)
  34. self.cache = DictCache() if cache is None else cache
  35. self.heuristic = heuristic
  36. self.cacheable_methods = cacheable_methods or ("GET",)
  37. controller_factory = controller_class or CacheController
  38. self.controller = controller_factory(
  39. self.cache, cache_etags=cache_etags, serializer=serializer
  40. )
  41. def send(
  42. self,
  43. request: PreparedRequest,
  44. stream: bool = False,
  45. timeout: None | float | tuple[float, float] | tuple[float, None] = None,
  46. verify: bool | str = True,
  47. cert: (None | bytes | str | tuple[bytes | str, bytes | str]) = None,
  48. proxies: Mapping[str, str] | None = None,
  49. cacheable_methods: Collection[str] | None = None,
  50. ) -> Response:
  51. """
  52. Send a request. Use the request information to see if it
  53. exists in the cache and cache the response if we need to and can.
  54. """
  55. cacheable = cacheable_methods or self.cacheable_methods
  56. if request.method in cacheable:
  57. try:
  58. cached_response = self.controller.cached_request(request)
  59. except zlib.error:
  60. cached_response = None
  61. if cached_response:
  62. return self.build_response(request, cached_response, from_cache=True)
  63. # check for etags and add headers if appropriate
  64. request.headers.update(self.controller.conditional_headers(request))
  65. resp = super().send(request, stream, timeout, verify, cert, proxies)
  66. return resp
  67. def build_response( # type: ignore[override]
  68. self,
  69. request: PreparedRequest,
  70. response: HTTPResponse,
  71. from_cache: bool = False,
  72. cacheable_methods: Collection[str] | None = None,
  73. ) -> Response:
  74. """
  75. Build a response by making a request or using the cache.
  76. This will end up calling send and returning a potentially
  77. cached response
  78. """
  79. cacheable = cacheable_methods or self.cacheable_methods
  80. if not from_cache and request.method in cacheable:
  81. # Check for any heuristics that might update headers
  82. # before trying to cache.
  83. if self.heuristic:
  84. response = self.heuristic.apply(response)
  85. # apply any expiration heuristics
  86. if response.status == 304:
  87. # We must have sent an ETag request. This could mean
  88. # that we've been expired already or that we simply
  89. # have an etag. In either case, we want to try and
  90. # update the cache if that is the case.
  91. cached_response = self.controller.update_cached_response(
  92. request, response
  93. )
  94. if cached_response is not response:
  95. from_cache = True
  96. # We are done with the server response, read a
  97. # possible response body (compliant servers will
  98. # not return one, but we cannot be 100% sure) and
  99. # release the connection back to the pool.
  100. response.read(decode_content=False)
  101. response.release_conn()
  102. response = cached_response
  103. # We always cache the 301 responses
  104. elif int(response.status) in PERMANENT_REDIRECT_STATUSES:
  105. self.controller.cache_response(request, response)
  106. else:
  107. # Wrap the response file with a wrapper that will cache the
  108. # response when the stream has been consumed.
  109. response._fp = CallbackFileWrapper( # type: ignore[assignment]
  110. response._fp, # type: ignore[arg-type]
  111. functools.partial(
  112. self.controller.cache_response, request, weakref.ref(response)
  113. ),
  114. )
  115. if response.chunked:
  116. super_update_chunk_length = response.__class__._update_chunk_length
  117. def _update_chunk_length(
  118. weak_self: weakref.ReferenceType[HTTPResponse],
  119. ) -> None:
  120. self = weak_self()
  121. if self is None:
  122. return
  123. super_update_chunk_length(self)
  124. if self.chunk_left == 0:
  125. self._fp._close() # type: ignore[union-attr]
  126. response._update_chunk_length = functools.partial( # type: ignore[method-assign]
  127. _update_chunk_length, weakref.ref(response)
  128. )
  129. resp: Response = super().build_response(request, response)
  130. # See if we should invalidate the cache.
  131. if request.method in self.invalidating_methods and resp.ok:
  132. assert request.url is not None
  133. cache_url = self.controller.cache_url(request.url)
  134. self.cache.delete(cache_url)
  135. # Give the request a from_cache attr to let people use it
  136. resp.from_cache = from_cache # type: ignore[attr-defined]
  137. return resp
  138. def close(self) -> None:
  139. self.cache.close()
  140. super().close() # type: ignore[no-untyped-call]