process_submissions.py

import os
import shutil
import json
import tqdm
import tempfile
import zipfile
from dataclasses import dataclass

# ================== CONFIG ================== #
INPUT_ZIP = 'submissions.zip'
OUTPUT_DIR = 'solutions'
# ============================================ #

COMMENT_MAP = {
    '.py': '#',
    '.cpp': '//',
    '.java': '//',
    '.c': '//',
    '.js': '//',
}

@dataclass
class Submission:
    id: int
    question_id: int
    lang: str
    lang_name: str
    time: str
    timestamp: int
    status: int
    status_display: str
    runtime: str
    url: str
    is_pending: str
    title: str
    memory: str
    code: str
    compare_result: str
    title_slug: str
    has_notes: bool
    flag_type: int

# read the submissions dir and extract the accepted submissions
# yields lists of paths to the accepted submissions
def get_accepted_submissions():
    for filename in os.listdir(INPUT_DIR):
        base = os.path.join(INPUT_DIR, filename, 'Accepted')
        if not os.path.exists(base):
            continue
        yield [os.path.join(INPUT_DIR, filename, 'Accepted', submission)
                for submission in os.listdir(base)]

# todo: use leetcoed api to get info about percentile, difficulty, etc
def get_info(submission_path : str) -> Submission:
    with open(os.path.join(submission_path, 'info.txt')) as f:
        data = json.load(f)
    return Submission(**data)

# Take fastest python solution, or fastest solution if no python solution
def heuristic(x : Submission):
    lang = x.lang_name
    runtime = int(x.runtime.split(" ")[0])
    if lang == 'Python': # take python
        return 100_000 - runtime
    else:
        return -runtime

def select_submission(submissions : list[str]):
    return max(submissions, key=lambda x : heuristic(get_info(x)))

# fix the incorrect comment generated by submission download script
def correct_comment(submission : str, info : Submission, language : str):
    if language not in COMMENT_MAP:
        print("Unknown language", language)
        return

    comment = COMMENT_MAP[language]
    with open(submission, 'r') as f:
        lines = f.readlines()
        lines[0] = f"{comment} Problem: {lines[0][3:]}"
        lines.insert(1, f"{comment} Runtime: {info.runtime}\n")
        # if desired, add more info here
    with open(submission, 'w') as f:
        f.writelines(lines)

def write_submission(submission_path : str):
    info = get_info(submission_path)
    base = None
    for file in os.listdir(submission_path):
        basename = os.path.basename(file)
        if os.path.splitext(basename)[0] == 'Solution':
            base = file
            break
    if base is None:
        return

    extension = os.path.splitext(base)[1]
    input_path = os.path.join(submission_path, base)
    output_path = os.path.join(OUTPUT_DIR, f'{info.question_id:04d}_{info.title_slug}{extension}')
    
    shutil.copy(input_path, output_path)
    correct_comment(output_path, info, extension)

if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)

with tempfile.TemporaryDirectory() as temp_dir:
    with zipfile.ZipFile('submissions.zip', 'r') as zip_ref:
        zip_ref.extractall(temp_dir)
    INPUT_DIR = temp_dir # kinda hacky but works -- maybe fix later
    submissions = list(get_accepted_submissions())
    for submissions in tqdm.tqdm(submissions):
        chosen = select_submission(submissions)
        write_submission(chosen)