scrape.py

"""Scrape Poe's complete works from poemuseum.org."""

import os
import re
import time

import requests
from bs4 import BeautifulSoup

DATA_DIR = os.path.join(os.path.dirname(__file__), "data")

# All poemuseum.org links from https://poemuseum.org/poes-complete-works/
STORIES = [
    "https://poemuseum.org/the-angel-of-the-odd/",
    "https://poemuseum.org/the-assignation/",
    "https://poemuseum.org/the-balloon-hoax/",
    "https://poemuseum.org/berenice/",
    "https://poemuseum.org/the-black-cat/",
    "https://poemuseum.org/bon-bon/",
    "https://poemuseum.org/the-cask-of-amontillado/",
    "https://poemuseum.org/the-colloquy-of-monos-and-una/",
    "https://poemuseum.org/the-conversation-of-eiros-and-charmion/",
    "https://poemuseum.org/loss-of-breath/",
    "https://poemuseum.org/a-descent-into-the-maelstrom/",
    "https://poemuseum.org/the-devil-in-the-belfry/",
    "https://poemuseum.org/diddling-considered-as-one-of-the-exact-sciences/",
    "https://poemuseum.org/the-domain-of-arnheim/",
    "https://poemuseum.org/the-duke-de-lomelette/",
    "https://poemuseum.org/eleonora/",
    "https://poemuseum.org/epimanes/",
    "https://poemuseum.org/the-facts-in-the-case-of-m-valdemar/",
    "https://poemuseum.org/the-fall-of-the-house-of-usher/",
    "https://poemuseum.org/the-gold-bug/",
    "https://poemuseum.org/the-unparalleled-adventure-of-one-hans-pfaall/",
    "https://poemuseum.org/hop-frog/",
    "https://poemuseum.org/how-to-write-a-blackwood-article/",
    "https://poemuseum.org/the-imp-of-the-perverse/",
    "https://poemuseum.org/the-island-of-the-fay/",
    "https://poemuseum.org/king-pest/",
    "https://poemuseum.org/landors-cottage/",
    "https://poemuseum.org/ligeia/",
    "https://poemuseum.org/the-light-house/",
    "https://poemuseum.org/lionizing/",
    "https://poemuseum.org/the-literary-life-of-thingum-bob-esq/",
    "https://poemuseum.org/the-man-of-the-crowd/",
    "https://poemuseum.org/the-masque-of-the-red-death/",
    "https://poemuseum.org/mellonta-tauta/",
    "https://poemuseum.org/mesmeric-revelation/",
    "https://poemuseum.org/metzengerstein/",
    "https://poemuseum.org/morella/",
    "https://poemuseum.org/morning-on-the-wissahiccon/",
    "https://poemuseum.org/ms-found-in-a-bottle/",
    "https://poemuseum.org/the-murders-in-the-rue-morgue/",
    "https://poemuseum.org/the-mystery-of-marie-roget/",
    "https://poemuseum.org/mystification/",
    "https://poemuseum.org/never-bet-the-devil-your-head/",
    "https://poemuseum.org/the-oblong-box/",
    "https://poemuseum.org/the-oval-portrait/",
    "https://poemuseum.org/peter-pendulum/",
    "https://poemuseum.org/the-pit-and-the-pendulum/",
    "https://poemuseum.org/politian/",
    "https://poemuseum.org/the-power-of-words/",
    "https://poemuseum.org/the-premature-burial/",
    "https://poemuseum.org/the-purloined-letter/",
    "https://poemuseum.org/shadow-a-fable/",
    "https://poemuseum.org/silence-a-fable/",
    "https://poemuseum.org/the-spectacles/",
    "https://poemuseum.org/the-sphinx/",
    "https://poemuseum.org/the-system-of-doctor-tarr-and-professor-fether/",
    "https://poemuseum.org/a-tale-of-jerusalem/",
    "https://poemuseum.org/a-tale-of-the-ragged-mountains/",
    "https://poemuseum.org/the-tell-tale-heart/",
    "https://poemuseum.org/thou-art-the-man/",
    "https://poemuseum.org/the-thousand-and-second-tale-of-scheherazade/",
    "https://poemuseum.org/a-succession-of-sundays/",
    "https://poemuseum.org/von-kempelen-and-his-discovery/",
    "https://poemuseum.org/why-the-little-frenchman-wears-his-arm-in-a-sling/",
    "https://poemuseum.org/william-wilson/",
    "https://poemuseum.org/x-ing-a-paragrab/",
]

POEMS = [
    "https://poemuseum.org/an-acrostic/",
    "https://poemuseum.org/al-aaraaf/",
    "https://poemuseum.org/alone/",
    "https://poemuseum.org/annabel-lee/",
    "https://poemuseum.org/bridal-ballad/",
    "https://poemuseum.org/the-bells/",
    "https://poemuseum.org/beloved-physician/",
    "https://poemuseum.org/catholic-hymn/",
    "https://poemuseum.org/the-coliseum/",
    "https://poemuseum.org/the-conqueror-worm/",
    "https://poemuseum.org/the-divine-right-of-kings/",
    "https://poemuseum.org/the-doomed-city/",
    "https://poemuseum.org/a-dream/",
    "https://poemuseum.org/dream-land/",
    "https://poemuseum.org/dreams/",
    "https://poemuseum.org/a-dream-within-a-dream/",
    "https://poemuseum.org/eldorado/",
    "https://poemuseum.org/elizabeth/",
    "https://poemuseum.org/eulalie/",
    "https://poemuseum.org/evening-star/",
    "https://poemuseum.org/fairyland/",
    "https://poemuseum.org/fanny/",
    "https://poemuseum.org/for-annie/",
    "https://poemuseum.org/the-happiest-day/",
    "https://poemuseum.org/the-haunted-palace/",
    "https://poemuseum.org/imitation/",
    "https://poemuseum.org/impromptu/",
    "https://poemuseum.org/the-sleeper/",
    "https://poemuseum.org/israfel/",
    "https://poemuseum.org/the-lake/",
    "https://poemuseum.org/latin-hymn/",
    "https://poemuseum.org/lenore/",
    "https://poemuseum.org/lines-on-joe-locke/",
    "https://poemuseum.org/lines-written-in-an-album/",
    "https://poemuseum.org/o-tempora-o-mores/",
    "https://poemuseum.org/a-paean/",
    "https://poemuseum.org/preface/",
    "https://poemuseum.org/the-raven/",
    "https://poemuseum.org/serenade/",
    "https://poemuseum.org/song-of-triumph/",
    "https://poemuseum.org/sonnet/",
    "https://poemuseum.org/silence-a-sonnet/",
    "https://poemuseum.org/sonnet-to-science/",
    "https://poemuseum.org/sonnet-to-zante/",
    "https://poemuseum.org/spiritual-song/",
    "https://poemuseum.org/stanzas/",
    "https://poemuseum.org/stanzas-to-f-s-o/",
    "https://poemuseum.org/tamerlane/",
    "https://poemuseum.org/to/",
    "https://poemuseum.org/to-2/",
    "https://poemuseum.org/to-song/",
    "https://poemuseum.org/to-helen/",
    "https://poemuseum.org/to-helen-2/",
    "https://poemuseum.org/to-her-whose-name-is-written-below/",
    "https://poemuseum.org/to-m/",
    "https://poemuseum.org/to-m-l-s/",
    "https://poemuseum.org/to-margaret/",
    "https://poemuseum.org/to-mary/",
    "https://poemuseum.org/to-marie-louise/",
    "https://poemuseum.org/to-my-mother/",
    "https://poemuseum.org/to-one-in-paradise/",
    "https://poemuseum.org/to-violet-vane/",
    "https://poemuseum.org/ulalume/",
    "https://poemuseum.org/spirits-of-the-dead/",
    "https://poemuseum.org/the-valley-of-unrest/",
]

ALL_URLS = STORIES + POEMS


def extract_text(html: str) -> str:
    """Extract the main content text from a poemuseum.org page."""
    soup = BeautifulSoup(html, "html.parser")

    # The content is in .entry-content (WordPress)
    content = soup.select_one(".entry-content")
    if not content:
        # Fallback: try article tag
        content = soup.select_one("article")
    if not content:
        return ""

    # Remove script/style/nav elements
    for tag in content.find_all(["script", "style", "nav", "iframe"]):
        tag.decompose()

    # Get text, preserving line breaks from <br> and <p>
    # Replace <br> with newlines
    for br in content.find_all("br"):
        br.replace_with("\n")

    text = content.get_text(separator="\n")

    # Clean up
    text = text.replace("\r\n", "\n")
    # Collapse runs of 3+ newlines
    text = re.sub(r"\n{3,}", "\n\n", text)
    # Strip trailing whitespace per line
    text = "\n".join(line.rstrip() for line in text.split("\n"))
    return text.strip()


def slug_from_url(url: str) -> str:
    """Extract slug from URL for filename."""
    parts = url.rstrip("/").split("/")
    return parts[-1]


def scrape_all(delay: float = 1.0) -> str:
    """Scrape all works, cache individually, return concatenated corpus."""
    os.makedirs(DATA_DIR, exist_ok=True)

    raw_dir = os.path.join(DATA_DIR, "poemuseum_raw")
    os.makedirs(raw_dir, exist_ok=True)

    all_texts = []
    failed = []

    for i, url in enumerate(ALL_URLS):
        slug = slug_from_url(url)
        cache_path = os.path.join(raw_dir, f"{slug}.txt")

        if os.path.exists(cache_path):
            with open(cache_path, "r") as f:
                text = f.read()
            if text:
                all_texts.append(text)
                continue

        print(f"[{i+1}/{len(ALL_URLS)}] Fetching {slug}...")
        try:
            resp = requests.get(url, timeout=30, headers={
                "User-Agent": "Mozilla/5.0 (educational project - building tiny LLM)"
            })
            resp.raise_for_status()
            text = extract_text(resp.text)

            if text:
                with open(cache_path, "w") as f:
                    f.write(text)
                all_texts.append(text)
                print(f"  {len(text):,} chars")
            else:
                print(f"  WARNING: no content extracted")
                failed.append(url)

        except Exception as e:
            print(f"  ERROR: {e}")
            failed.append(url)

        time.sleep(delay)

    if failed:
        print(f"\nFailed to fetch {len(failed)} URLs:")
        for url in failed:
            print(f"  {url}")

    corpus = "\n\n---\n\n".join(all_texts)

    # Save combined corpus
    corpus_path = os.path.join(DATA_DIR, "poe_corpus.txt")
    with open(corpus_path, "w") as f:
        f.write(corpus)

    print(f"\nTotal works: {len(all_texts)}")
    print(f"Total corpus: {len(corpus):,} chars")
    print(f"Saved to {corpus_path}")

    return corpus


if __name__ == "__main__":
    scrape_all()