|
6 | 6 | import traceback |
7 | 7 | import io |
8 | 8 | from dataclasses import dataclass |
| 9 | +from email.message import Message |
| 10 | +from email.utils import collapse_rfc2231_value |
9 | 11 | from importlib.metadata import entry_points |
10 | 12 | from typing import Any, List, Dict, Optional, Union, BinaryIO |
11 | 13 | from pathlib import Path |
|
51 | 53 | ) |
52 | 54 |
|
53 | 55 |
|
| 56 | +def _get_content_disposition_filename(content_disposition: str) -> Optional[str]: |
| 57 | + message = Message() |
| 58 | + message["content-disposition"] = content_disposition |
| 59 | + |
| 60 | + fallback_filename: Optional[str] = None |
| 61 | + extended_filename: Optional[str] = None |
| 62 | + for key, value in message.get_params(header="content-disposition", unquote=True): |
| 63 | + if key != "filename": |
| 64 | + continue |
| 65 | + if isinstance(value, tuple): |
| 66 | + extended_filename = collapse_rfc2231_value(value) |
| 67 | + elif fallback_filename is None: |
| 68 | + fallback_filename = value |
| 69 | + |
| 70 | + return extended_filename or fallback_filename |
| 71 | + |
| 72 | + |
54 | 73 | # Lower priority values are tried first. |
55 | 74 | PRIORITY_SPECIFIC_FILE_FORMAT = ( |
56 | 75 | 0.0 # e.g., .docx, .pdf, .xlsx, Or specific pages, e.g., wikipedia |
@@ -512,9 +531,10 @@ def convert_response( |
512 | 531 | filename: Optional[str] = None |
513 | 532 | extension: Optional[str] = None |
514 | 533 | if "content-disposition" in response.headers: |
515 | | - m = re.search(r"filename=([^;]+)", response.headers["content-disposition"]) |
516 | | - if m: |
517 | | - filename = m.group(1).strip("\"'") |
| 534 | + filename = _get_content_disposition_filename( |
| 535 | + response.headers["content-disposition"] |
| 536 | + ) |
| 537 | + if filename is not None: |
518 | 538 | _, _extension = os.path.splitext(filename) |
519 | 539 | if len(_extension) > 0: |
520 | 540 | extension = _extension |
@@ -726,9 +746,9 @@ def _get_stream_info_guesses( |
726 | 746 | # If it's text, also guess the charset |
727 | 747 | charset = None |
728 | 748 | if result.prediction.output.is_text: |
729 | | - # Read the first 4k to guess the charset |
| 749 | + # Read the first 64k to guess the charset |
730 | 750 | file_stream.seek(cur_pos) |
731 | | - stream_page = file_stream.read(4096) |
| 751 | + stream_page = file_stream.read(65536) |
732 | 752 | charset_result = charset_normalizer.from_bytes(stream_page).best() |
733 | 753 |
|
734 | 754 | if charset_result is not None: |
|
0 commit comments