adapter.py 4.6 KB

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