adapter.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133
  1. import types
  2. import functools
  3. import zlib
  4. from pip._vendor.requests.adapters import HTTPAdapter
  5. from .controller import CacheController
  6. from .cache import DictCache
  7. from .filewrapper import CallbackFileWrapper
  8. class CacheControlAdapter(HTTPAdapter):
  9. invalidating_methods = {"PUT", "DELETE"}
  10. def __init__(
  11. self,
  12. cache=None,
  13. cache_etags=True,
  14. controller_class=None,
  15. serializer=None,
  16. heuristic=None,
  17. cacheable_methods=None,
  18. *args,
  19. **kw
  20. ):
  21. super(CacheControlAdapter, self).__init__(*args, **kw)
  22. self.cache = DictCache() if cache is None else cache
  23. self.heuristic = heuristic
  24. self.cacheable_methods = cacheable_methods or ("GET",)
  25. controller_factory = controller_class or CacheController
  26. self.controller = controller_factory(
  27. self.cache, cache_etags=cache_etags, serializer=serializer
  28. )
  29. def send(self, request, cacheable_methods=None, **kw):
  30. """
  31. Send a request. Use the request information to see if it
  32. exists in the cache and cache the response if we need to and can.
  33. """
  34. cacheable = cacheable_methods or self.cacheable_methods
  35. if request.method in cacheable:
  36. try:
  37. cached_response = self.controller.cached_request(request)
  38. except zlib.error:
  39. cached_response = None
  40. if cached_response:
  41. return self.build_response(request, cached_response, from_cache=True)
  42. # check for etags and add headers if appropriate
  43. request.headers.update(self.controller.conditional_headers(request))
  44. resp = super(CacheControlAdapter, self).send(request, **kw)
  45. return resp
  46. def build_response(
  47. self, request, response, from_cache=False, cacheable_methods=None
  48. ):
  49. """
  50. Build a response by making a request or using the cache.
  51. This will end up calling send and returning a potentially
  52. cached response
  53. """
  54. cacheable = cacheable_methods or self.cacheable_methods
  55. if not from_cache and request.method in cacheable:
  56. # Check for any heuristics that might update headers
  57. # before trying to cache.
  58. if self.heuristic:
  59. response = self.heuristic.apply(response)
  60. # apply any expiration heuristics
  61. if response.status == 304:
  62. # We must have sent an ETag request. This could mean
  63. # that we've been expired already or that we simply
  64. # have an etag. In either case, we want to try and
  65. # update the cache if that is the case.
  66. cached_response = self.controller.update_cached_response(
  67. request, response
  68. )
  69. if cached_response is not response:
  70. from_cache = True
  71. # We are done with the server response, read a
  72. # possible response body (compliant servers will
  73. # not return one, but we cannot be 100% sure) and
  74. # release the connection back to the pool.
  75. response.read(decode_content=False)
  76. response.release_conn()
  77. response = cached_response
  78. # We always cache the 301 responses
  79. elif response.status == 301:
  80. self.controller.cache_response(request, response)
  81. else:
  82. # Wrap the response file with a wrapper that will cache the
  83. # response when the stream has been consumed.
  84. response._fp = CallbackFileWrapper(
  85. response._fp,
  86. functools.partial(
  87. self.controller.cache_response, request, response
  88. ),
  89. )
  90. if response.chunked:
  91. super_update_chunk_length = response._update_chunk_length
  92. def _update_chunk_length(self):
  93. super_update_chunk_length()
  94. if self.chunk_left == 0:
  95. self._fp._close()
  96. response._update_chunk_length = types.MethodType(
  97. _update_chunk_length, response
  98. )
  99. resp = super(CacheControlAdapter, self).build_response(request, response)
  100. # See if we should invalidate the cache.
  101. if request.method in self.invalidating_methods and resp.ok:
  102. cache_url = self.controller.cache_url(request.url)
  103. self.cache.delete(cache_url)
  104. # Give the request a from_cache attr to let people use it
  105. resp.from_cache = from_cache
  106. return resp
  107. def close(self):
  108. self.cache.close()
  109. super(CacheControlAdapter, self).close()