manual_slop/scripts/twitter_threads/fetch_thread.py

"""Acquire a Twitter/X thread into a ThreadData: gallery-dl (URL) or html.parser (local HTML)."""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from dataclasses import asdict
from html.parser import HTMLParser
from pathlib import Path
from typing import Any
from urllib.parse import urlparse, urlunparse

try:
 from scripts.twitter_threads.error_types import ErrorInfo, make_error, Result, PostMetrics, PostData, ThreadData
except ImportError:
 from error_types import ErrorInfo, make_error, Result, PostMetrics, PostData, ThreadData


def _normalize_url(url: str) -> str:
 parts = urlparse(url)
 return urlunparse((parts.scheme, parts.netloc, parts.path, parts.params, "", ""))


def _to_int(value: str | None, default: int | None) -> int | None:
 if value is None or value == "":
  return default
 try:
  return int(value)
 except ValueError:
  return default


class _ThreadHTMLParser(HTMLParser):
 def __init__(self) -> None:
  super().__init__(convert_charrefs=True)
  self.posts: list[PostData] = []
  self.source_url: str = ""
  self._attrs: dict[str, str] = {}
  self._media: list[str] = []
  self._timestamp: str = ""
  self._text_parts: list[str] = []
  self._in_article: bool = False
  self._in_text: bool = False
  self._text_depth: int = 0

 def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
  a = {k: (v or "") for k, v in attrs}
  if tag == "link" and a.get("rel") == "canonical" and a.get("href"):
   self.source_url = a["href"]
  if tag == "article" and "data-post-id" in a:
   self._in_article = True
   self._attrs = a
   self._media = []
   self._timestamp = ""
   self._text_parts = []
   self._in_text = False
   self._text_depth = 0
  if not self._in_article:
   return
  if tag == "time" and a.get("datetime"):
   self._timestamp = a["datetime"]
  if tag == "div" and a.get("data-testid") == "tweetText":
   self._in_text = True
   self._text_depth = 1
   self._text_parts = []
  elif tag == "div" and self._in_text:
   self._text_depth += 1
  if tag in ("img", "source", "video") and a.get("src"):
   self._media.append(a["src"])

 def handle_data(self, data: str) -> None:
  if self._in_text:
   self._text_parts.append(data)

 def handle_endtag(self, tag: str) -> None:
  if tag == "div" and self._in_text:
   self._text_depth -= 1
   if self._text_depth <= 0:
    self._in_text = False
  if tag == "article" and self._in_article:
   self.posts.append(self._build())
   self._in_article = False

 def _build(self) -> PostData:
  a = self._attrs
  metrics = PostMetrics(
   reply_count=_to_int(a.get("data-reply-count"), 0),
   repost_count=_to_int(a.get("data-repost-count"), 0),
   like_count=_to_int(a.get("data-like-count"), 0),
   view_count=_to_int(a.get("data-view-count"), None),
  )
  return PostData(
   post_id=a.get("data-post-id", ""),
   author=a.get("data-author", ""),
   handle=a.get("data-handle", ""),
   text="".join(self._text_parts).strip(),
   timestamp=self._timestamp,
   media_urls=tuple(self._media),
   reply_to_id=a.get("data-reply-to"),
   quote_of_id=a.get("data-quote-of"),
   metrics=metrics,
  )


def fetch_thread_from_html(html_path: Path) -> Result[ThreadData]:
 try:
  text = Path(html_path).read_text(encoding="utf-8")
 except OSError as e:
  return Result.err(make_error("ReadError", "fetch_thread_from_html", str(e)))
 parser = _ThreadHTMLParser()
 try:
  parser.feed(text)
  parser.close()
 except Exception as e:
  return Result.err(make_error("ParseError", "fetch_thread_from_html", str(e)))
 if not parser.posts:
  return Result.err(make_error("ParseError", "fetch_thread_from_html", "no <article> elements found"))
 posts = tuple(parser.posts)
 return Result.ok(ThreadData(root_post_id=posts[0].post_id, posts=posts, source_url=parser.source_url))


def _find_tweet_meta(entry: Any) -> dict[str, Any] | None:
 if isinstance(entry, dict):
  if "tweet_id" in entry or ("content" in entry and "author" in entry):
   return entry
  return None
 if isinstance(entry, (list, tuple)):
  for item in entry:
   meta = _find_tweet_meta(item)
   if meta is not None:
    return meta
 return None


def _media_url_from_entry(entry: Any) -> str:
 if isinstance(entry, (list, tuple)):
  for item in entry:
   if isinstance(item, str) and ("pbs.twimg.com" in item or "video.twimg.com" in item):
    return item
 return ""


def _post_from_meta(meta: dict[str, Any]) -> PostData:
 author = meta.get("author") or {}
 if isinstance(author, dict):
  author_name = str(author.get("name") or author.get("nick") or "")
  handle = str(author.get("nick") or author.get("name") or "")
 else:
  author_name = str(author)
  handle = str(author)
 view = meta.get("view_count")
 metrics = PostMetrics(
  reply_count=int(meta.get("reply_count") or 0),
  repost_count=int(meta.get("retweet_count") or 0),
  like_count=int(meta.get("favorite_count") or 0),
  view_count=int(view) if view is not None else None,
 )
 return PostData(
  post_id=str(meta.get("tweet_id") or meta.get("id") or ""),
  author=author_name,
  handle=handle,
  text=str(meta.get("content") or ""),
  timestamp=str(meta.get("date") or ""),
  media_urls=(),
  reply_to_id=str(meta["reply_id"]) if meta.get("reply_id") else None,
  quote_of_id=str(meta["quote_id"]) if meta.get("quote_id") else None,
  metrics=metrics,
 )


def _posts_from_gallery_dl(data: Any) -> list[PostData]:
 by_id: dict[str, PostData] = {}
 media_by_id: dict[str, list[str]] = {}
 order: list[str] = []
 entries = data if isinstance(data, list) else []
 for entry in entries:
  meta = _find_tweet_meta(entry)
  if meta is None:
   continue
  tid = str(meta.get("tweet_id") or meta.get("id") or "")
  if not tid:
   continue
  if tid not in by_id:
   by_id[tid] = _post_from_meta(meta)
   media_by_id[tid] = []
   order.append(tid)
  media_url = _media_url_from_entry(entry)
  if media_url:
   media_by_id[tid].append(media_url)
 posts: list[PostData] = []
 for tid in order:
  base = by_id[tid]
  posts.append(PostData(
   post_id=base.post_id,
   author=base.author,
   handle=base.handle,
   text=base.text,
   timestamp=base.timestamp,
   media_urls=tuple(media_by_id[tid]),
   reply_to_id=base.reply_to_id,
   quote_of_id=base.quote_of_id,
   metrics=base.metrics,
  ))
 return posts


def fetch_thread_from_url(url: str, cookies_path: Path | None = None) -> Result[ThreadData]:
 clean = _normalize_url(url)
 args = ["gallery-dl", "--dump-json"]
 if cookies_path is not None:
  args += ["--cookies", str(cookies_path)]
 args.append(clean)
 try:
  completed = subprocess.run(args, capture_output=True, text=True)
 except OSError as e:
  return Result.err(make_error("GalleryDlError", "fetch_thread_from_url", str(e)))
 if completed.returncode != 0:
  return Result.err(make_error("GalleryDlError", "fetch_thread_from_url", completed.stderr[:500]))
 try:
  data = json.loads(completed.stdout)
 except json.JSONDecodeError as e:
  return Result.err(make_error("JsonParseError", "fetch_thread_from_url", str(e)))
 posts = _posts_from_gallery_dl(data)
 if not posts:
  return Result.err(make_error("ParseError", "fetch_thread_from_url", "no tweet metadata in gallery-dl output"))
 return Result.ok(ThreadData(root_post_id=posts[0].post_id, posts=tuple(posts), source_url=clean))


def thread_to_dict(thread: ThreadData) -> dict[str, Any]:
 return asdict(thread)


def main(argv: list[str] | None = None) -> int:
 parser = argparse.ArgumentParser(description="Fetch a Twitter/X thread into a ThreadData JSON.")
 parser.add_argument("source", help="X.com URL or local HTML file path")
 parser.add_argument("--output", type=Path, required=True, help="Output directory")
 parser.add_argument("--cookies", type=Path, default=None, help="cookies.txt for gallery-dl auth")
 args = parser.parse_args(argv)
 if args.source.startswith("http://") or args.source.startswith("https://"):
  result = fetch_thread_from_url(args.source, args.cookies)
 else:
  result = fetch_thread_from_html(Path(args.source))
 if not result.is_ok:
  sys.stderr.write(f"error: {result.error.kind}: {result.error.detail}" + "`r`n")
  return 1
 thread = result.data
 out_dir = args.output / thread.root_post_id
 try:
  out_dir.mkdir(parents=True, exist_ok=True)
  out_path = out_dir / "thread_data.json"
  out_path.write_text(json.dumps(thread_to_dict(thread), indent=2, ensure_ascii=False), encoding="utf-8")
 except OSError as e:
  sys.stderr.write(f"error: WriteError: {e}" + "`r`n")
  return 1
 print(str(out_path))
 return 0


if __name__ == "__main__":
 sys.exit(main())