def reindent_docstring(node, indent_level=1):
"""
Reindent the docstring
:param node: AST node
:type node: ```ast.AST```
:param indent_level: docstring indentation level whence: 0=no_tabs, 1=one tab; 2=two tabs
:type indent_level: ```int```
:returns: Node with reindent docstring
:rtype: ```ast.AST```
"""
doc_str = ast.get_docstring(node)
if doc_str is not None:
node.body[0] = ast.Expr(
set_value("\n{tab}{s}\n{tab}".format(
tab=tab * abs(indent_level),
s="\n".join(
map(
lambda line: "{sep}{line}".format(sep=tab * 2,
line=line.lstrip())
if line.startswith(tab)
and len(line) > len(tab) and line[len(tab):line.lstrip(
).find(" ") + len(tab)].rstrip(":s") not in frozenset(
(False, ) + TOKENS.rest) else line,
reindent(doc_str).splitlines(),
)),
)))
return node
def test_to_function_emit_as_kwonlyargs(self) -> None:
"""
Tests whether `function` produces method with keyword only arguments
"""
function_def = deepcopy(
next(
filter(
rpartial(isinstance, FunctionDef),
ast.parse(class_with_method_types_str.replace("self", "self, *"))
.body[0]
.body,
)
)
)
# Reindent docstring
function_def.body[0].value.value = "\n{tab}{docstring}\n{tab}".format(
tab=tab,
docstring=reindent(
deindent(ast.get_docstring(function_def)),
),
)
function_name = function_def.name
function_type = get_function_type(function_def)
gen_ast = emit.function(
parse.docstring(docstring_str),
function_name=function_name,
function_type=function_type,
emit_default_doc=False,
inline_types=True,
emit_separating_tab=PY3_8,
indent_level=0,
emit_as_kwonlyargs=True,
)
gen_ast.body[0].value.value = "\n{tab}{docstring}".format(
tab=tab, docstring=reindent(deindent(ast.get_docstring(gen_ast)))
)
run_ast_test(
self,
gen_ast=gen_ast,
gold=function_def,
)
def test_to_function(self) -> None:
"""
Tests whether `function` produces method from `class_with_method_types_ast` given `docstring_str`
"""
function_def = deepcopy(
next(
filter(
rpartial(isinstance, FunctionDef), class_with_method_types_ast.body
)
)
)
# Reindent docstring
function_def.body[0].value.value = "\n{tab}{docstring}\n{tab}".format(
tab=tab,
docstring=reindent(
deindent(ast.get_docstring(function_def).translate("\n\t")),
),
)
function_name = function_def.name
function_type = get_function_type(function_def)
gen_ast = emit.function(
parse.docstring(docstring_str),
function_name=function_name,
function_type=function_type,
emit_default_doc=False,
inline_types=True,
emit_separating_tab=PY3_8,
indent_level=0,
emit_as_kwonlyargs=False,
)
gen_ast.body[0].value.value = "\n{tab}{docstring}".format(
tab=tab, docstring=reindent(deindent(ast.get_docstring(gen_ast)))
)
run_ast_test(
self,
gen_ast=gen_ast,
gold=function_def,
)
def test_to_function_with_docstring_types(self) -> None:
"""
Tests that `function` can generate a function_def with types in docstring
"""
# Sanity check
run_ast_test(
self,
class_with_method_ast,
gold=ast.parse(class_with_method_str).body[0],
)
function_def = deepcopy(
next(filter(rpartial(isinstance, FunctionDef), class_with_method_ast.body))
)
# Reindent docstring
function_def.body[0].value.value = "\n{tab}{docstring}\n{tab}".format(
tab=tab, docstring=reindent(ast.get_docstring(function_def))
)
ir = parse.function(function_def)
gen_ast = emit.function(
ir,
function_name=function_def.name,
function_type=get_function_type(function_def),
emit_default_doc=False,
inline_types=False,
indent_level=1,
emit_separating_tab=True,
emit_as_kwonlyargs=False,
)
gen_ast.body[0].value.value = "\n{tab}{docstring}\n{tab}".format(
tab=tab, docstring=reindent(ast.get_docstring(gen_ast))
)
run_ast_test(
self,
gen_ast=gen_ast,
gold=function_def,
)
def ast_parse(
source,
filename="<unknown>",
mode="exec",
skip_annotate=False,
skip_docstring_remit=False,
):
"""
Convert the AST input to Python source string
:param source: Python source
:type source: ```str```
:param filename: Filename being parsed
:type filename: ```str```
:param mode: 'exec' to compile a module, 'single' to compile a, single (interactive) statement,
or 'eval' to compile an expression.
:type mode: ```Literal['exec', 'single', 'eval']```
:param skip_annotate: Don't run `annotate_ancestry`
:type skip_annotate: ```bool```
:param skip_docstring_remit: Don't parse & emit the docstring as a replacement for current docstring
:type skip_docstring_remit: ```bool```
:returns: AST node
:rtype: node: ```AST```
"""
parsed_ast = parse(source, filename=filename, mode=mode)
if not skip_annotate:
annotate_ancestry(parsed_ast)
if not skip_docstring_remit and isinstance(
parsed_ast, (Module, ClassDef, FunctionDef, AsyncFunctionDef)):
docstring = get_docstring(parsed_ast)
if docstring is None:
return parsed_ast
# Reindent docstring
parsed_ast.body[0].value.value = "\n{tab}{docstring}\n{tab}".format(
tab=tab, docstring=reindent(docstring))
return parsed_ast
def test_from_class_with_body_in_method_to_method_with_body(self) -> None:
""" Tests if this can make the roundtrip from a full function to a full function """
annotate_ancestry(class_with_method_and_body_types_ast)
function_def = next(
filter(
rpartial(isinstance, FunctionDef),
class_with_method_and_body_types_ast.body,
)
)
# Reindent docstring
function_def.body[0].value.value = "\n{tab}{docstring}\n{tab}".format(
tab=tab, docstring=reindent(ast.get_docstring(function_def))
)
ir = parse.function(
find_in_ast(
"C.function_name".split("."),
class_with_method_and_body_types_ast,
),
)
gen_ast = emit.function(
ir,
emit_default_doc=False,
function_name="function_name",
function_type="self",
indent_level=1,
emit_separating_tab=True,
emit_as_kwonlyargs=False,
)
# emit.file(gen_ast, os.path.join(os.path.dirname(__file__), "delme.py"), mode="wt")
run_ast_test(
self,
gen_ast=gen_ast,
gold=function_def,
)