16
16
import base64
17
17
import json
18
18
import mimetypes
19
+ from collections import defaultdict
19
20
from pathlib import Path
20
21
from types import SimpleNamespace
21
22
from typing import TYPE_CHECKING , Any , Callable , Dict , List , Optional , Union , cast
22
23
from urllib import parse
23
24
24
25
from playwright ._impl ._api_structures import (
26
+ Headers ,
27
+ HeadersArray ,
25
28
RemoteAddr ,
26
29
RequestSizes ,
27
30
ResourceTiming ,
34
37
from_nullable_channel ,
35
38
)
36
39
from playwright ._impl ._event_context_manager import EventContextManagerImpl
37
- from playwright ._impl ._helper import ContinueParameters , Header , locals_to_params
40
+ from playwright ._impl ._helper import ContinueParameters , locals_to_params
38
41
from playwright ._impl ._wait_helper import WaitHelper
39
42
40
43
if TYPE_CHECKING : # pragma: no cover
@@ -64,8 +67,8 @@ def __init__(
64
67
"responseStart" : - 1 ,
65
68
"responseEnd" : - 1 ,
66
69
}
67
- self ._headers : List [ Header ] = self ._initializer ["headers" ]
68
- self ._all_headers_future : Optional [asyncio .Future [List [ Header ] ]] = None
70
+ self ._provisional_headers = RawHeaders ( self ._initializer ["headers" ])
71
+ self ._all_headers_future : Optional [asyncio .Future [RawHeaders ]] = None
69
72
70
73
def __repr__ (self ) -> str :
71
74
return f"<Request url={ self .url !r} method={ self .method !r} >"
@@ -115,10 +118,6 @@ def post_data_buffer(self) -> Optional[bytes]:
115
118
return None
116
119
return base64 .b64decode (b64_content )
117
120
118
- @property
119
- def headers (self ) -> Dict [str , str ]:
120
- return headers_array_to_object (self ._headers , True )
121
-
122
121
async def response (self ) -> Optional ["Response" ]:
123
122
return from_nullable_channel (await self ._channel .send ("response" ))
124
123
@@ -145,25 +144,27 @@ def failure(self) -> Optional[str]:
145
144
def timing (self ) -> ResourceTiming :
146
145
return self ._timing
147
146
148
- async def all_headers (self ) -> Dict [str , str ]:
149
- return headers_array_to_object (await self ._get_headers_if_needed (), True )
147
+ @property
148
+ def headers (self ) -> Headers :
149
+ return self ._provisional_headers .headers ()
150
150
151
- async def headers_array (self ) -> List [List [str ]]:
152
- return list (
153
- map (
154
- lambda header : [header ["name" ], header ["value" ]],
155
- await self ._get_headers_if_needed (),
156
- )
157
- )
151
+ async def all_headers (self ) -> Headers :
152
+ return (await self ._actual_headers ()).headers ()
153
+
154
+ async def headers_array (self ) -> HeadersArray :
155
+ return (await self ._actual_headers ()).headers_array ()
156
+
157
+ async def header_value (self , name : str ) -> Optional [str ]:
158
+ return (await self ._actual_headers ()).get (name )
158
159
159
- async def _get_headers_if_needed (self ) -> List [ Header ] :
160
+ async def _actual_headers (self ) -> "RawHeaders" :
160
161
if not self ._all_headers_future :
161
162
self ._all_headers_future = asyncio .Future ()
162
163
response = await self .response ()
163
164
if not response :
164
- return self ._headers
165
+ return self ._provisional_headers
165
166
headers = await response ._channel .send ("rawRequestHeaders" )
166
- self ._all_headers_future .set_result (headers )
167
+ self ._all_headers_future .set_result (RawHeaders ( headers ) )
167
168
return await self ._all_headers_future
168
169
169
170
@@ -256,10 +257,10 @@ def __init__(
256
257
self ._request ._timing ["connectEnd" ] = timing ["connectEnd" ]
257
258
self ._request ._timing ["requestStart" ] = timing ["requestStart" ]
258
259
self ._request ._timing ["responseStart" ] = timing ["responseStart" ]
259
- self ._headers = headers_array_to_object (
260
- cast (List [ Header ] , self ._initializer ["headers" ]), True
260
+ self ._provisional_headers = RawHeaders (
261
+ cast (HeadersArray , self ._initializer ["headers" ])
261
262
)
262
- self ._raw_headers_future : Optional [asyncio .Future [List [ Header ] ]] = None
263
+ self ._raw_headers_future : Optional [asyncio .Future [RawHeaders ]] = None
263
264
self ._finished_future : asyncio .Future [bool ] = asyncio .Future ()
264
265
265
266
def __repr__ (self ) -> str :
@@ -284,25 +285,26 @@ def status_text(self) -> str:
284
285
return self ._initializer ["statusText" ]
285
286
286
287
@property
287
- def headers (self ) -> Dict [ str , str ] :
288
- return self ._headers . copy ()
288
+ def headers (self ) -> Headers :
289
+ return self ._provisional_headers . headers ()
289
290
290
- async def all_headers (self ) -> Dict [ str , str ] :
291
- return headers_array_to_object (await self ._get_headers_if_needed (), True )
291
+ async def all_headers (self ) -> Headers :
292
+ return (await self ._actual_headers ()). headers ( )
292
293
293
- async def headers_array (self ) -> List [List [str ]]:
294
- return list (
295
- map (
296
- lambda header : [header ["name" ], header ["value" ]],
297
- await self ._get_headers_if_needed (),
298
- )
299
- )
294
+ async def headers_array (self ) -> HeadersArray :
295
+ return (await self ._actual_headers ()).headers_array ()
296
+
297
+ async def header_value (self , name : str ) -> Optional [str ]:
298
+ return (await self ._actual_headers ()).get (name )
299
+
300
+ async def header_values (self , name : str ) -> List [str ]:
301
+ return (await self ._actual_headers ()).get_all (name )
300
302
301
- async def _get_headers_if_needed (self ) -> List [ Header ] :
303
+ async def _actual_headers (self ) -> "RawHeaders" :
302
304
if not self ._raw_headers_future :
303
305
self ._raw_headers_future = asyncio .Future ()
304
- headers = cast (List [ Header ] , await self ._channel .send ("rawResponseHeaders" ))
305
- self ._raw_headers_future .set_result (headers )
306
+ headers = cast (HeadersArray , await self ._channel .send ("rawResponseHeaders" ))
307
+ self ._raw_headers_future .set_result (RawHeaders ( headers ) )
306
308
return await self ._raw_headers_future
307
309
308
310
async def server_addr (self ) -> Optional [RemoteAddr ]:
@@ -420,12 +422,32 @@ def _on_close(self) -> None:
420
422
self .emit (WebSocket .Events .Close , self )
421
423
422
424
423
- def serialize_headers (headers : Dict [str , str ]) -> List [ Header ] :
425
+ def serialize_headers (headers : Dict [str , str ]) -> HeadersArray :
424
426
return [{"name" : name , "value" : value } for name , value in headers .items ()]
425
427
426
428
427
- def headers_array_to_object (headers : List [Header ], lower_case : bool ) -> Dict [str , str ]:
428
- return {
429
- (header ["name" ].lower () if lower_case else header ["name" ]): header ["value" ]
430
- for header in headers
431
- }
429
+ class RawHeaders :
430
+ def __init__ (self , headers : HeadersArray ) -> None :
431
+ self ._headers_array = headers
432
+ self ._headers_map : Dict [str , Dict [str , bool ]] = defaultdict (dict )
433
+ for header in headers :
434
+ self ._headers_map [header ["name" ].lower ()][header ["value" ]] = True
435
+
436
+ def get (self , name : str ) -> Optional [str ]:
437
+ values = self .get_all (name )
438
+ if not values :
439
+ return None
440
+ separator = "\n " if name .lower () == "set-cookie" else ", "
441
+ return separator .join (values )
442
+
443
+ def get_all (self , name : str ) -> List [str ]:
444
+ return list (self ._headers_map [name .lower ()].keys ())
445
+
446
+ def headers (self ) -> Dict [str , str ]:
447
+ result = {}
448
+ for name in self ._headers_map .keys ():
449
+ result [name ] = cast (str , self .get (name ))
450
+ return result
451
+
452
+ def headers_array (self ) -> HeadersArray :
453
+ return self ._headers_array
0 commit comments