Create Leetcode Python file from Leetcode url

leetcode_file.py

import argparse
import json
import os
import re

import requests
from bs4 import BeautifulSoup


TYPING_STRINGS = ["List", "Optional"]


def extract_function_name(code):
    match = re.search(r"def (\w+)\(", code)
    if not match:
        function_name = ""
        print("No function name found")
    else:
        function_name = match.group(1)
    return function_name


def extract_typing_imports(code):
    type_imports = []
    for type_string in TYPING_STRINGS:
        if re.search(type_string, code):
            type_imports.append(type_string)
    return type_imports


def create_leetcode_file(url, overwrite):
    problem_name = url.strip("/").split("/")[-1]
    filename = f"{problem_name}.py"

    if os.path.isfile(filename) and not overwrite:
        print(f"File {filename} already exists. " "Use -o flag to overwrite.")
        return

    response = requests.get(url)

    if response.status_code != 200:
        print(
            "Failed to download the Leetcode page,"
            f"status code: {response.status_code}"
        )
        return

    soup = BeautifulSoup(response.content, "html.parser")

    # Find the JSON data within the script tags
    scripts = soup.find_all("script")
    code_json = None
    for script in scripts:
        if "langSlug" in script.text:
            code_json = re.search(
                r"({\"lang\":\"Python3\",\"langSlug\":"
                r"\"python3\",\"code\":.*?})",
                script.text,
            )
            if code_json:
                code_json = code_json.group(1)
                break

    if not code_json:
        print("Failed to extract the code JSON data.")
        return

    code_data = json.loads(code_json)
    code = code_data.get("code", "")

    function_name = extract_function_name(code)
    typing_imports = extract_typing_imports(code)

    with open(f"{problem_name}.py", "w") as f:
        f.write(f"# Leetcode {url}\n")
        if typing_imports:
            f.write(f"from typing import {', '.join(typing_imports)}\n")
        f.write("\n" * 2)
        f.write(code)
        f.write("pass")
        f.write("\n" * 3)
        f.write("s = Solution()\n")
        f.write(f"print(s.{function_name}())\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Create a LeetCode Python file with problem boilerplate."
    )
    parser.add_argument("url", type=str, help="URL of the Leetcode problem.")
    parser.add_argument(
        "-o",
        "--overwrite",
        action="store_true",
        help="Overwrite file if it exists",
    )

    args = parser.parse_args()
    create_leetcode_file(args.url, args.overwrite)