| import logging |
| import os |
| import re |
| from typing import Optional |
| from plugins.common import DownloadCancelled |
| from plugins.direct_link import DirectLinkError, is_supported_direct_host, resolve_direct_link |
| from plugins.stream import download_stream_url |
|
|
| logger = logging.getLogger("hf_worker.plugins.direct") |
|
|
|
|
| def download_direct( |
| url: str, |
| output_dir: str, |
| custom_name: Optional[str] = None, |
| progress_fn=None, |
| cancel_event=None |
| ) -> str: |
| """Descarga cualquier enlace directo HTTP/HTTPS por streaming para el worker, resolviendo hosts soportados.""" |
| if cancel_event and cancel_event.is_set(): |
| raise DownloadCancelled() |
|
|
| download_url = url |
| headers = {} |
|
|
| if is_supported_direct_host(url): |
| logger.info(f"Resolviendo enlace multi-host: {url}") |
| try: |
| download_url, headers = resolve_direct_link(url) |
| logger.info(f"Enlace resuelto con éxito: {download_url[:80]}...") |
| except DirectLinkError as e: |
| logger.warning(f"Error resolviendo host ({e}), intentando con URL original...") |
| download_url = url |
|
|
| clean_name = None |
| if custom_name: |
| clean_name = str(custom_name).strip() |
| if clean_name.startswith(("http://", "https://", "ftp://")) or "://" in clean_name: |
| clean_name = None |
| else: |
| clean_name = os.path.basename(clean_name).strip() |
| clean_name = re.sub(r'[\/\\:\*\?"<>\|\x00-\x1f]', '_', clean_name).strip(' ._') |
| if not clean_name or clean_name in ("downloaded_file", "archivo"): |
| clean_name = None |
|
|
| if clean_name: |
| out_target = os.path.join(output_dir, clean_name) |
| else: |
| out_target = output_dir |
|
|
| dl_file, _ = download_stream_url( |
| download_url, |
| out_target, |
| progress_fn=progress_fn, |
| cancel_event=cancel_event, |
| headers=headers if headers else None, |
| ) |
| return dl_file |
|
|
