AtCoderStudyBooster 0.3__py3-none-any.whl → 0.3.2__py3-none-any.whl

atcdr/generate.py

@@ -2,105 +2,94 @@ import json
 import os
 import re
 
+import rich_click as click
 from rich.console import Console
+from rich.panel import Panel
+from rich.syntax import Syntax
 
-from atcdr.test import (
-	LabeledTestCaseResult,
-	ResultStatus,
-	create_testcases_from_html,
-	judge_code_from,
-	render_results,
-)
-from atcdr.util.execute import execute_files
+from atcdr.test import ResultStatus, TestRunner, create_renderable_test_info
+from atcdr.util.fileops import add_file_selector
 from atcdr.util.filetype import (
-	FILE_EXTENSIONS,
-	Filename,
-	Lang,
-	lang2str,
-	str2lang,
+    FILE_EXTENSIONS,
+    Filename,
+    Lang,
+    lang2str,
+    str2lang,
 )
-from atcdr.util.gpt import ChatGPT, set_api_key
-from atcdr.util.problem import make_problem_markdown
+from atcdr.util.gpt import ChatGPT, Model, set_api_key
+from atcdr.util.parse import ProblemHTML
 
 
 def get_code_from_gpt_output(output: str) -> str:
-	pattern = re.compile(r'```(?:\w+)?\s*(.*?)\s*```', re.DOTALL)
-	match = pattern.search(output)
-	return match.group(1) if match else ''
-
-
-def render_result_for_GPT(lresult: LabeledTestCaseResult) -> str:
-	output = f'{lresult.label} of Test:\n'
-	result = lresult.result
-	testcase = lresult.testcase
-
-	if result.passed == ResultStatus.AC:
-		output += 'Accepted !! \n'
-
-	elif result.passed == ResultStatus.WA:
-		output += (
-			f'Wrong Answer\n'
-			f'Output:\n{result.output}\n'
-			f'Expected Output:\n{testcase.output}\n'
-		)
-
-	elif result.passed == ResultStatus.RE:
-		output += f'[RE] Runtime Error\n  Output:\n{result.output}'
-
-	elif result.passed == ResultStatus.TLE:
-		output += '[TLE] Time Limit Exceeded\n'
-
-	elif result.passed == ResultStatus.CE:
-		output += f'[CE] Compile Error\n  Output:\n{result.output}'
-
-	elif result.passed == ResultStatus.MLE:
-		output += '[ME] Memory Limit Exceeded\n'
-
-	return output
-
-
-def generate_code(file: Filename, lang: Lang) -> None:
-	console = Console()
-	with open(file, 'r') as f:
-		html_content = f.read()
-	md = make_problem_markdown(html_content, 'en')
-
-	if set_api_key() is None:
-		return
-	gpt = ChatGPT(
-		system_prompt=f"""You are an excellent programmer. You solve problems in competitive programming.When a user provides you with a problem from a programming contest called AtCoder, including the Problem,Constraints, Input, Output, Input Example, and Output Example, please carefully consider these and solve the problem.Make sure that your output code block contains no more than two blocks. Pay close attention to the Input, Input Example, Output, and Output Example.Create the solution in {lang2str(lang)}.""",
-	)
-	with console.status(f'{gpt.model.value}がコードを生成しています...'):
-		reply = gpt.tell(md)
-
-	code = get_code_from_gpt_output(reply)
-
-	saved_filename = (
-		os.path.splitext(file)[0] + f'_by_{gpt.model.value}' + FILE_EXTENSIONS[lang]
-	)
-	with open(saved_filename, 'w') as f:
-		console.print(
-			f'[green][+][/green] {gpt.model.value} の出力したコードを保存しました：{f.name}'
-		)
-		f.write(code)
-
-	console.print(f'[info] AI利用にかかったAPIコスト: {gpt.sum_cost}')
+    pattern = re.compile(r'```(?:\w+)?\s*(.*?)\s*```', re.DOTALL)
+    match = pattern.search(output)
+    return match.group(1) if match else ''
+
+
+def render_result_for_GPT(
+    test: TestRunner,
+) -> tuple[str, bool]:
+    results = list(test)
+
+    match test.info.summary:
+        case ResultStatus.AC:
+            return 'Accepted', True
+        case ResultStatus.CE:
+            return f'Compile Error \n {test.info.compiler_message}', False
+        case _:
+            message_for_gpt = ''.join(
+                f'\n{result.label} => {result.result.passed.value}\nInput :\n{result.testcase.input}\nOutput :\n{result.result.output}\nExpected :\n{result.testcase.output}\n'
+                if result.result.passed == ResultStatus.WA
+                else f'\n{result.label} => {result.result.passed.value}\nInput :\n{result.testcase.input}\nOutput :\n{result.result.output}\n'
+                for result in results
+            )
+            return message_for_gpt, False
+
+
+def generate_code(file: Filename, lang: Lang, model: Model) -> None:
+    console = Console()
+    with open(file, 'r') as f:
+        html_content = f.read()
+    md = ProblemHTML(html_content).make_problem_markdown('en')
+
+    if set_api_key() is None:
+        return
+    gpt = ChatGPT(
+        system_prompt=f"""You are an excellent programmer. You solve problems in competitive programming.When a user provides you with a problem from a programming contest called AtCoder, including the Problem,Constraints, Input, Output, Input Example, and Output Example, please carefully consider these and solve the problem.Make sure that your output code block contains no more than two blocks. Pay close attention to the Input, Input Example, Output, and Output Example.Create the solution in {lang2str(lang)}.""",
+        model=model,
+    )
+    with console.status(f'コード生成中 (by {gpt.model.value})'):
+        reply = gpt.tell(md)
+
+    code = get_code_from_gpt_output(reply)
+    console.print('[green][+][/green] コードの生成に成功しました. ')
+    console.rule(f'{gpt.model.value}による{lang2str(lang)}コード')
+    console.print(Syntax(code=code, lexer=lang2str(lang)))
+
+    saved_filename = (
+        os.path.splitext(file)[0] + f'_by_{gpt.model.value}' + FILE_EXTENSIONS[lang]
+    )
+    with open(saved_filename, 'w') as f:
+        console.print(
+            f'[green][+][/green] {gpt.model.value} の出力したコードを保存しました：{f.name}'
+        )
+        f.write(code)
 
 
 def generate_template(file: Filename, lang: Lang) -> None:
-	console = Console()
-	with open(file, 'r') as f:
-		html_content = f.read()
-	md = make_problem_markdown(html_content, 'en')
-
-	if set_api_key() is None:
-		return
-	gpt = ChatGPT(
-		system_prompt='You are a highly skilled programmer. Your role is to create a template code for competitive programming.',
-		temperature=0.0,
-	)
-
-	propmpt = f"""
+    console = Console()
+    with open(file, 'r') as f:
+        html_content = f.read()
+    md = ProblemHTML(html_content).make_problem_markdown('en')
+
+    if set_api_key() is None:
+        return
+    gpt = ChatGPT(
+        system_prompt='You are a highly skilled programmer. Your role is to create a template code for competitive programming.',
+        temperature=0.0,
+    )
+
+    propmpt = f"""
 The user will provide a problem from a programming contest called AtCoder. This problem will include the Problem Statement, Constraints, Input, Output, Input Example, and Output Example. You should focus on the Constraints and Input sections to create the template in {lang2str(lang)}.
 
 - First, create the part of the code that handles input. Then, you should read ###Input Block and ###Constraints Block.
@@ -109,112 +98,105 @@ The user will provide a problem from a programming contest called AtCoder. This
 
 You must not solve the problem. Please faithfully reproduce the variable names defined in the problem.
     """
-	with console.status(f'{lang2str(lang)}のテンプレートを生成しています...'):
-		reply = gpt.tell(md + propmpt)
-	code = get_code_from_gpt_output(reply)
-
-	savaed_filename = os.path.splitext(file)[0] + FILE_EXTENSIONS[lang]
-	with open(savaed_filename, 'x') as f:
-		console.print(
-			f'[green][+][/green] テンプレートファイルを作成 :{savaed_filename}'
-		)
-		f.write(code)
-
-	console.print(f'[info] AI利用にかかったAPIコスト: {gpt.sum_cost}')
-
-
-def solve_problem(file: Filename, lang: Lang) -> None:
-	console = Console()
-	with open(file, 'r') as f:
-		html_content = f.read()
-	md = make_problem_markdown(html_content, 'en')
-	labeled_cases = create_testcases_from_html(html_content)
-
-	if set_api_key() is None:
-		return
-	gpt = ChatGPT(
-		system_prompt=f"""You are a brilliant programmer. Your task is to solve an AtCoder problem. AtCoder is a platform that hosts programming competitions where participants write programs to solve algorithmic challenges.Please solve the problem in {lang2str(lang)}.""",
-	)
-
-	file_without_ext = os.path.splitext(file)[0]
-
-	for i in range(1, 4):
-		with console.status(f'{i}回目のコード生成 (by {gpt.model.value})...'):
-			test_report = ''
-			if i == 1:
-				reply = gpt.tell(md)
-			else:
-				reply = gpt.tell(f"""The following is the test report for the code you provided:
-	{test_report}
-	Please provide an updated version of the code in {lang2str(lang)}.""")
-
-		code = get_code_from_gpt_output(reply)
-
-		saved_filename = (
-			f'{i}_'
-			+ file_without_ext
-			+ f'_by_{gpt.model.value}'
-			+ FILE_EXTENSIONS[lang]
-		)
-		with open(saved_filename, 'w') as f:
-			console.print(
-				f'[green][+][/green] {gpt.model.value} の出力したコードを保存しました：{f.name}'
-			)
-			f.write(code)
-
-		labeled_results = judge_code_from(labeled_cases, saved_filename)
-		test_report = '\n'.join(
-			render_result_for_GPT(lresult) for lresult in labeled_results
-		)
-
-		console.rule(f'{i}回目のコード生成でのテスト結果')
-		render_results(saved_filename, labeled_results)
-
-		if all(
-			labeled_result.result.passed == ResultStatus.AC
-			for labeled_result in labeled_results
-		):
-			console.print('[green]コードのテストに成功![/green]')
-			break
-
-	with open(
-		'log_'
-		+ file_without_ext
-		+ f'_by_{gpt.model.value}'
-		+ FILE_EXTENSIONS[Lang.JSON],
-		'w',
-	) as f:
-		console.print(
-			f'[green][+][/green] {gpt.model.value}の出力のログを保存しました：{f.name}'
-		)
-		f.write(json.dumps(gpt.messages, indent=2))
-	console.print(f'AI利用にかかったAPIコスト:{gpt.sum_cost}')
-	return
-
-
-def generate(
-	*source: str,
-	lang: str = 'Python',
-	without_test: bool = False,
-	template: bool = False,
-) -> None:
-	la = str2lang(lang)
-
-	if template:
-		execute_files(
-			*source,
-			func=lambda file: generate_template(file, la),
-			target_filetypes=[Lang.HTML],
-		)
-	elif without_test:
-		execute_files(
-			*source,
-			func=lambda file: generate_code(file, la),
-			target_filetypes=[Lang.HTML],
-		)
-	else:
-		execute_files(
-			*source,
-			func=lambda file: solve_problem(file, la),
-			target_filetypes=[Lang.HTML],
-		)
+    with console.status(f'{lang2str(lang)}のテンプレートを生成しています...'):
+        reply = gpt.tell(md + propmpt)
+    code = get_code_from_gpt_output(reply)
+
+    savaed_filename = os.path.splitext(file)[0] + FILE_EXTENSIONS[lang]
+    with open(savaed_filename, 'x') as f:
+        console.print(
+            f'[green][+][/green] テンプレートファイルを作成 :{savaed_filename}'
+        )
+        f.write(code)
+
+
+def solve_problem(file: Filename, lang: Lang, model: Model) -> None:
+    console = Console()
+    with open(file, 'r') as f:
+        html = ProblemHTML(f.read())
+
+    md = html.make_problem_markdown('en')
+    labeled_cases = html.load_labeled_testcase()
+
+    if set_api_key() is None:
+        return
+    gpt = ChatGPT(
+        system_prompt=f"""You are a brilliant programmer. Your task is to solve an AtCoder problem. AtCoder is a platform that hosts programming competitions where participants write programs to solve algorithmic challenges.Please solve the problem in {lang2str(lang)}.""",
+        model=model,
+    )
+
+    file_without_ext = os.path.splitext(file)[0]
+
+    for i in range(1, 4):
+        with console.status(f'{i}回目のコード生成中 (by {gpt.model.value})'):
+            if i == 1:
+                test_report = ''
+                reply = gpt.tell(md)
+            else:
+                prompt = f"""The following is the test report for the code you provided:
+                {test_report}
+Please provide an updated version of the code in {lang2str(lang)}."""
+                console.print(
+                    f'[green][+][/] 次のプロンプトを{gpt.model.value}に与え,再生成します'
+                )
+                console.print(Panel(prompt))
+                reply = gpt.tell(prompt)
+
+        code = get_code_from_gpt_output(reply)
+
+        saved_filename = (
+            f'{i}_'
+            + file_without_ext
+            + f'_by_{gpt.model.value}'
+            + FILE_EXTENSIONS[lang]
+        )
+        with open(saved_filename, 'w') as f:
+            console.print(f'[green][+][/] コードの生成に成功しました！：{f.name}')
+            f.write(code)
+
+        with console.status(
+            f'{gpt.model.value}が生成したコードをテスト中', spinner='circleHalves'
+        ):
+            test = TestRunner(saved_filename, labeled_cases)
+            test_report, is_ac = render_result_for_GPT(test)
+
+        console.print(create_renderable_test_info(test.info))
+
+        if is_ac:
+            console.print('[green][+][/] コードのテストに成功!')
+            break
+        else:
+            console.print('[red][-][/] コードのテストに失敗!')
+
+    with open(
+        'log_'
+        + file_without_ext
+        + f'_by_{gpt.model.value}'
+        + FILE_EXTENSIONS[Lang.JSON],
+        'w',
+    ) as f:
+        console.print(
+            f'[green][+][/] {gpt.model.value}の出力のログを保存しました：{f.name}'
+        )
+        f.write(json.dumps(gpt.messages, indent=2))
+    return
+
+
+@click.command(short_help='コードを生成')
+@add_file_selector('files', filetypes=[Lang.HTML])
+@click.option('--lang', default='Python', help='出力するプログラミング言語')
+@click.option('--model', default=Model.GPT41_MINI.value, help='使用するGPTモデル')
+@click.option('--without-test', is_flag=True, help='テストケースを省略して生成')
+@click.option('--template', is_flag=True, help='テンプレートを生成')
+def generate(files, lang, model, without_test, template):
+    """HTMLファイルからコード生成またはテンプレート出力を行います。"""
+    la = str2lang(lang)
+    model_enum = Model(model)
+
+    for path in files:
+        if template:
+            generate_template(path, la)
+        elif without_test:
+            generate_code(path, la, model_enum)
+        else:
+            solve_problem(path, la, model_enum)

atcdr/login.py

@@ -0,0 +1,136 @@
+import threading
+import time
+
+import rich_click as click
+import webview
+from requests import Session
+from rich.console import Console
+
+from atcdr.util.session import load_session, save_session, validate_session
+
+ATCODER_LOGIN_URL = 'https://atcoder.jp/login'
+ATCODER_HOME_URL = 'https://atcoder.jp/home'
+
+console = Console()
+
+
+@click.command(short_help='AtCoderへログイン')
+def login() -> None:
+    """AtCoderへログインします."""
+    session = load_session()
+    if validate_session(session):
+        console.print('[green][+][/] すでにログインしています. ')
+        return
+
+    # Prompt in CLI
+    username = console.input('[cyan]ユーザー名: [/]').strip()
+    password = console.input('[cyan]パスワード: [/]').strip()
+
+    window = webview.create_window('AtCoder Login', ATCODER_LOGIN_URL, hidden=False)
+
+    def on_start():
+        js_fill = f"""
+        document.getElementById('username').value = '{username}';
+        document.getElementById('password').value = '{password}';
+        """
+        window.evaluate_js(js_fill)
+
+        def poll_and_submit():
+            with console.status(
+                'キャプチャー認証を解決してください', spinner='circleHalves'
+            ):
+                while True:
+                    try:
+                        token = window.evaluate_js(
+                            'document.querySelector(\'input[name=\\"cf-turnstile-response\\"]\').value'
+                        )
+                        if token:
+                            console.print('[green][+][/] ログインします')
+                            window.evaluate_js(
+                                "document.getElementById('submit').click();"
+                            )
+                            break
+                    except Exception:
+                        pass
+
+                    time.sleep(0.5)
+
+            with console.status('ログインの結果の待機中...', spinner='circleHalves'):
+                while True:
+                    try:
+                        current_url = window.get_current_url()
+                    except Exception:
+                        current_url = None
+
+                    if current_url and current_url.startswith(ATCODER_HOME_URL):
+                        console.print('[green][+][/] ログイン成功!')
+
+                        session = Session()
+                        session = move_cookies_from_webview_to_session(window, session)
+                        save_session(session)
+                        window.destroy()
+                        break
+
+                    try:
+                        err = window.evaluate_js(
+                            'Array.from(document.querySelectorAll('
+                            '\'div.alert.alert-danger[role="alert"]\'))'
+                            ".map(e=>e.textContent.trim()).filter(t=>t).join(' ')"
+                        )
+                        err = err.replace('\n', '').replace('\r', '').replace('\t', '')
+                    except Exception:
+                        err = ''
+
+                    if err:
+                        console.print(f'[red][-][/] エラー: {err}')
+                        session = Session()
+                        session = move_cookies_from_webview_to_session(window, session)
+                        save_session(session)
+                        window.destroy()
+                        return
+
+                    time.sleep(0.5)
+
+        t = threading.Thread(target=poll_and_submit, daemon=True)
+        t.start()
+
+    webview.start(on_start, private_mode=False)
+
+
+def move_cookies_from_webview_to_session(
+    window: webview.Window, session: Session
+) -> Session:
+    cookie_list = window.get_cookies()
+    for cookie_obj in cookie_list:
+        for cookie_name, morsel in cookie_obj.items():
+            # morselからデータを取得
+            value = morsel.value
+
+            domain = morsel.get('domain')
+            if domain is None:
+                domain = '.atcoder.jp'
+
+            path = morsel.get('path', '/')
+            secure = 'secure' in morsel
+
+            expires = None  # __NSTaggedDateオブジェクトを回避
+
+            http_only = 'httponly' in morsel
+
+            # HttpOnlyをrestに含める
+            rest = {}
+            if http_only:
+                rest['HttpOnly'] = True
+
+            # セッションにクッキーを設定
+            session.cookies.set(
+                name=cookie_name,
+                value=value,
+                domain=domain,
+                path=path,
+                secure=secure,
+                expires=expires,  # Noneを渡す
+                rest=rest,
+            )
+
+    return session

atcdr/logout.py

@@ -0,0 +1,27 @@
+import rich_click as click
+import webview
+from rich import print
+
+from atcdr.util.session import delete_session, load_session, validate_session
+
+ATCODER_LOGIN_URL = 'https://atcoder.jp/login'
+
+
+@click.command(short_help='AtCoderへログアウト')
+def logout() -> None:
+    """AtCoderからログアウトします."""
+    session = load_session()
+    if not validate_session(session):
+        print('[red][-][/] ログインしていません.')
+        return
+
+    delete_session()
+    print('[green][+][/] ログアウトしました.')
+
+    window = webview.create_window('AtCoder Logout', ATCODER_LOGIN_URL, hidden=True)
+
+    def on_start():
+        window.clear_cookies()
+        window.destroy()
+
+    webview.start(on_start, private_mode=False)

atcdr/markdown.py

@@ -1,39 +1,43 @@
 import os
 
+import rich_click as click
 from rich.console import Console
 from rich.markdown import Markdown
 
-from atcdr.util.execute import execute_files
+from atcdr.util.fileops import add_file_selector
 from atcdr.util.filetype import FILE_EXTENSIONS, Lang
-from atcdr.util.problem import make_problem_markdown
+from atcdr.util.parse import ProblemHTML
 
 
 def save_markdown(html_path: str, lang: str) -> None:
-	console = Console()
-	with open(html_path, 'r', encoding='utf-8') as f:
-		html = f.read()
-	md = make_problem_markdown(html, lang)
-	file_without_ext = os.path.splitext(html_path)[0]
-	md_path = file_without_ext + FILE_EXTENSIONS[Lang.MARKDOWN]
-
-	with open(md_path, 'w', encoding='utf-8') as f:
-		f.write(md)
-		console.print('[green][+][/green] Markdownファイルを作成しました.')
-
-
-def print_markdown(md_path: str) -> None:
-	console = Console()
-	with open(md_path, 'r', encoding='utf-8') as f:
-		md = f.read()
-	console.print(Markdown(md))
-
-
-def markdown(*args: str, lang: str = 'ja', save: bool = False) -> None:
-	if save:
-		execute_files(
-			*args,
-			func=lambda html_path: save_markdown(html_path, lang),
-			target_filetypes=[Lang.HTML],
-		)
-	else:
-		execute_files(*args, func=print_markdown, target_filetypes=[Lang.MARKDOWN])
+    console = Console()
+    with open(html_path, 'r', encoding='utf-8') as f:
+        html = ProblemHTML(f.read())
+    md = html.make_problem_markdown(lang)
+    file_without_ext = os.path.splitext(html_path)[0]
+    md_path = file_without_ext + FILE_EXTENSIONS[Lang.MARKDOWN]
+
+    with open(md_path, 'w', encoding='utf-8') as f:
+        f.write(md)
+        console.print('[green][+][/green] Markdownファイルを作成しました.')
+
+
+def print_markdown(html_path: str, lang: str) -> None:
+    console = Console()
+    with open(html_path, 'r', encoding='utf-8') as f:
+        html = ProblemHTML(f.read())
+    md = html.make_problem_markdown(lang)
+    console.print(Markdown(md))
+
+
+@click.command(short_help='Markdown形式で問題を表示します')
+@add_file_selector('files', filetypes=[Lang.HTML])
+@click.option('--lang', default='ja', help='出力する言語を指定')
+@click.option('--save', is_flag=True, help='変換結果をファイルに保存')
+def markdown(files, lang, save):
+    """Markdown形式で問題を表示します"""
+    for path in files:
+        if save:
+            save_markdown(path, lang)
+        else:
+            print_markdown(path, lang)