From bec5778dd0e3bd57008e4175a7d2be539540be36 Mon Sep 17 00:00:00 2001 From: shenchucheng Date: Sun, 6 Aug 2023 17:56:24 +0800 Subject: [PATCH 1/5] add write docstring action --- metagpt/actions/write_docstring.py | 169 ++++++++++++++++++ metagpt/utils/pycst.py | 169 ++++++++++++++++++ requirements.txt | 1 + tests/metagpt/actions/test_write_docstring.py | 32 ++++ tests/metagpt/utils/test_pycst.py | 136 ++++++++++++++ 5 files changed, 507 insertions(+) create mode 100644 metagpt/actions/write_docstring.py create mode 100644 metagpt/utils/pycst.py create mode 100644 tests/metagpt/actions/test_write_docstring.py create mode 100644 tests/metagpt/utils/test_pycst.py diff --git a/metagpt/actions/write_docstring.py b/metagpt/actions/write_docstring.py new file mode 100644 index 000000000..091bd0a82 --- /dev/null +++ b/metagpt/actions/write_docstring.py @@ -0,0 +1,169 @@ +import ast +import contextlib +from typing import Literal + +from metagpt.actions.action import Action +from metagpt.utils.common import OutputParser +from metagpt.utils.pycst import merge_docstring + +PYTHON_DOCSTRING_SYSTEM = '''### Requirements +1. Add docstrings to the given code following the {style} style. +2. Remove all private members whose names start with an underscore, such as `_test` and `__init__`. +3. Replace the function body with an Ellipsis object(...) to reduce output. +4. If the types are already annotated, there is no need to include them in the docstring. +5. Only output Python code and avoid including any other text. + +### Input Example +```python +def function_with_pep484_type_annotations(param1: int) -> bool: + return isinstanc(param1, int) + +class ExampleError(Exception): + def __init__(self, msg: str): + self.msg = msg +``` + +### Output Example +```python{example}``` +''' + +# https://www.sphinx-doc.org/en/master/usage/extensions/napoleon.html + +PYTHON_DOCSTRING_EXAMPLE_GOOGLE = ''' +def function_with_pep484_type_annotations(param1: int) -> bool: + """Example function with PEP 484 type annotations. + + Extended description of function. + + Args: + param1: The first parameter. + + Returns: + The return value. True for success, False otherwise. + """ + ... + +class ExampleError(Exception): + """Exceptions are documented in the same way as classes. + + The __init__ method was documented in the class level docstring. + + Args: + msg: Human readable string describing the exception. + + Attributes: + msg: Human readable string describing the exception. + """ + ... +''' + +PYTHON_DOCSTRING_EXAMPLE_NUMPY = ''' +def function_with_pep484_type_annotations(param1: int) -> bool: + """ + Example function with PEP 484 type annotations. + + Extended description of function. + + Parameters + ---------- + param1 + The first parameter. + + Returns + ------- + bool + The return value. True for success, False otherwise. + """ + ... + +class ExampleError(Exception): + """ + Exceptions are documented in the same way as classes. + + The __init__ method was documented in the class level docstring. + + Parameters + ---------- + msg + Human readable string describing the exception. + + Attributes + ---------- + msg + Human readable string describing the exception. + """ + ... +''' + +PYTHON_DOCSTRING_EXAMPLE_SPHINX = ''' +def function_with_pep484_type_annotations(param1: int) -> bool: + """Example function with PEP 484 type annotations. + + Extended description of function. + + :param param1: The first parameter. + :type param1: int + + :return: The return value. True for success, False otherwise. + :rtype: bool + """ + ... + +class ExampleError(Exception): + """Exceptions are documented in the same way as classes. + + The __init__ method was documented in the class level docstring. + + :param msg: Human-readable string describing the exception. + :type msg: str + """ + ... +''' + +_python_docstring_style = { + "google": PYTHON_DOCSTRING_EXAMPLE_GOOGLE, + "numpy": PYTHON_DOCSTRING_EXAMPLE_NUMPY, + "sphinx": PYTHON_DOCSTRING_EXAMPLE_SPHINX, +} + + +class WriteDocstring(Action): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.desc = "Write docstring for code." + + async def run( + self, code: str, + system_text: str = PYTHON_DOCSTRING_SYSTEM, + style: Literal["google", "numpy", "sphinx"] = "google", + ) -> str: + system_text = system_text.format(style=style, example=_python_docstring_style[style]) + simplified_code = _simplify_python_code(code) + documented_code = await self._aask(simplified_code, [system_text]) + with contextlib.suppress(Exception): + documented_code = OutputParser.parse_code(documented_code) + return merge_docstring(code, documented_code) + + +def _simplify_python_code(code: str) -> None: + code_tree = ast.parse(code) + code_tree.body = [i for i in code_tree.body if not isinstance(i, ast.Expr)] + if isinstance(code_tree.body[-1], ast.If): + code_tree.body.pop() + return ast.unparse(code_tree) + + +if __name__ == "__main__": + import fire + + async def run(filename: str, overwrite: bool = False, style: Literal["google", "numpy", "sphinx"] = "google"): + with open(filename) as f: + code = f.read() + code = await WriteDocstring().run(code, style=style) + if overwrite: + with open(filename, "w") as f: + f.write(code) + return code + + fire.Fire(run) diff --git a/metagpt/utils/pycst.py b/metagpt/utils/pycst.py new file mode 100644 index 000000000..c2eb532ab --- /dev/null +++ b/metagpt/utils/pycst.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +from typing import Union + +import libcst as cst +from libcst._nodes.module import Module + +DocstringNode = Union[cst.Module, cst.ClassDef, cst.FunctionDef] + + +def get_docstring_statement(body: DocstringNode) -> cst.SimpleStatementLine: + """Extracts the docstring from the body of a node. + + Args: + body: The body of a node. + + Returns: + The docstring statement if it exists, None otherwise. + """ + if isinstance(body, cst.Module): + body = body.body + else: + body = body.body.body + + if not body: + return + + statement = body[0] + if not isinstance(statement, cst.SimpleStatementLine): + return + + expr = statement + while isinstance(expr, (cst.BaseSuite, cst.SimpleStatementLine)): + if len(expr.body) == 0: + return None + expr = expr.body[0] + + if not isinstance(expr, cst.Expr): + return None + + val = expr.value + if not isinstance(val, (cst.SimpleString, cst.ConcatenatedString)): + return None + + evaluated_value = val.evaluated_value + if isinstance(evaluated_value, bytes): + return None + + return statement + + +class DocstringCollector(cst.CSTVisitor): + """A visitor class for collecting docstrings from a CST. + + Attributes: + stack: A list to keep track of the current path in the CST. + docstrings: A dictionary mapping paths in the CST to their corresponding docstrings. + """ + def __init__(self): + self.stack: list[str] = [] + self.docstrings: dict[tuple[str, ...], cst.SimpleStatementLine] = {} + + def visit_Module(self, node: cst.Module) -> bool | None: + self.stack.append("") + + def leave_Module(self, node: cst.Module) -> None: + return self._leave(node) + + def visit_ClassDef(self, node: cst.ClassDef) -> bool | None: + self.stack.append(node.name.value) + + def leave_ClassDef(self, node: cst.ClassDef) -> None: + return self._leave(node) + + def visit_FunctionDef(self, node: cst.FunctionDef) -> bool | None: + self.stack.append(node.name.value) + + def leave_FunctionDef(self, node: cst.FunctionDef) -> None: + return self._leave(node) + + def _leave(self, node: DocstringNode) -> None: + key = tuple(self.stack) + self.stack.pop() + if hasattr(node, "decorators") and any(i.decorator.value == "overload" for i in node.decorators): + return + + statement = get_docstring_statement(node) + if statement: + self.docstrings[key] = statement + + +class DocstringTransformer(cst.CSTTransformer): + """A transformer class for replacing docstrings in a CST. + + Attributes: + stack: A list to keep track of the current path in the CST. + docstrings: A dictionary mapping paths in the CST to their corresponding docstrings. + """ + def __init__( + self, + docstrings: dict[tuple[str, ...], cst.SimpleStatementLine], + ): + self.stack: list[str] = [] + self.docstrings = docstrings + + def visit_Module(self, node: cst.Module) -> bool | None: + self.stack.append("") + + def leave_Module(self, original_node: Module, updated_node: Module) -> Module: + return self._leave(original_node, updated_node) + + def visit_ClassDef(self, node: cst.ClassDef) -> bool | None: + self.stack.append(node.name.value) + + def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.CSTNode: + return self._leave(original_node, updated_node) + + def visit_FunctionDef(self, node: cst.FunctionDef) -> bool | None: + self.stack.append(node.name.value) + + def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.CSTNode: + return self._leave(original_node, updated_node) + + def _leave(self, original_node: DocstringNode, updated_node: DocstringNode) -> DocstringNode: + key = tuple(self.stack) + self.stack.pop() + + if hasattr(updated_node, "decorators") and any((i.decorator.value == "overload") for i in updated_node.decorators): + return updated_node + + statement = self.docstrings.get(key) + if not statement: + return updated_node + + original_statement = get_docstring_statement(original_node) + + if isinstance(updated_node, cst.Module): + body = updated_node.body + if original_statement: + return updated_node.with_changes(body=(body[0], statement, *body[1:])) + else: + updated_node = updated_node.with_changes(body=(statement, cst.EmptyLine(), *body)) + return updated_node + + body = updated_node.body.body + if original_statement: + return updated_node.with_changes(body=updated_node.body.with_changes(body=(body[0], statement, *body[1:]))) + else: + return updated_node.with_changes(body=updated_node.body.with_changes(body=(statement, *body))) + + +def merge_docstring(code: str, documented_code: str) -> str: + """Merges the docstrings from the documented code into the original code. + + Args: + code: The original code. + documented_code: The documented code. + + Returns: + The original code with the docstrings from the documented code. + """ + code_tree = cst.parse_module(code) + documented_code_tree = cst.parse_module(documented_code) + + visitor = DocstringCollector() + documented_code_tree.visit(visitor) + transformer = DocstringTransformer(visitor.docstrings) + modified_tree = code_tree.visit(transformer) + return modified_tree.code diff --git a/requirements.txt b/requirements.txt index 32a436962..452e2d092 100644 --- a/requirements.txt +++ b/requirements.txt @@ -35,3 +35,4 @@ tqdm==4.64.0 anthropic==0.3.6 typing-inspect==0.8.0 typing_extensions==4.5.0 +libcst==1.0.1 diff --git a/tests/metagpt/actions/test_write_docstring.py b/tests/metagpt/actions/test_write_docstring.py new file mode 100644 index 000000000..82d96e1a6 --- /dev/null +++ b/tests/metagpt/actions/test_write_docstring.py @@ -0,0 +1,32 @@ +import pytest + +from metagpt.actions.write_docstring import WriteDocstring + +code = ''' +def add_numbers(a: int, b: int): + return a + b + + +class Person: + def __init__(self, name: str, age: int): + self.name = name + self.age = age + + def greet(self): + return f"Hello, my name is {self.name} and I am {self.age} years old." +''' + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("style", "part"), + [ + ("google", "Args:"), + ("numpy", "Parameters"), + ("sphinx", ":param name:"), + ], + ids=["google", "numpy", "sphinx"] +) +async def test_write_docstring(style: str, part: str): + ret = await WriteDocstring().run(code, style=style) + assert part in ret diff --git a/tests/metagpt/utils/test_pycst.py b/tests/metagpt/utils/test_pycst.py new file mode 100644 index 000000000..07352eac2 --- /dev/null +++ b/tests/metagpt/utils/test_pycst.py @@ -0,0 +1,136 @@ +from metagpt.utils import pycst + +code = ''' +#!/usr/bin/env python +# -*- coding: utf-8 -*- +from typing import overload + +@overload +def add_numbers(a: int, b: int): + ... + +@overload +def add_numbers(a: float, b: float): + ... + +def add_numbers(a: int, b: int): + return a + b + + +class Person: + def __init__(self, name: str, age: int): + self.name = name + self.age = age + + def greet(self): + return f"Hello, my name is {self.name} and I am {self.age} years old." +''' + +documented_code = ''' +""" +This is an example module containing a function and a class definition. +""" + + +def add_numbers(a: int, b: int): + """This function is used to add two numbers and return the result. + + Parameters: + a: The first integer. + b: The second integer. + + Returns: + int: The sum of the two numbers. + """ + return a + b + +class Person: + """This class represents a person's information, including name and age. + + Attributes: + name: The person's name. + age: The person's age. + """ + + def __init__(self, name: str, age: int): + """Creates a new instance of the Person class. + + Parameters: + name: The person's name. + age: The person's age. + """ + ... + + def greet(self): + """ + Returns a greeting message including the name and age. + + Returns: + str: The greeting message. + """ + ... +''' + + +merged_code = ''' +#!/usr/bin/env python +# -*- coding: utf-8 -*- +""" +This is an example module containing a function and a class definition. +""" + +from typing import overload + +@overload +def add_numbers(a: int, b: int): + ... + +@overload +def add_numbers(a: float, b: float): + ... + +def add_numbers(a: int, b: int): + """This function is used to add two numbers and return the result. + + Parameters: + a: The first integer. + b: The second integer. + + Returns: + int: The sum of the two numbers. + """ + return a + b + + +class Person: + """This class represents a person's information, including name and age. + + Attributes: + name: The person's name. + age: The person's age. + """ + def __init__(self, name: str, age: int): + """Creates a new instance of the Person class. + + Parameters: + name: The person's name. + age: The person's age. + """ + self.name = name + self.age = age + + def greet(self): + """ + Returns a greeting message including the name and age. + + Returns: + str: The greeting message. + """ + return f"Hello, my name is {self.name} and I am {self.age} years old." +''' + + +def test_merge_docstring(): + data = pycst.merge_docstring(code, documented_code) + print(data) + assert data == merged_code From 6c0531347f90c1e4d2680c2f628fb3629c6bbf32 Mon Sep 17 00:00:00 2001 From: shenchucheng Date: Sun, 6 Aug 2023 19:45:06 +0800 Subject: [PATCH 2/5] add docs for metagpt/actions/write_docstring.py --- metagpt/actions/write_docstring.py | 32 +++++++++++++++++++++++++----- metagpt/utils/common.py | 5 +++-- metagpt/utils/pycst.py | 9 +++------ 3 files changed, 33 insertions(+), 13 deletions(-) diff --git a/metagpt/actions/write_docstring.py b/metagpt/actions/write_docstring.py index 091bd0a82..370d020de 100644 --- a/metagpt/actions/write_docstring.py +++ b/metagpt/actions/write_docstring.py @@ -8,10 +8,9 @@ from metagpt.utils.pycst import merge_docstring PYTHON_DOCSTRING_SYSTEM = '''### Requirements 1. Add docstrings to the given code following the {style} style. -2. Remove all private members whose names start with an underscore, such as `_test` and `__init__`. -3. Replace the function body with an Ellipsis object(...) to reduce output. -4. If the types are already annotated, there is no need to include them in the docstring. -5. Only output Python code and avoid including any other text. +2. Replace the function body with an Ellipsis object(...) to reduce output. +3. If the types are already annotated, there is no need to include them in the docstring. +4. Extract only class, function or the docstrings for the module parts from the given Python code, avoiding any other text. ### Input Example ```python @@ -128,6 +127,11 @@ _python_docstring_style = { class WriteDocstring(Action): + """This class is used to write docstrings for code. + + Attributes: + desc: A string describing the action. + """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -138,15 +142,33 @@ class WriteDocstring(Action): system_text: str = PYTHON_DOCSTRING_SYSTEM, style: Literal["google", "numpy", "sphinx"] = "google", ) -> str: + """Writes docstrings for the given code and system text in the specified style. + + Args: + code: A string of Python code. + system_text: A string of system text. + style: A string specifying the style of the docstring. Can be 'google', 'numpy', or 'sphinx'. + + Returns: + The Python code with docstrings added. + """ system_text = system_text.format(style=style, example=_python_docstring_style[style]) simplified_code = _simplify_python_code(code) - documented_code = await self._aask(simplified_code, [system_text]) + documented_code = await self._aask(f"```python\n{simplified_code}\n```", [system_text]) with contextlib.suppress(Exception): documented_code = OutputParser.parse_code(documented_code) return merge_docstring(code, documented_code) def _simplify_python_code(code: str) -> None: + """Simplifies the given Python code by removing expressions and the last if statement. + + Args: + code: A string of Python code. + + Returns: + The simplified Python code. + """ code_tree = ast.parse(code) code_tree.body = [i for i in code_tree.body if not isinstance(i, ast.Expr)] if isinstance(code_tree.body[-1], ast.If): diff --git a/metagpt/utils/common.py b/metagpt/utils/common.py index 0cd73ec0b..d695fd4c4 100644 --- a/metagpt/utils/common.py +++ b/metagpt/utils/common.py @@ -49,7 +49,7 @@ class OutputParser: @classmethod def parse_code(cls, text: str, lang: str = "") -> str: - pattern = rf'```{lang}.*?\s+(.*?)```' + pattern = rf'```{lang}.*?\s+(.*)```' match = re.search(pattern, text, re.DOTALL) if match: code = match.group(1) @@ -231,7 +231,8 @@ def print_members(module, indent=0): elif inspect.ismethod(obj): print(f'{prefix}Method: {name}') + def parse_recipient(text): - pattern = "## Send To:\s*([A-Za-z]+)\s*?" # hard code for now + pattern = r"## Send To:\s*([A-Za-z]+)\s*?" # hard code for now recipient = re.search(pattern, text) return recipient.group(1) if recipient else "" diff --git a/metagpt/utils/pycst.py b/metagpt/utils/pycst.py index c2eb532ab..afd85a547 100644 --- a/metagpt/utils/pycst.py +++ b/metagpt/utils/pycst.py @@ -137,16 +137,13 @@ class DocstringTransformer(cst.CSTTransformer): if isinstance(updated_node, cst.Module): body = updated_node.body if original_statement: - return updated_node.with_changes(body=(body[0], statement, *body[1:])) + return updated_node.with_changes(body=(statement, *body[1:])) else: updated_node = updated_node.with_changes(body=(statement, cst.EmptyLine(), *body)) return updated_node - body = updated_node.body.body - if original_statement: - return updated_node.with_changes(body=updated_node.body.with_changes(body=(body[0], statement, *body[1:]))) - else: - return updated_node.with_changes(body=updated_node.body.with_changes(body=(statement, *body))) + body = updated_node.body.body[1:] if original_statement else updated_node.body.body + return updated_node.with_changes(body=updated_node.body.with_changes(body=(statement, *body))) def merge_docstring(code: str, documented_code: str) -> str: From cfb26155c6d7589f4ddcf2a6d407044f75fa15df Mon Sep 17 00:00:00 2001 From: shenchucheng <60704484+shenchucheng@users.noreply.github.com> Date: Mon, 7 Aug 2023 01:29:25 +0800 Subject: [PATCH 3/5] Update write_docstring.py to fix wrong word --- metagpt/actions/write_docstring.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/metagpt/actions/write_docstring.py b/metagpt/actions/write_docstring.py index 370d020de..cc4d7d904 100644 --- a/metagpt/actions/write_docstring.py +++ b/metagpt/actions/write_docstring.py @@ -15,7 +15,7 @@ PYTHON_DOCSTRING_SYSTEM = '''### Requirements ### Input Example ```python def function_with_pep484_type_annotations(param1: int) -> bool: - return isinstanc(param1, int) + return isinstance(param1, int) class ExampleError(Exception): def __init__(self, msg: str): From 50cdba30525518e142f5f7d0686ac5fc2617cb65 Mon Sep 17 00:00:00 2001 From: shenchucheng <60704484+shenchucheng@users.noreply.github.com> Date: Mon, 7 Aug 2023 10:16:00 +0800 Subject: [PATCH 4/5] Update write_docstring.py to make the template of the prompt more clear --- metagpt/actions/write_docstring.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/metagpt/actions/write_docstring.py b/metagpt/actions/write_docstring.py index cc4d7d904..db1928872 100644 --- a/metagpt/actions/write_docstring.py +++ b/metagpt/actions/write_docstring.py @@ -23,7 +23,9 @@ class ExampleError(Exception): ``` ### Output Example -```python{example}``` +```python +{example} +``` ''' # https://www.sphinx-doc.org/en/master/usage/extensions/napoleon.html @@ -120,9 +122,9 @@ class ExampleError(Exception): ''' _python_docstring_style = { - "google": PYTHON_DOCSTRING_EXAMPLE_GOOGLE, - "numpy": PYTHON_DOCSTRING_EXAMPLE_NUMPY, - "sphinx": PYTHON_DOCSTRING_EXAMPLE_SPHINX, + "google": PYTHON_DOCSTRING_EXAMPLE_GOOGLE.strip(), + "numpy": PYTHON_DOCSTRING_EXAMPLE_NUMPY.strip(), + "sphinx": PYTHON_DOCSTRING_EXAMPLE_SPHINX.strip(), } From 8b1fff22155fb0844376e76f1eaa55d175ae015e Mon Sep 17 00:00:00 2001 From: shenchucheng Date: Mon, 7 Aug 2023 20:58:01 +0800 Subject: [PATCH 5/5] add docs for write_docstring.py and parse python code with retry --- metagpt/actions/write_docstring.py | 27 ++++++++++++++++++++--- metagpt/utils/common.py | 20 ++++++++++++++++- tests/metagpt/utils/test_output_parser.py | 18 ++++++++++++++- 3 files changed, 60 insertions(+), 5 deletions(-) diff --git a/metagpt/actions/write_docstring.py b/metagpt/actions/write_docstring.py index db1928872..5c7815793 100644 --- a/metagpt/actions/write_docstring.py +++ b/metagpt/actions/write_docstring.py @@ -1,5 +1,27 @@ +"""Code Docstring Generator. + +This script provides a tool to automatically generate docstrings for Python code. It uses the specified style to create +docstrings for the given code and system text. + +Usage: + python3 -m metagpt.actions.write_docstring [--overwrite] [--style=] + +Arguments: + filename The path to the Python file for which you want to generate docstrings. + +Options: + --overwrite If specified, overwrite the original file with the code containing docstrings. + --style= Specify the style of the generated docstrings. + Valid values: 'google', 'numpy', or 'sphinx'. + Default: 'google' + +Example: + python3 -m metagpt.actions.write_docstring startup.py --overwrite False --style=numpy + +This script uses the 'fire' library to create a command-line interface. It generates docstrings for the given Python code using +the specified docstring style and adds them to the code. +""" import ast -import contextlib from typing import Literal from metagpt.actions.action import Action @@ -157,8 +179,7 @@ class WriteDocstring(Action): system_text = system_text.format(style=style, example=_python_docstring_style[style]) simplified_code = _simplify_python_code(code) documented_code = await self._aask(f"```python\n{simplified_code}\n```", [system_text]) - with contextlib.suppress(Exception): - documented_code = OutputParser.parse_code(documented_code) + documented_code = OutputParser.parse_python_code(documented_code) return merge_docstring(code, documented_code) diff --git a/metagpt/utils/common.py b/metagpt/utils/common.py index d695fd4c4..7f090cf63 100644 --- a/metagpt/utils/common.py +++ b/metagpt/utils/common.py @@ -6,6 +6,7 @@ @File : common.py """ import ast +import contextlib import inspect import os import re @@ -49,7 +50,7 @@ class OutputParser: @classmethod def parse_code(cls, text: str, lang: str = "") -> str: - pattern = rf'```{lang}.*?\s+(.*)```' + pattern = rf'```{lang}.*?\s+(.*?)```' match = re.search(pattern, text, re.DOTALL) if match: code = match.group(1) @@ -78,6 +79,23 @@ class OutputParser: else: tasks = text.split("\n") return tasks + + @staticmethod + def parse_python_code(text: str) -> str: + for pattern in ( + r'(.*?```python.*?\s+)?(?P.*)(```.*?)', + r'(.*?```python.*?\s+)?(?P.*)', + ): + match = re.search(pattern, text, re.DOTALL) + if not match: + continue + code = match.group("code") + if not code: + continue + with contextlib.suppress(Exception): + ast.parse(code) + return code + raise ValueError("Invalid python code") @classmethod def parse_data(cls, data): diff --git a/tests/metagpt/utils/test_output_parser.py b/tests/metagpt/utils/test_output_parser.py index 155297860..c56cff6fa 100644 --- a/tests/metagpt/utils/test_output_parser.py +++ b/tests/metagpt/utils/test_output_parser.py @@ -19,7 +19,7 @@ def test_parse_blocks(): def test_parse_code(): - test_text = "```python\nprint('Hello, world!')\n```" + test_text = "```python\nprint('Hello, world!')```" expected_result = "print('Hello, world!')" assert OutputParser.parse_code(test_text, 'python') == expected_result @@ -27,6 +27,22 @@ def test_parse_code(): OutputParser.parse_code(test_text, 'java') +def test_parse_python_code(): + expected_result = "print('Hello, world!')" + assert OutputParser.parse_python_code("```python\nprint('Hello, world!')```") == expected_result + assert OutputParser.parse_python_code("```python\nprint('Hello, world!')") == expected_result + assert OutputParser.parse_python_code("print('Hello, world!')") == expected_result + assert OutputParser.parse_python_code("print('Hello, world!')```") == expected_result + assert OutputParser.parse_python_code("print('Hello, world!')```") == expected_result + expected_result = "print('```Hello, world!```')" + assert OutputParser.parse_python_code("```python\nprint('```Hello, world!```')```") == expected_result + assert OutputParser.parse_python_code("The code is: ```python\nprint('```Hello, world!```')```") == expected_result + assert OutputParser.parse_python_code("xxx.\n```python\nprint('```Hello, world!```')```\nxxx") == expected_result + + with pytest.raises(ValueError): + OutputParser.parse_python_code("xxx =") + + def test_parse_str(): test_text = "name = 'Alice'" expected_result = 'Alice'