diff --git a/src/dokyu/parser/docstring/arg.py b/src/dokyu/parser/docstring/arg.py index dcee2e6..9723806 100644 --- a/src/dokyu/parser/docstring/arg.py +++ b/src/dokyu/parser/docstring/arg.py @@ -10,12 +10,22 @@ def to_argument(arg: list[Any]) -> Parser[str, Argument]: + is_optional = False name, type_, description_ = arg - type = type_[0] if type_ else None + + if type_: + type = type_[0].strip() + assert isinstance(type, str) + if type.endswith("optional"): + is_optional = True + type = type[: type.rfind(",")].strip() + else: + type = None + description_first_line = description_[0] description_rest_lines = description_[1] description = " ".join([description_first_line] + description_rest_lines) - return success(Argument(name=name, type=type, description=description)) + return success(Argument(name=name, type=type, is_optional=is_optional, description=description)) def to_arguments(args: list[Any]) -> Parser[str, list[Argument]]: @@ -25,7 +35,7 @@ def to_arguments(args: list[Any]) -> Parser[str, list[Argument]]: class ArgsParser(ParserContext, whitespace=None): next_line = "\n" & indent blank_line = "\n" - args_head = lit("Args:") + args_head = lit("Args:") | lit("Parameters:") arg_type = reg(r"[^)]+") arg_doc = optional_whitespace >> reg(r".*") & rep("\n" >> double_indent >> reg(r".*")) # TODO(SigureMo): Deal *args and **kwargs diff --git a/src/dokyu/schema/docstring.py b/src/dokyu/schema/docstring.py index 211ced4..9d8fd04 100644 --- a/src/dokyu/schema/docstring.py +++ b/src/dokyu/schema/docstring.py @@ -6,4 +6,5 @@ class Argument(BaseModel): name: str type: Optional[str] + is_optional: bool description: str diff --git a/tests/test_parser/test_docstring/test_arg.py b/tests/test_parser/test_docstring/test_arg.py index 0c873e8..44ce71f 100644 --- a/tests/test_parser/test_docstring/test_arg.py +++ b/tests/test_parser/test_docstring/test_arg.py @@ -42,10 +42,11 @@ def test_parse_arg_type(arg_type_string: str): arg_2 (list[int|tuple[int, ...]] | tuple[int, ...]): Arg 2 is a list of integers. """, [ - Argument(name="arg_1", type="int", description="Arg 1 is an integer."), + Argument(name="arg_1", type="int", is_optional=False, description="Arg 1 is an integer."), Argument( name="arg_2", type="list[int|tuple[int, ...]] | tuple[int, ...]", + is_optional=False, description="Arg 2 is a list of integers.", ), ], @@ -57,7 +58,7 @@ def test_parse_arg_type(arg_type_string: str): arg_1 (int): Arg 1 is an integer. """, [ - Argument(name="arg_1", type="int", description="Arg 1 is an integer."), + Argument(name="arg_1", type="int", is_optional=False, description="Arg 1 is an integer."), ], ), ( @@ -71,6 +72,7 @@ def test_parse_arg_type(arg_type_string: str): Argument( name="arg_1", type="np.ndarray[int] | int", + is_optional=False, description="Arg 1 is an integer. This line is too long and will be wrapped with next line. and next line.", ), ], @@ -89,15 +91,26 @@ def test_parse_arg_type(arg_type_string: str): Argument( name="arg_1", type="np.ndarray[int] | int", + is_optional=False, description="Arg 1 is an integer. This line is too long and will be wrapped with next line. and next line.", ), Argument( name="arg_2", type="np.ndarray[int] | int", + is_optional=False, description="Arg 1 is an integer. This line is too long and will be wrapped with next line. and next line.", ), ], ), + ( + """\ +Parameters: + arg_1 (int, optional): Arg 1 is an integer. +""", + [ + Argument(name="arg_1", type="int", is_optional=True, description="Arg 1 is an integer."), + ], + ), ], ) def test_parse_args(args_string: str, expected: list[Argument]):