adapter.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134
  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 = set(['PUT', 'DELETE'])
  10. def __init__(self, cache=None,
  11. cache_etags=True,
  12. controller_class=None,
  13. serializer=None,
  14. heuristic=None,
  15. cacheable_methods=None,
  16. *args, **kw):
  17. super(CacheControlAdapter, self).__init__(*args, **kw)
  18. self.cache = cache or DictCache()
  19. self.heuristic = heuristic
  20. self.cacheable_methods = cacheable_methods or ('GET',)
  21. controller_factory = controller_class or CacheController
  22. self.controller = controller_factory(
  23. self.cache,
  24. cache_etags=cache_etags,
  25. serializer=serializer,
  26. )
  27. def send(self, request, cacheable_methods=None, **kw):
  28. """
  29. Send a request. Use the request information to see if it
  30. exists in the cache and cache the response if we need to and can.
  31. """
  32. cacheable = cacheable_methods or self.cacheable_methods
  33. if request.method in cacheable:
  34. try:
  35. cached_response = self.controller.cached_request(request)
  36. except zlib.error:
  37. cached_response = None
  38. if cached_response:
  39. return self.build_response(request, cached_response,
  40. from_cache=True)
  41. # check for etags and add headers if appropriate
  42. request.headers.update(
  43. self.controller.conditional_headers(request)
  44. )
  45. resp = super(CacheControlAdapter, self).send(request, **kw)
  46. return resp
  47. def build_response(self, request, response, from_cache=False,
  48. cacheable_methods=None):
  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,
  88. request,
  89. response,
  90. )
  91. )
  92. if response.chunked:
  93. super_update_chunk_length = response._update_chunk_length
  94. def _update_chunk_length(self):
  95. super_update_chunk_length()
  96. if self.chunk_left == 0:
  97. self._fp._close()
  98. response._update_chunk_length = types.MethodType(_update_chunk_length, response)
  99. resp = super(CacheControlAdapter, self).build_response(
  100. request, response
  101. )
  102. # See if we should invalidate the cache.
  103. if request.method in self.invalidating_methods and resp.ok:
  104. cache_url = self.controller.cache_url(request.url)
  105. self.cache.delete(cache_url)
  106. # Give the request a from_cache attr to let people use it
  107. resp.from_cache = from_cache
  108. return resp
  109. def close(self):
  110. self.cache.close()
  111. super(CacheControlAdapter, self).close()