110 lines
3.2 KiB
Python
110 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import json
|
|
import sys
|
|
|
|
import requests
|
|
from bs4 import BeautifulSoup
|
|
|
|
|
|
def parse_problem_url(contest_id: str, problem_letter: str) -> str:
|
|
task_id = f"{contest_id}_{problem_letter}"
|
|
return f"https://atcoder.jp/contests/{contest_id}/tasks/{task_id}"
|
|
|
|
|
|
def scrape(url: str) -> list[tuple[str, str]]:
|
|
try:
|
|
headers = {
|
|
"User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
|
|
}
|
|
|
|
response = requests.get(url, headers=headers, timeout=10)
|
|
response.raise_for_status()
|
|
|
|
soup = BeautifulSoup(response.text, "html.parser")
|
|
|
|
tests = []
|
|
|
|
sample_headers = soup.find_all(
|
|
"h3", string=lambda x: x and "sample" in x.lower() if x else False
|
|
)
|
|
|
|
i = 0
|
|
while i < len(sample_headers):
|
|
header = sample_headers[i]
|
|
if "input" in header.get_text().lower():
|
|
input_pre = header.find_next("pre")
|
|
if input_pre and i + 1 < len(sample_headers):
|
|
next_header = sample_headers[i + 1]
|
|
if "output" in next_header.get_text().lower():
|
|
output_pre = next_header.find_next("pre")
|
|
if output_pre:
|
|
input_text = input_pre.get_text().strip().replace("\r", "")
|
|
output_text = (
|
|
output_pre.get_text().strip().replace("\r", "")
|
|
)
|
|
if input_text and output_text:
|
|
tests.append((input_text, output_text))
|
|
i += 2
|
|
continue
|
|
i += 1
|
|
|
|
return tests
|
|
|
|
except Exception as e:
|
|
print(f"Error scraping AtCoder: {e}", file=sys.stderr)
|
|
return []
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) != 3:
|
|
result = {
|
|
"success": False,
|
|
"error": "Usage: atcoder.py <contest_id> <problem_letter>",
|
|
"problem_id": None,
|
|
}
|
|
print(json.dumps(result))
|
|
sys.exit(1)
|
|
|
|
contest_id = sys.argv[1]
|
|
problem_letter = sys.argv[2]
|
|
problem_id = contest_id + problem_letter.lower()
|
|
|
|
url = parse_problem_url(contest_id, problem_letter)
|
|
print(f"Scraping: {url}", file=sys.stderr)
|
|
|
|
tests = scrape(url)
|
|
|
|
if not tests:
|
|
result = {
|
|
"success": False,
|
|
"error": f"No tests found for {contest_id} {problem_letter}",
|
|
"problem_id": problem_id,
|
|
"url": url,
|
|
}
|
|
print(json.dumps(result))
|
|
sys.exit(1)
|
|
|
|
test_cases = []
|
|
for input_data, output_data in tests:
|
|
test_cases.append({"input": input_data, "output": output_data})
|
|
|
|
if test_cases:
|
|
combined_input = (
|
|
str(len(test_cases)) + "\n" + "\n".join(tc["input"] for tc in test_cases)
|
|
)
|
|
combined_output = "\n".join(tc["output"] for tc in test_cases)
|
|
test_cases = [{"input": combined_input, "output": combined_output}]
|
|
|
|
result = {
|
|
"success": True,
|
|
"problem_id": problem_id,
|
|
"url": url,
|
|
"test_cases": test_cases,
|
|
}
|
|
|
|
print(json.dumps(result))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|