diff --git a/conductor/tracks/twitter_threads_extraction_20260705/extract_corpus.py b/conductor/tracks/twitter_threads_extraction_20260705/extract_corpus.py index 8bbea57e..5dd70be5 100644 --- a/conductor/tracks/twitter_threads_extraction_20260705/extract_corpus.py +++ b/conductor/tracks/twitter_threads_extraction_20260705/extract_corpus.py @@ -18,7 +18,6 @@ from __future__ import annotations import json import shutil -import subprocess import sys from pathlib import Path @@ -26,6 +25,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[3])) from scripts.twitter_threads.error_types import ThreadData from scripts.twitter_threads.fetch_thread import fetch_thread_from_url, thread_to_dict +from scripts.twitter_threads.download_media import download_media, media_files_for_post from scripts.twitter_threads.render_markdown import render_markdown TRACK = Path(__file__).resolve().parent @@ -62,12 +62,6 @@ def ensure_cookies() -> None: COOKIES_NS.write_text("\n".join(out) + "\n", encoding="utf-8") -def gdl_download(url: str, media_dir: Path) -> None: - media_dir.mkdir(parents=True, exist_ok=True) - subprocess.run( - ["gallery-dl", "--cookies", str(COOKIES_NS), "-o", "text-tweets=true", "-o", "conversations=true", "-D", str(media_dir), url], - capture_output=True, text=True, - ) def groups(items: dict[str, dict]) -> list[list[str]]: @@ -116,9 +110,10 @@ def main() -> int: td = ThreadData(root_post_id=primary, posts=tuple(union), source_url=items[primary]["url"]) base = CORPUS / primary media_dir = base / "media" + media_dir.mkdir(parents=True, exist_ok=True) for m in members: - gdl_download(items[m]["url"], media_dir) - media_names = {p.post_id: sorted(f.name for f in media_dir.glob(f"{p.post_id}_*")) for p in td.posts} + download_media(items[m]["td"], media_dir, COOKIES_NS) + media_names = {p.post_id: media_files_for_post(media_dir, p.post_id) for p in td.posts} media_names = {k: v for k, v in media_names.items() if v} n_media = sum(len(v) for v in media_names.values()) data = thread_to_dict(td) diff --git a/scripts/twitter_threads/download_media.py b/scripts/twitter_threads/download_media.py index 0754e849..da8ad40f 100644 --- a/scripts/twitter_threads/download_media.py +++ b/scripts/twitter_threads/download_media.py @@ -1,70 +1,44 @@ -"""Download media URLs from ThreadData via stdlib urllib into a local directory.""" +"""Download a Twitter/X thread's media via gallery-dl (handles auth, 403, and video/GIF -> mp4).""" from __future__ import annotations import argparse import json +import subprocess import sys from pathlib import Path -from urllib.request import urlopen -from urllib.error import URLError -from urllib.parse import urlparse, parse_qs try: - from scripts.twitter_threads.error_types import ErrorInfo, make_error, Result, PostData, ThreadData, thread_from_dict + from scripts.twitter_threads.error_types import ErrorInfo, make_error, Result, ThreadData, thread_from_dict except ImportError: - from error_types import ErrorInfo, make_error, Result, PostData, ThreadData, thread_from_dict + from error_types import ErrorInfo, make_error, Result, ThreadData, thread_from_dict -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]]: +def download_media(thread: ThreadData, output_dir: Path, cookies_path: Path | None = None) -> 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) + args = ["gallery-dl", "-o", "text-tweets=true", "-o", "conversations=true", "-D", str(output_dir)] + if cookies_path is not None: + args += ["--cookies", str(cookies_path)] + args.append(thread.source_url) + try: + completed = subprocess.run(args, capture_output=True, text=True) + except OSError as e: + return Result.err(make_error("GalleryDlError", "download_media", str(e))) + if completed.returncode != 0: + return Result.err(make_error("GalleryDlError", "download_media", completed.stderr[:500])) + post_ids = {p.post_id for p in thread.posts} + files = [f for f in output_dir.rglob("*") if f.is_file() and f.name.split("_")[0] in post_ids] + return Result.ok(sorted(files, key=lambda f: f.name)) + +def media_files_for_post(output_dir: Path, post_id: str) -> list[str]: + return sorted(f.name for f in output_dir.glob(f"{post_id}_*") if f.is_file()) def main(argv: list[str] | None = None) -> int: - parser = argparse.ArgumentParser(description="Download media for a Twitter/X thread_data.json.") + parser = argparse.ArgumentParser(description="Download media for a Twitter/X thread_data.json via gallery-dl.") parser.add_argument("--input", type=Path, required=True, help="thread_data.json path") parser.add_argument("--output", type=Path, required=True, help="media output directory") + parser.add_argument("--cookies", type=Path, default=None, help="cookies.txt for gallery-dl auth") args = parser.parse_args(argv) try: raw = args.input.read_text(encoding="utf-8") @@ -77,7 +51,7 @@ def main(argv: list[str] | None = None) -> int: sys.stderr.write(f"error: JsonParseError: {e}\n") return 1 thread = thread_from_dict(data) - result = download_media(thread, args.output) + result = download_media(thread, args.output, args.cookies) if not result.is_ok: sys.stderr.write(f"error: {result.error.kind}: {result.error.detail}\n") return 1 diff --git a/scripts/twitter_threads/render_markdown.py b/scripts/twitter_threads/render_markdown.py index 98ab734d..057bc3cd 100644 --- a/scripts/twitter_threads/render_markdown.py +++ b/scripts/twitter_threads/render_markdown.py @@ -43,6 +43,14 @@ def _render_yaml(thread: ThreadData, title: str) -> list[str]: "---", ] +def _media_line(index: int, name: str) -> str: + ext = name.rsplit(".", 1)[-1].lower() if "." in name else "" + if ext in ("mp4", "mov", "webm", "m4v"): + return f'' + if ext in ("jpg", "jpeg", "png", "gif", "webp", "bmp"): + return f"![Media {index}](./media/{name})" + return f"[Media {index}](./media/{name})" + def _render_post(post: PostData, n: int, id_to_index: dict[str, int], media_names: dict[str, list[str]]) -> list[str]: lines: list[str] = [] header = f"## Post {n} ({post.timestamp})" @@ -58,7 +66,7 @@ def _render_post(post: PostData, n: int, id_to_index: dict[str, int], media_name names = media_names.get(post.post_id, []) if names: for k, name in enumerate(names, start=1): - lines.append(f"[Media {k}](./media/{name})") + lines.append(_media_line(k, name)) lines.append("") return lines diff --git a/tests/test_twitter_threads_media.py b/tests/test_twitter_threads_media.py index e5f4401a..ad753bac 100644 --- a/tests/test_twitter_threads_media.py +++ b/tests/test_twitter_threads_media.py @@ -1,82 +1,62 @@ -"""Tests for scripts.twitter_threads.download_media: naming and HTTP download contract.""" -from __future__ import annotations +"""Tests for scripts.twitter_threads.download_media gallery-dl subprocess 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 +from scripts.twitter_threads.download_media import download_media, media_files_for_post -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: +def _post(post_id: 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) + return PostData(post_id=post_id, author="Tim", handle="NOTimothyLottes", text="t", timestamp="2024-02-01T00:00:00Z", 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 _thread(post_id: str) -> ThreadData: + return ThreadData(root_post_id=post_id, posts=(_post(post_id),), source_url=f"https://x.com/x/status/{post_id}") -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_collects_matching_files(tmp_path: Path) -> None: + mediadir = tmp_path / "m" + mediadir.mkdir() + (mediadir / "123_1.jpg").write_bytes(b"x") + (mediadir / "123_2.png").write_bytes(b"x") + (mediadir / "999_1.jpg").write_bytes(b"x") + with patch("scripts.twitter_threads.download_media.subprocess.run", MagicMock(return_value=MagicMock(returncode=0))): + r = download_media(_thread("123"), mediadir) + assert r.is_ok + names = [p.name for p in r.data] + assert "123_1.jpg" in names + assert "123_2.png" in names + assert "999_1.jpg" not in names -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_invokes_gallery_dl_with_url_and_cookies(tmp_path: Path) -> None: + mediadir = tmp_path / "m" + ck = tmp_path / "cookies.txt" + ck.write_text("x", encoding="utf-8") + with patch("scripts.twitter_threads.download_media.subprocess.run", MagicMock(return_value=MagicMock(returncode=0))) as mock: + download_media(_thread("55"), mediadir, ck) + assert mock.called + args = mock.call_args.args[0] + assert "gallery-dl" in args[0] or args[0] == "gallery-dl" + assert "https://x.com/x/status/55" in args + assert "--cookies" in args + assert str(ck) in args -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_gallery_dl_error(tmp_path: Path) -> None: + with patch("scripts.twitter_threads.download_media.subprocess.run", MagicMock(return_value=MagicMock(returncode=1, stderr="boom"))): + r = download_media(_thread("123"), tmp_path / "m") + assert r.is_ok is False + assert r.error.kind == "GalleryDlError" -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 +def test_media_files_for_post(tmp_path: Path) -> None: + d = tmp_path / "m" + d.mkdir() + (d / "123_1.jpg").write_bytes(b"x") + (d / "123_2.png").write_bytes(b"x") + (d / "999_1.jpg").write_bytes(b"x") + assert media_files_for_post(d, "123") == ["123_1.jpg", "123_2.png"] diff --git a/tests/test_twitter_threads_pipeline.py b/tests/test_twitter_threads_pipeline.py index 49f5ec2b..7decb949 100644 --- a/tests/test_twitter_threads_pipeline.py +++ b/tests/test_twitter_threads_pipeline.py @@ -3,13 +3,13 @@ from __future__ import annotations import json from pathlib import Path -from unittest.mock import patch +from unittest.mock import patch, MagicMock import pytest from scripts.twitter_threads.error_types import PostMetrics, PostData, ThreadData, thread_from_dict from scripts.twitter_threads.fetch_thread import fetch_thread_from_html, thread_to_dict -from scripts.twitter_threads.download_media import media_names_for_post, download_media, main as download_main +from scripts.twitter_threads.download_media import download_media, main as download_main from scripts.twitter_threads.render_markdown import render_markdown, main as render_main @@ -59,7 +59,7 @@ def test_pipeline_html_to_markdown(tmp_path: Path) -> None: assert res.is_ok td = res.data assert td is not None - media_names = {p.post_id: media_names_for_post(p) for p in td.posts if p.media_urls} + media_names = {"10": ["10_img1.jpg"]} out = tmp_path / "thread.md" rr = render_markdown(td, media_names, out) assert rr.is_ok @@ -67,7 +67,7 @@ def test_pipeline_html_to_markdown(tmp_path: Path) -> None: assert "## Post 1" in text assert "## Post 2" in text assert "reply to Post 1" in text - assert "[Media 1](./media/10_img1.jpg)" in text + assert "![Media 1](./media/10_img1.jpg)" in text def test_download_cli(tmp_path: Path) -> None: @@ -76,10 +76,12 @@ def test_download_cli(tmp_path: Path) -> None: thread_data_json = tmp_path / "thread_data.json" thread_data_json.write_text(json.dumps(thread_to_dict(t)), encoding="utf-8") mediadir = tmp_path / "media" - with patch("scripts.twitter_threads.download_media.urlopen", return_value=_FakeResp(b"x")): + mediadir.mkdir() + (mediadir / "5_1.jpg").write_bytes(b"x") + with patch("scripts.twitter_threads.download_media.subprocess.run", MagicMock(return_value=MagicMock(returncode=0))): code = download_main(["--input", str(thread_data_json), "--output", str(mediadir)]) assert code == 0 - assert (mediadir / "5_img1.jpg").exists() + assert (mediadir / "5_1.jpg").exists() def test_render_cli(tmp_path: Path) -> None: diff --git a/tests/test_twitter_threads_render.py b/tests/test_twitter_threads_render.py index 2a9302af..11dc5c3a 100644 --- a/tests/test_twitter_threads_render.py +++ b/tests/test_twitter_threads_render.py @@ -77,9 +77,9 @@ def test_render_media_links(tmp_path: Path) -> None: result = render_markdown(thread, media_names, out) assert result.is_ok text = out.read_text(encoding="utf-8") - assert "[Media 1](./media/123_img1.jpg)" in text - assert "[Media 2](./media/123_img2.jpg)" in text - assert "[Media 3](./media/123_vid1.mp4)" in text + assert "![Media 1](./media/123_img1.jpg)" in text + assert "![Media 2](./media/123_img2.jpg)" in text + assert '' in text def test_render_view_count_null(tmp_path: Path) -> None: