|
7 | 7 | import httpx |
8 | 8 |
|
9 | 9 |
|
10 | | -def upload_file(fh: io.IOBase, output_file_prefix: Optional[str] = None) -> str: |
| 10 | +def upload_file(file: io.IOBase, output_file_prefix: Optional[str] = None) -> str: |
11 | 11 | """ |
12 | 12 | Upload a file to the server. |
13 | 13 |
|
14 | 14 | Args: |
15 | | - fh: A file handle to upload. |
| 15 | + file: A file handle to upload. |
16 | 16 | output_file_prefix: A string to prepend to the output file name. |
17 | 17 | Returns: |
18 | 18 | str: A URL to the uploaded file. |
19 | 19 | """ |
20 | 20 | # Lifted straight from cog.files |
21 | 21 |
|
22 | | - fh.seek(0) |
| 22 | + file.seek(0) |
23 | 23 |
|
24 | 24 | if output_file_prefix is not None: |
25 | | - name = getattr(fh, "name", "output") |
| 25 | + name = getattr(file, "name", "output") |
26 | 26 | url = output_file_prefix + os.path.basename(name) |
27 | | - resp = httpx.put(url, files={"file": fh}, timeout=None) # type: ignore |
| 27 | + resp = httpx.put(url, files={"file": file}, timeout=None) # type: ignore |
28 | 28 | resp.raise_for_status() |
| 29 | + |
29 | 30 | return url |
30 | 31 |
|
31 | | - b = fh.read() |
32 | | - # The file handle is strings, not bytes |
33 | | - if isinstance(b, str): |
34 | | - b = b.encode("utf-8") |
35 | | - encoded_body = base64.b64encode(b) |
36 | | - if getattr(fh, "name", None): |
37 | | - # despite doing a getattr check here, mypy complains that io.IOBase has no attribute name |
38 | | - mime_type = mimetypes.guess_type(fh.name)[0] # type: ignore |
39 | | - else: |
40 | | - mime_type = "application/octet-stream" |
41 | | - s = encoded_body.decode("utf-8") |
42 | | - return f"data:{mime_type};base64,{s}" |
| 32 | + body = file.read() |
| 33 | + # Ensure the file handle is in bytes |
| 34 | + body = body.encode("utf-8") if isinstance(body, str) else body |
| 35 | + encoded_body = base64.b64encode(body).decode("utf-8") |
| 36 | + # Use getattr to avoid mypy complaints about io.IOBase having no attribute name |
| 37 | + mime_type = ( |
| 38 | + mimetypes.guess_type(getattr(file, "name", ""))[0] or "application/octet-stream" |
| 39 | + ) |
| 40 | + return f"data:{mime_type};base64,{encoded_body}" |
0 commit comments