def test_doctrans_function_from_docstring_to_annotated(self) -> None:
"""Tests `DocTrans` converts docstring function to type annotated function"""
original_node = annotate_ancestry(deepcopy(function_type_in_docstring))
doc_trans = DocTrans(
docstring_format="rest",
type_annotations=True,
existing_type_annotations=False,
whole_ast=original_node,
)
gen_ast = doc_trans.visit(original_node)
run_ast_test(self, gen_ast, gold=function_type_annotated)
def test_class_with_internal_converts_to_annotated(self) -> None:
"""Tests that class, function, and class variable hierarchy is correctly converts to annotated"""
original_node = annotate_ancestry(
deepcopy(class_with_internal_type_commented_and_docstring_typed))
doc_trans = DocTrans(
docstring_format="rest",
type_annotations=True,
existing_type_annotations=False,
whole_ast=original_node,
)
gen_ast = doc_trans.visit(original_node)
run_ast_test(self, gen_ast=gen_ast, gold=class_with_internal_annotated)
def test_doctrans_assign_to_assign(self) -> None:
"""
Tests that `AnnAssign` converts to `Assign`
"""
original_node = annotate_ancestry(deepcopy(assign_with_type_comment))
doc_trans = DocTrans(
docstring_format="rest",
type_annotations=False,
existing_type_annotations=False,
whole_ast=original_node,
)
gen_ast = doc_trans.visit(original_node)
run_ast_test(self, gen_ast=gen_ast, gold=assign_with_type_comment)
def test_module_docstring(self) -> None:
"""Tests that module gets the right new docstring"""
module_node = Module(body=[Expr(set_value("\nModule\n"))],
stmt=None,
type_ignores=[])
original = deepcopy(module_node)
doc_trans = DocTrans(
docstring_format="rest",
type_annotations=True,
existing_type_annotations=True,
whole_ast=module_node,
)
doc_trans.visit_Module(module_node)
run_ast_test(self, gen_ast=module_node, gold=original)
def test__get_ass_typ(self) -> None:
"""Tests that _get_ass_typ returns when location isn't set"""
original_node = annotate_ancestry(deepcopy(assign_with_type_comment))
doc_trans = DocTrans(
docstring_format="rest",
type_annotations=True,
existing_type_annotations=True,
whole_ast=original_node,
)
del original_node._location
run_ast_test(
self,
gen_ast=doc_trans._get_ass_typ(original_node),
gold=Name("int", Load()),
)
def doctrans(filename, docstring_format, type_annotations):
"""
Transform the docstrings found within provided filename to intended docstring_format
:param filename: Python file to convert docstrings within. Edited in place.
:type filename: ```str```
:param docstring_format: Format of docstring
:type docstring_format: ```Literal['rest', 'numpydoc', 'google']```
:param type_annotations: True to have type annotations (3.6+), False to place in docstring
:type type_annotations: ```bool```
"""
with open(filename, "rt") as f:
node = ast_parse(f.read(), skip_docstring_remit=False)
orig_node = deepcopy(node)
node = DocTrans(
docstring_format=docstring_format,
type_annotations=type_annotations,
existing_type_annotations=has_type_annotations(node),
whole_ast=orig_node,
).visit(node)
if not cmp_ast(node, orig_node):
emit.file(node, filename, mode="wt", skip_black=True)