def test_get_function_type(self) -> None:
"""Test get_function_type returns the right type"""
self.assertEqual(
get_function_type(
FunctionDef(
args=arguments(
args=[set_arg("something else")],
arg=None,
),
arguments_args=None,
identifier_name=None,
stmt=None,
)
),
"static",
)
self.assertEqual(
get_function_type(
FunctionDef(
args=arguments(args=[], arg=None),
arguments_args=None,
identifier_name=None,
stmt=None,
)
),
"static",
)
self.assertEqual(
get_function_type(
FunctionDef(
args=arguments(
args=[set_arg("self")],
arg=None,
),
arguments_args=None,
identifier_name=None,
stmt=None,
)
),
"self",
)
self.assertEqual(
get_function_type(
FunctionDef(
args=arguments(
args=[set_arg("cls")],
arg=None,
),
arguments_args=None,
identifier_name=None,
stmt=None,
)
),
"cls",
)
def _handle_function(self, node, doc_str):
"""
Handle functions
:param node: AsyncFunctionDef | FunctionDef
:type node: ```Union[AsyncFunctionDef, FunctionDef]```
:param doc_str: The docstring
:type doc_str: ```Optional[str]```
:returns: Same type as input with args, returns, and docstring potentially modified
:rtype: ```Union[AsyncFunctionDef, FunctionDef]```
"""
ir = parse_docstring(doc_str)
ir_merge(ir, parse.function(node))
ir["name"] = node.name
indent_level = max(
len(node._location) - 1,
1) # function docstrings always have at least 1 indent level
_doc_str = emit.docstring(
ir,
emit_types=not self.type_annotations,
emit_default_doc=False,
docstring_format=self.docstring_format,
indent_level=indent_level,
)
if _doc_str.isspace():
if doc_str is not None:
del node.body[0]
else:
set_docstring(_doc_str, False, node)
if self.type_annotations:
# Add annotations
if ir["params"]:
node.args.args = list(
map(
lambda _arg: set_arg(
_arg.arg,
annotation=to_annotation(
_arg.annotation
if _arg.annotation is not None else ir[
"params"][_arg.arg].get("typ")),
),
node.args.args,
))
if ("return_type" in (ir.get("returns") or iter(()))
and ir["returns"]["return_type"].get("typ") is not None):
node.returns = to_annotation(
ir["returns"]["return_type"]["typ"])
else:
# Remove annotations
node.args.args = list(
map(set_arg, map(attrgetter("arg"), node.args.args)))
node.returns = None
node.body = list(map(self.visit, node.body))
return node
def test_node_to_dict(self) -> None:
"""
Tests `node_to_dict`
"""
self.assertDictEqual(
node_to_dict(set_arg(arg="a", annotation=Name("int", Load()))),
dict(
annotation="int",
arg="a",
**{"expr": None, "type_comment": None} if PY_GTE_3_8 else {}
),
)
def test_find_in_ast_no_val(self) -> None:
"""Tests that `find_in_ast` correctly gives AST node from
`def class C(object): def function_name(self,dataset_name: str,…)`"""
run_ast_test(
self,
find_in_ast(
"C.function_name.dataset_name".split("."),
class_with_optional_arg_method_ast,
),
set_arg(
annotation=Name(
"str",
Load(),
),
arg="dataset_name",
),
)
def test_emit_arg(self) -> None:
"""Tests that `arg` is emitted from `emit_arg`"""
self.assertIsInstance(
class_with_method_and_body_types_ast.body[1].args.args[1], arg
)
self.assertIsInstance(
emit_arg(class_with_method_and_body_types_ast.body[1].args.args[1]), arg
)
assign = Assign(
targets=[Name("yup", Store())],
value=set_value("nup"),
expr=None,
**maybe_type_comment
)
gen_ast = emit_arg(assign)
self.assertIsInstance(gen_ast, arg)
run_ast_test(
self,
gen_ast=gen_ast,
gold=set_arg("yup"),
)
def test_find_in_ast_with_val(self) -> None:
"""Tests that `find_in_ast` correctly gives AST node from
`def class C(object): def function_name(self,dataset_name: str='foo',…)`"""
gen_ast = find_in_ast(
"C.function_name.dataset_name".split("."),
class_with_method_and_body_types_ast,
)
self.assertIsInstance(gen_ast.default, Constant if PY_GTE_3_8 else Str)
self.assertEqual(get_value(gen_ast.default), "~/tensorflow_datasets")
run_ast_test(
self,
gen_ast,
set_arg(
annotation=Name(
"str",
Load(),
),
arg="dataset_name",
),
)
value=assign_with_type_comment.value,
simple=1,
target=assign_with_type_comment.targets[0],
type_comment=None,
expr=None,
expr_target=None,
expr_annotation=None,
lineno=None,
)
function_type_annotated = FunctionDef(
name="sum",
args=arguments(
posonlyargs=[],
args=[
set_arg(arg="a", annotation=Name("int", Load())),
set_arg(arg="b", annotation=Name("int", Load())),
],
kwonlyargs=[],
kw_defaults=[],
defaults=[],
vararg=None,
kwarg=None,
),
body=[
ann_assign_with_annotation,
Return(value=Name("res", Load())),
],
decorator_list=[],
lineno=None,
returns=Name("int", Load()),
def _make_call_meth(body, return_type, param_names, docstring_format,
word_wrap):
"""
Construct a `__call__` method from the provided `body`
:param body: The body, probably from a FunctionDef.body
:type body: ```List[AST]```
:param return_type: The return type of the parent symbol (probably class). Used to fill in `__call__` return.
:type return_type: ```Optional[str]```
:param param_names: Container of AST `id`s to match for rename
:type param_names: ```Optional[Iterator[str]]```
:param docstring_format: Format of docstring
:type docstring_format: ```Literal['rest', 'numpydoc', 'google']```
:param word_wrap: Whether to word-wrap. Set `DOCTRANS_LINE_LENGTH` to configure length.
:type word_wrap: ```bool```
:returns: Internal function for `__call__`
:rtype: ```FunctionDef```
"""
body_len = len(body)
if body_len:
if isinstance(body, dict):
body = list(
filter(
None,
(
None if body.get("doc") in none_types else Expr(
set_value(
emit_param_str(
(
"return_type",
{
"doc":
multiline(
indent_all_but_first(
body["doc"]))
},
),
style=docstring_format,
word_wrap=word_wrap,
))),
RewriteName(param_names).visit(
Return(
get_value(
ast.parse(return_type.strip("`")).body[0]),
expr=None,
)) if code_quoted(body["default"]) else Return(
set_value(body["default"]), expr=None),
),
))
return (ast.fix_missing_locations(
FunctionDef(args=arguments(
args=[set_arg("self")],
defaults=[],
kw_defaults=[],
kwarg=None,
kwonlyargs=[],
posonlyargs=[],
vararg=None,
arg=None,
),
body=body,
decorator_list=[],
name="__call__",
returns=None,
arguments_args=None,
identifier_name=None,
stmt=None,
lineno=None,
**maybe_type_comment)) if body else None)
def argparse_function(
intermediate_repr,
emit_default_doc=False,
function_name="set_cli_args",
function_type="static",
wrap_description=False,
word_wrap=True,
docstring_format="rest",
):
"""
Convert to an argparse FunctionDef
:param intermediate_repr: a dictionary of form
{ "name": Optional[str],
"type": Optional[str],
"doc": Optional[str],
"params": OrderedDict[str, {'typ': str, 'doc': Optional[str], 'default': Any}]
"returns": Optional[OrderedDict[Literal['return_type'],
{'typ': str, 'doc': Optional[str], 'default': Any}),)]] }
:type intermediate_repr: ```dict```
:param emit_default_doc: Whether help/docstring should include 'With default' text
:type emit_default_doc: ```bool```
:param function_name: name of function_def
:type function_name: ```str```
:param function_type: Type of function, static is static or global method, others just become first arg
:type function_type: ```Literal['self', 'cls', 'static']```
:param wrap_description: Whether to word-wrap the description. Set `DOCTRANS_LINE_LENGTH` to configure length.
:type wrap_description: ```bool```
:param word_wrap: Whether to word-wrap. Set `DOCTRANS_LINE_LENGTH` to configure length.
:type word_wrap: ```bool```
:param docstring_format: Format of docstring
:type docstring_format: ```Literal['rest', 'numpydoc', 'google']```
:returns: AST node for function definition which constructs argparse
:rtype: ```FunctionDef```
"""
function_name = function_name or intermediate_repr["name"]
function_type = function_type or intermediate_repr["type"]
internal_body = get_internal_body(
target_name=function_name,
target_type=function_type,
intermediate_repr=intermediate_repr,
)
return FunctionDef(
args=arguments(
args=[set_arg("argument_parser")],
# None if function_type in frozenset((None, "static"))
# else set_arg(function_type),
defaults=[],
kw_defaults=[],
kwarg=None,
kwonlyargs=[],
posonlyargs=[],
vararg=None,
arg=None,
),
body=list(
chain.from_iterable((
iter((
Expr(
set_value(
indent(
docstring(
{
"doc":
"Set CLI arguments",
"params":
OrderedDict(((
"argument_parser",
{
"doc": "argument parser",
"typ": "ArgumentParser",
},
), )),
"returns":
OrderedDict(((
"return_type",
{
"doc":
"argument_parser, {returns_doc}"
.format(
returns_doc=
intermediate_repr["returns"]
["return_type"]["doc"])
if intermediate_repr["returns"]
["return_type"].get("doc") else
"argument_parser",
"typ":
"Tuple[ArgumentParser, {typ}]".
format(typ=intermediate_repr[
"returns"]["return_type"]
["typ"]),
} if "return_type" in (
(intermediate_repr or {}
).get("returns") or iter(())
) and intermediate_repr["returns"]
["return_type"].get("typ") not in
none_types else {
"doc": "argument_parser",
"typ": "ArgumentParser",
},
), ), ),
},
docstring_format=docstring_format,
word_wrap=word_wrap,
),
tab,
) + tab)),
Assign(targets=[
Attribute(
Name("argument_parser", Load()),
"description",
Store(),
)
],
value=set_value(
(fill if wrap_description else identity)(
intermediate_repr["doc"])),
lineno=None,
expr=None,
**maybe_type_comment),
)),
filter(
None,
(
*((map(
partial(
param2argparse_param,
word_wrap=word_wrap,
emit_default_doc=emit_default_doc,
),
intermediate_repr["params"].items(),
)) if "params" in intermediate_repr else ()),
*(internal_body[
2 if len(internal_body) > 1 and isinstance(
internal_body[1], Assign) and internal_body[1].
targets[0].id == "argument_parser" else 1:] if
internal_body and isinstance(internal_body[0], Expr)
and isinstance(get_value(internal_body[0].value),
str) else internal_body),
None if internal_body
and isinstance(internal_body[-1], Return) else (Return(
value=Tuple(
ctx=Load(),
elts=[
Name("argument_parser", Load()),
set_value(intermediate_repr["returns"]
["return_type"]["default"])
if code_quoted(intermediate_repr["returns"]
["return_type"]["default"])
else ast.parse(intermediate_repr["returns"]
["return_type"]
["default"]).body[0].value,
],
expr=None,
),
expr=None,
) if "default" in (
intermediate_repr.get("returns") or {
"return_type": iter(())
})["return_type"] else Return(value=Name(
"argument_parser", Load()),
expr=None)),
),
),
))),
decorator_list=[],
name=function_name,
returns=None,
lineno=None,
arguments_args=None,
identifier_name=None,
stmt=None,
**maybe_type_comment)
def function(
intermediate_repr,
function_name,
function_type,
word_wrap=True,
emit_default_doc=False,
docstring_format="rest",
indent_level=2,
emit_separating_tab=PY3_8,
type_annotations=True,
emit_as_kwonlyargs=True,
emit_original_whitespace=False,
):
"""
Construct a function from our IR
:param intermediate_repr: a dictionary of form
{ "name": Optional[str],
"type": Optional[str],
"doc": Optional[str],
"params": OrderedDict[str, {'typ': str, 'doc': Optional[str], 'default': Any}]
"returns": Optional[OrderedDict[Literal['return_type'],
{'typ': str, 'doc': Optional[str], 'default': Any}),)]] }
:type intermediate_repr: ```dict```
:param function_name: name of function_def
:type function_name: ```Optional[str]```
:param function_type: Type of function, static is static or global method, others just become first arg
:type function_type: ```Optional[Literal['self', 'cls', 'static']]```
:param docstring_format: Format of docstring
:type docstring_format: ```Literal['rest', 'numpydoc', 'google']```
:param word_wrap: Whether to word-wrap. Set `DOCTRANS_LINE_LENGTH` to configure length.
:type word_wrap: ```bool```
:param emit_default_doc: Whether help/docstring should include 'With default' text
:type emit_default_doc: ```bool```
:param indent_level: docstring indentation level whence: 0=no_tabs, 1=one tab; 2=two tabs
:type indent_level: ```int```
:param emit_separating_tab: docstring decider for whether to put a tab between :param and return and desc
:type emit_separating_tab: ```bool```
:param type_annotations: True to have type annotations (3.6+), False to place in docstring
:type type_annotations: ```bool```
:param emit_as_kwonlyargs: Whether argument(s) emitted must be keyword only
:type emit_as_kwonlyargs: ```bool```
:param emit_original_whitespace: Whether to emit an original whitespace (in docstring) or strip it out
:type emit_original_whitespace: ```bool```
:returns: AST node for function definition
:rtype: ```FunctionDef```
"""
params_no_kwargs = tuple(
filter(
lambda param: not param[0].endswith("kwargs"),
intermediate_repr["params"].items(),
))
function_name = function_name or intermediate_repr["name"]
function_type = function_type or intermediate_repr["type"]
args = ([] if function_type in frozenset(
(None, "static")) else [set_arg(function_type)])
from cdd.emitter_utils import ast_parse_fix
args_from_params = list(
map(
lambda param: set_arg(
annotation=(Name(param[1]["typ"], Load()) if param[1]["typ"] in
simple_types else ast_parse_fix(param[1]["typ"]))
if type_annotations and "typ" in param[1] else None,
arg=param[0],
),
params_no_kwargs,
), )
defaults_from_params = list(
map(
lambda param: set_value(None) if param[1].get("default") in
none_types else set_value(param[1].get("default")),
params_no_kwargs,
))
if emit_as_kwonlyargs:
kwonlyargs, kw_defaults, defaults = args_from_params, defaults_from_params, []
else:
kwonlyargs, kw_defaults, defaults = [], [], defaults_from_params
args += args_from_params
internal_body = get_internal_body(
target_name=function_name,
target_type=function_type,
intermediate_repr=intermediate_repr,
)
return_val = (Return(
value=ast.parse(intermediate_repr["returns"]["return_type"]
["default"].strip("`")).body[0].value,
expr=None,
) if (intermediate_repr.get("returns") or {
"return_type": {}
})["return_type"].get("default") else None)
return FunctionDef(
args=arguments(
args=args,
defaults=defaults,
kw_defaults=kw_defaults,
kwarg=next(
map(
lambda param: set_arg(param[0]),
filter(
lambda param: param[0].endswith("kwargs"),
intermediate_repr["params"].items(),
),
),
None,
),
kwonlyargs=kwonlyargs,
posonlyargs=[],
vararg=None,
arg=None,
),
body=list(
filter(
None,
(
Expr(
set_value(
docstring(
intermediate_repr,
docstring_format=docstring_format,
emit_default_doc=emit_default_doc,
emit_original_whitespace=
emit_original_whitespace,
emit_separating_tab=emit_separating_tab,
emit_types=not type_annotations,
indent_level=indent_level,
word_wrap=word_wrap,
))),
*(internal_body[:-1] if internal_body
and isinstance(internal_body[-1], Return) and return_val
else internal_body),
return_val,
),
)),
decorator_list=[],
name=function_name,
returns=(ast.parse(
intermediate_repr["returns"]["return_type"]["typ"]).body[0].value
if type_annotations and (intermediate_repr.get("returns") or {
"return_type": {}
})["return_type"].get("typ") else None),
lineno=None,
arguments_args=None,
identifier_name=None,
stmt=None,
**maybe_type_comment)
"ModelCheckpoint",
"ProgbarLogger",
"ReduceLROnPlateau",
"RemoteMonitor",
"TensorBoard",
"TerminateOnNaN",
),
action="append",
help="Collection of callables that are run inside the training loop",
)
return argument_parser
'''.format(_cli_doc_str=_cli_doc_str, header_doc_str=docstring_header_str)
argparse_func_ast = fix_missing_locations(
FunctionDef(args=arguments(
args=[set_arg("argument_parser")],
defaults=[],
kw_defaults=[],
kwarg=None,
kwonlyargs=[],
posonlyargs=[],
vararg=None,
arg=None,
),
body=[
_cli_doc_expr,
Assign(targets=[
Attribute(
Name("argument_parser", Load()),
"description",
Store(),
)
class_with_method_and_body_types_ast = fix_missing_locations(
ClassDef(
name="C",
bases=[Name("object", Load())],
keywords=[],
body=[
Expr(set_value(" C class (mocked!) ")),
FunctionDef(
name="function_name",
args=arguments(
posonlyargs=[],
vararg=None,
args=[
set_arg("self"),
set_arg(arg="dataset_name",
annotation=Name("str", Load())),
set_arg(
arg="tfds_dir",
annotation=Subscript(
Name("Optional", Load()),
set_slice(Name("str", Load())),
Load(),
),
),
set_arg(
arg="K",
annotation=Subscript(
Name("Literal", Load()),
set_slice(
expr_annotation=None,
),
AnnAssign(
annotation=Name("str", Load()),
simple=1,
target=Name("return_type", Store()),
value=set_value(
"```K.mean(math_ops.square(math_ops.maximum(1.0 - y_true * y_pred, 0.0)), axis=-1)```"
),
expr=None,
expr_target=None,
expr_annotation=None,
),
FunctionDef(
args=arguments(
args=[set_arg("self")],
defaults=[],
kw_defaults=[],
kwarg=None,
kwonlyargs=[],
vararg=None,
arg=None,
posonlyargs=[],
),
body=[
Assign(
targets=[
Attribute(
Name("self", Load()),
"y_pred",
Load(),