diff --git a/scripts/twitter_threads/download_media.py b/scripts/twitter_threads/download_media.py new file mode 100644 index 00000000..3a042951 --- /dev/null +++ b/scripts/twitter_threads/download_media.py @@ -0,0 +1,56 @@ +"""Download media URLs from ThreadData via stdlib urllib into a local directory.""" +from __future__ import annotations +from pathlib import Path +from urllib.request import urlopen +from urllib.error import URLError +from urllib.parse import urlparse, parse_qs +from scripts.twitter_threads.error_types import ErrorInfo, make_error, Result, PostData, ThreadData + +def _classify(url: str) -> tuple[str, str]: + parsed = urlparse(url) + path = parsed.path + if "video.twimg.com" in parsed.netloc or path.endswith(".mp4"): + if "tweet_video" in path: + return ("gif", "mp4") + return ("vid", "mp4") + qs = parse_qs(parsed.query) + fmt = qs.get("format", [""])[0] + if fmt: + return ("img", fmt) + if "." in path: + return ("img", path.rsplit(".", 1)[-1]) + return ("img", "jpg") + +def media_names_for_post(post: PostData) -> list[str]: + counters: dict[str, int] = {} + names: list[str] = [] + for url in post.media_urls: + kind, ext = _classify(url) + counters[kind] = counters.get(kind, 0) + 1 + names.append(f"{post.post_id}_{kind}{counters[kind]}.{ext}") + return names + +def download_media(thread: ThreadData, output_dir: Path) -> Result[list[Path]]: + try: + output_dir.mkdir(parents=True, exist_ok=True) + except OSError as e: + return Result.err(make_error("WriteError", "download_media", str(e))) + paths: list[Path] = [] + for post in thread.posts: + names = media_names_for_post(post) + for url, name in zip(post.media_urls, names): + target = output_dir / name + if target.exists() and target.stat().st_size > 0: + paths.append(target) + continue + try: + with urlopen(url) as resp: + data = resp.read() + except URLError as e: + return Result.err(make_error("HttpError", "download_media", str(e))) + try: + target.write_bytes(data) + except OSError as e: + return Result.err(make_error("WriteError", "download_media", str(e))) + paths.append(target) + return Result.ok(paths) diff --git a/tests/test_twitter_threads_media.py b/tests/test_twitter_threads_media.py new file mode 100644 index 00000000..e5f4401a --- /dev/null +++ b/tests/test_twitter_threads_media.py @@ -0,0 +1,82 @@ +"""Tests for scripts.twitter_threads.download_media: naming and HTTP download contract.""" +from __future__ import annotations + +from pathlib import Path +from unittest.mock import patch, MagicMock +from urllib.error import URLError + +import pytest + +from scripts.twitter_threads.error_types import PostMetrics, PostData, ThreadData +from scripts.twitter_threads.download_media import media_names_for_post, download_media + + +class _FakeResp: + def __init__(self, data: bytes) -> None: + self._data = data + def read(self) -> bytes: + return self._data + def __enter__(self) -> "_FakeResp": + return self + def __exit__(self, *a: object) -> bool: + return False + + +def _post(post_id: str, media_urls: tuple[str, ...]) -> PostData: + metrics = PostMetrics(reply_count=0, repost_count=0, like_count=0, view_count=None) + return PostData(post_id=post_id, author="Tim", handle="NOTimothyLottes", text="t", timestamp="2024-02-01T00:00:00Z", media_urls=media_urls, reply_to_id=None, quote_of_id=None, metrics=metrics) + + +def test_media_names_for_post() -> None: + post = _post("123", ("https://pbs.twimg.com/media/A?format=jpg&name=large", "https://pbs.twimg.com/media/B?format=png&name=large")) + assert media_names_for_post(post) == ["123_img1.jpg", "123_img2.png"] + + +def test_media_names_video() -> None: + post = _post("123", ("https://video.twimg.com/ext_tw_video/1/pu/vid/1280x720/1234567890.mp4",)) + assert media_names_for_post(post) == ["123_vid1.mp4"] + + +def test_download_naming(tmp_path: Path) -> None: + post = _post("123", ("https://pbs.twimg.com/media/A?format=jpg&name=large", "https://pbs.twimg.com/media/B?format=png&name=large")) + thread = ThreadData(root_post_id="123", posts=(post,), source_url="https://x.com/x/x/123") + with patch("scripts.twitter_threads.download_media.urlopen", return_value=_FakeResp(b"bytes")): + result = download_media(thread, tmp_path) + assert result.is_ok + assert (tmp_path / "123_img1.jpg").exists() + assert (tmp_path / "123_img2.png").exists() + assert result.data is not None + names = [p.name for p in result.data] + assert "123_img1.jpg" in names + assert "123_img2.png" in names + + +def test_download_video_naming(tmp_path: Path) -> None: + post = _post("123", ("https://video.twimg.com/ext_tw_video/1/pu/vid/1280x720/1234567890.mp4",)) + thread = ThreadData(root_post_id="123", posts=(post,), source_url="https://x.com/x/x/123") + with patch("scripts.twitter_threads.download_media.urlopen", return_value=_FakeResp(b"vid")): + result = download_media(thread, tmp_path) + assert (tmp_path / "123_vid1.mp4").exists() + + +def test_download_idempotent(tmp_path: Path) -> None: + post = _post("123", ("https://pbs.twimg.com/media/A?format=jpg&name=large",)) + thread = ThreadData(root_post_id="123", posts=(post,), source_url="https://x.com/x/x/123") + name = media_names_for_post(post)[0] + (tmp_path / name).write_bytes(b"already") + mock_urlopen = MagicMock() + with patch("scripts.twitter_threads.download_media.urlopen", mock_urlopen): + result = download_media(thread, tmp_path) + assert result.is_ok + mock_urlopen.assert_not_called() + assert (tmp_path / name).read_bytes() == b"already" + + +def test_download_http_error(tmp_path: Path) -> None: + post = _post("123", ("https://pbs.twimg.com/media/A?format=jpg&name=large",)) + thread = ThreadData(root_post_id="123", posts=(post,), source_url="https://x.com/x/x/123") + with patch("scripts.twitter_threads.download_media.urlopen", side_effect=URLError("boom")): + result = download_media(thread, tmp_path) + assert result.is_ok is False + assert result.error is not None + assert result.error.kind == "HttpError" \ No newline at end of file