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