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