|
15 | 15 | import os |
16 | 16 | import sys |
17 | 17 | from pathlib import Path |
18 | | -from typing import TYPE_CHECKING, List, Optional, Union |
| 18 | +from typing import TYPE_CHECKING, Dict, List, Optional, Union, cast |
19 | 19 |
|
20 | 20 | if sys.version_info >= (3, 8): # pragma: no cover |
21 | 21 | from typing import TypedDict |
22 | 22 | else: # pragma: no cover |
23 | 23 | from typing_extensions import TypedDict |
24 | 24 |
|
25 | 25 | from playwright._impl._connection import Channel, from_channel |
26 | | -from playwright._impl._helper import Error, async_readfile |
| 26 | +from playwright._impl._helper import Error |
27 | 27 | from playwright._impl._writable_stream import WritableStream |
28 | 28 |
|
29 | 29 | if TYPE_CHECKING: # pragma: no cover |
|
34 | 34 | SIZE_LIMIT_IN_BYTES = 50 * 1024 * 1024 |
35 | 35 |
|
36 | 36 |
|
37 | | -class InputFilesList(TypedDict): |
| 37 | +class InputFilesList(TypedDict, total=False): |
38 | 38 | streams: Optional[List[Channel]] |
39 | 39 | localPaths: Optional[List[str]] |
40 | | - files: Optional[List[FilePayload]] |
| 40 | + payloads: Optional[List[Dict[str, Union[str, bytes]]]] |
41 | 41 |
|
42 | 42 |
|
43 | 43 | async def convert_input_files( |
44 | 44 | files: Union[str, Path, FilePayload, List[Union[str, Path]], List[FilePayload]], |
45 | 45 | context: "BrowserContext", |
46 | 46 | ) -> InputFilesList: |
47 | | - file_list = files if isinstance(files, list) else [files] |
| 47 | + items = files if isinstance(files, list) else [files] |
48 | 48 |
|
49 | | - total_buffer_size_exceeds_limit = ( |
50 | | - sum( |
51 | | - [ |
52 | | - len(f.get("buffer", "")) |
53 | | - for f in file_list |
54 | | - if not isinstance(f, (str, Path)) |
55 | | - ] |
56 | | - ) |
57 | | - > SIZE_LIMIT_IN_BYTES |
58 | | - ) |
59 | | - if total_buffer_size_exceeds_limit: |
60 | | - raise Error( |
61 | | - "Cannot set buffer larger than 50Mb, please write it to a file and pass its path instead." |
62 | | - ) |
63 | | - |
64 | | - total_file_size_exceeds_limit = ( |
65 | | - sum([os.stat(f).st_size for f in file_list if isinstance(f, (str, Path))]) |
66 | | - > SIZE_LIMIT_IN_BYTES |
67 | | - ) |
68 | | - if total_file_size_exceeds_limit: |
| 49 | + if any([isinstance(item, (str, Path)) for item in items]): |
| 50 | + if not all([isinstance(item, (str, Path)) for item in items]): |
| 51 | + raise Error("File paths cannot be mixed with buffers") |
69 | 52 | if context._channel._connection.is_remote: |
70 | 53 | streams = [] |
71 | | - for file in file_list: |
72 | | - assert isinstance(file, (str, Path)) |
| 54 | + for item in items: |
| 55 | + assert isinstance(item, (str, Path)) |
| 56 | + last_modified_ms = int(os.path.getmtime(item) * 1000) |
73 | 57 | stream: WritableStream = from_channel( |
74 | 58 | await context._channel.send( |
75 | | - "createTempFile", {"name": os.path.basename(file)} |
| 59 | + "createTempFile", |
| 60 | + { |
| 61 | + "name": os.path.basename(item), |
| 62 | + "lastModifiedMs": last_modified_ms, |
| 63 | + }, |
76 | 64 | ) |
77 | 65 | ) |
78 | | - await stream.copy(file) |
| 66 | + await stream.copy(item) |
79 | 67 | streams.append(stream._channel) |
80 | | - return InputFilesList(streams=streams, localPaths=None, files=None) |
81 | | - local_paths = [] |
82 | | - for p in file_list: |
83 | | - assert isinstance(p, (str, Path)) |
84 | | - local_paths.append(str(Path(p).absolute().resolve())) |
85 | | - return InputFilesList(streams=None, localPaths=local_paths, files=None) |
| 68 | + return InputFilesList(streams=streams) |
| 69 | + return InputFilesList( |
| 70 | + localPaths=[ |
| 71 | + str(Path(cast(Union[str, Path], item)).absolute().resolve()) |
| 72 | + for item in items |
| 73 | + ] |
| 74 | + ) |
86 | 75 |
|
87 | | - return InputFilesList( |
88 | | - streams=None, localPaths=None, files=await _normalize_file_payloads(files) |
| 76 | + file_payload_exceeds_size_limit = ( |
| 77 | + sum([len(f.get("buffer", "")) for f in items if not isinstance(f, (str, Path))]) |
| 78 | + > SIZE_LIMIT_IN_BYTES |
89 | 79 | ) |
| 80 | + if file_payload_exceeds_size_limit: |
| 81 | + raise Error( |
| 82 | + "Cannot set buffer larger than 50Mb, please write it to a file and pass its path instead." |
| 83 | + ) |
90 | 84 |
|
91 | | - |
92 | | -async def _normalize_file_payloads( |
93 | | - files: Union[str, Path, FilePayload, List[Union[str, Path]], List[FilePayload]] |
94 | | -) -> List: |
95 | | - file_list = files if isinstance(files, list) else [files] |
96 | | - file_payloads: List = [] |
97 | | - for item in file_list: |
98 | | - if isinstance(item, (str, Path)): |
99 | | - file_payloads.append( |
100 | | - { |
101 | | - "name": os.path.basename(item), |
102 | | - "buffer": base64.b64encode(await async_readfile(item)).decode(), |
103 | | - } |
104 | | - ) |
105 | | - else: |
106 | | - file_payloads.append( |
107 | | - { |
108 | | - "name": item["name"], |
109 | | - "mimeType": item["mimeType"], |
110 | | - "buffer": base64.b64encode(item["buffer"]).decode(), |
111 | | - } |
112 | | - ) |
113 | | - |
114 | | - return file_payloads |
| 85 | + return InputFilesList( |
| 86 | + payloads=[ |
| 87 | + { |
| 88 | + "name": item["name"], |
| 89 | + "mimeType": item["mimeType"], |
| 90 | + "buffer": base64.b64encode(item["buffer"]).decode(), |
| 91 | + } |
| 92 | + for item in cast(List[FilePayload], items) |
| 93 | + ] |
| 94 | + ) |
0 commit comments