Commit 1dab3c5
authored
Preserve docstrings in the new parser (#21890)
Old parser has some special treatment for docstrings. Match that in the
new parser. Not adding tests, since some existing tests fail when
switched to new parser, see #218231 parent cea9058 commit 1dab3c5
1 file changed
Lines changed: 12 additions & 0 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -723,6 +723,8 @@ def read_func_def(state: State, data: ReadBuffer) -> FuncDef: | |||
| 723 | 723 | # TODO: This seems wasteful, can we avoid it? | |
| 724 | 724 | func_def.unanalyzed_type = typ.copy_modified() | |
| 725 | 725 | expect_end_tag(data) | |
| 726 | + if state.options.include_docstrings: | ||
| 727 | + func_def.docstring = get_docstring(body) | ||
| 726 | 728 | return func_def | |
| 727 | 729 | ||
| 728 | 730 | ||
@@ -769,9 +771,19 @@ def read_class_def(state: State, data: ReadBuffer) -> ClassDef: | |||
| 769 | 771 | ) | |
| 770 | 772 | check_type_param_defaults(state, type_params, class_def.line, class_def.column) | |
| 771 | 773 | expect_end_tag(data) | |
| 774 | + if state.options.include_docstrings: | ||
| 775 | + class_def.docstring = get_docstring(body) | ||
| 772 | 776 | return class_def | |
| 773 | 777 | ||
| 774 | 778 | ||
| 779 | + def get_docstring(body: Block) -> str | None: | ||
| 780 | + if body.body and isinstance(body.body[0], ExpressionStmt): | ||
| 781 | + expr = body.body[0].expr | ||
| 782 | + if isinstance(expr, StrExpr): | ||
| 783 | + return expr.value | ||
| 784 | + return None | ||
| 785 | + | ||
| 786 | + | ||
| 775 | 787 | def read_type_alias_stmt(state: State, data: ReadBuffer) -> TypeAliasStmt: | |
| 776 | 788 | """Read PEP 695 type alias statement.""" | |
| 777 | 789 | name = read_expression(state, data) | |
0 commit comments