download_wikidocs.py

import time
import requests
from bs4 import BeautifulSoup
import re
from markdownify import markdownify as md
import pandas as pd
import argparse


def extract_content(url: str):
    response = requests.get(url)
    soup = BeautifulSoup(response.content, "html.parser")

    page_subject = soup.select_one("#load_content .page-subject")
    page_content = soup.select_one("#load_content .page-content")
    markdown_content = md(
        str(page_subject) + str(page_content),
        heading_style="ATX",
        bullets="-",
        strong_em_symbol="*",
        code_language="python",
        escape_asterisks=False,
        escape_underscores=False,
    )
    normalized_text = re.sub(r"\n{2}", "\n", markdown_content)

    return normalized_text


def main(ebook_url):
    base_url = "https://wikidocs.net"

    # book_id 추출
    book_id = ebook_url.split("/")[-1]

    # 페이지 소스 가져오기
    response = requests.get(ebook_url)
    response.raise_for_status()  # 예외 처리
    soup = BeautifulSoup(response.content, "html.parser")

    # 목차에서 'a' 태그만 가져오기
    toc = soup.select(".list-group-toc a[href^='javascript:page(']")

    # 추출한 데이터 저장할 리스트
    data_list = []
    for item in toc:
        title = item.get_text(strip=True)
        page_id = item.get("href").split("page(")[-1].rstrip(")")
        link = f"{base_url}/{page_id}"
        data_list.append({"title": title, "link": link})

    # 데이터 리스트를 순회하며 콘텐츠 추출
    for item in data_list[1:]:
        item["content"] = extract_content(item["link"])
        time.sleep(1)  # 페이지 로드를 위해 대기

    # 데이터프레임으로 변환
    df = pd.DataFrame(data_list)
    df = df.dropna(subset=["content"])

    # 데이터프레임을 parquet 파일로 저장
    parquet_filename = f"wikidocs_{book_id}.parquet"
    df.to_parquet(parquet_filename, index=False)

    print(f"파일이 성공적으로 저장되었습니다: {parquet_filename}")


if __name__ == "__main__":
    # 명령어 줄 인자 처리
    parser = argparse.ArgumentParser(description="Wikidocs ebook URL을 입력하세요.")
    parser.add_argument("ebook_url", type=str, help="Wikidocs ebook URL")
    args = parser.parse_args()

    main(args.ebook_url)