def test_annotate_ancestry(self) -> None:
""" Tests that `annotate_ancestry` properly decorates """
node = Module(
body=[
AnnAssign(
annotation=Name(
"str",
Load(),
),
simple=1,
target=Name("dataset_name", Store()),
value=set_value("~/tensorflow_datasets"),
expr=None,
expr_target=None,
expr_annotation=None,
),
Assign(annotation=None,
simple=1,
targets=[Name("epochs", Store())],
value=set_value("333"),
expr=None,
expr_target=None,
expr_annotation=None,
**maybe_type_comment),
],
stmt=None,
)
self.assertFalse(hasattr(node.body[0], "_location"))
self.assertFalse(hasattr(node.body[1], "_location"))
annotate_ancestry(node)
self.assertEqual(node.body[0]._location, ["dataset_name"])
self.assertEqual(node.body[1]._location, ["epochs"])
def test_param2argparse_param_default_list(self) -> None:
"""
Tests that param2argparse_param works to change the type based on the default
whence said default is a list
"""
run_ast_test(
gen_ast=param2argparse_param(("byo", {
"default": [],
"typ": "str"
}), ),
gold=Expr(
Call(
args=[set_value("--byo")],
func=Attribute(
Name("argument_parser", Load()),
"add_argument",
Load(),
),
keywords=[
keyword(arg="action",
value=set_value("append"),
identifier=None),
keyword(arg="required",
value=set_value(True),
identifier=None),
],
expr=None,
expr_func=None,
), ),
test_case_instance=self,
)
def test_parse_out_param_fails(self) -> None:
""" Test that parse_out_param throws NotImplementedError when unsupported type given """
self.assertRaises(
NotImplementedError,
lambda: parse_out_param(
Expr(
Call(
args=[set_value("--num")],
func=Attribute(
Name("argument_parser", Load()),
"add_argument",
Load(),
),
keywords=[
keyword(
arg="type",
value=Subscript(
expr_context_ctx=None,
expr_slice=None,
expr_value=None,
),
identifier=None,
),
keyword(
arg="required",
value=set_value(True),
identifier=None,
),
],
expr=None,
expr_func=None,
))),
)
def test_param2argparse_param_none_default(self) -> None:
"""
Tests that param2argparse_param works to reparse the default
"""
run_ast_test(
gen_ast=param2argparse_param(("yup", {
"default": NoneStr
})),
gold=Expr(
Call(
args=[set_value("--yup")],
func=Attribute(
Name("argument_parser", Load()),
"add_argument",
Load(),
),
keywords=[
keyword(arg="type",
value=Name("str", Load()),
identifier=None),
keyword(arg="default",
value=set_value(None),
identifier=None),
],
expr=None,
expr_func=None,
)),
test_case_instance=self,
)
def test_parse_to_scalar(self) -> None:
""" Test various inputs and outputs for `parse_to_scalar` """
for fst, snd in (
(5, 5),
("5", "5"),
(set_value(5), 5),
(ast.Expr(None), NoneStr),
):
self.assertEqual(parse_to_scalar(fst), snd)
self.assertEqual(
get_value(parse_to_scalar(ast.parse("[5]").body[0]).elts[0]), 5
)
self.assertTrue(
cmp_ast(
parse_to_scalar(ast.parse("[5]").body[0]),
List([set_value(5)], Load()),
)
)
self.assertEqual(parse_to_scalar(ast.parse("[5]")), "[5]")
parse_to_scalar(ast.parse("[5]").body[0])
self.assertRaises(NotImplementedError, parse_to_scalar, memoryview(b""))
self.assertRaises(NotImplementedError, parse_to_scalar, memoryview(b""))
def test_param2argparse_param_default_function(self) -> None:
"""
Tests that param2argparse_param works to change the type based on the default
whence said default is an in-memory function
"""
function_str = (
"from operator import add\n"
"def adder(a, b):\n"
"{tab}return add(a, b)".format(tab=tab)
)
adder = getattr(
inspectable_compile(function_str),
"adder",
)
pickled_adder = pickle.dumps(adder) # eww
run_ast_test(
gen_ast=param2argparse_param(
(
"byo",
{
"default": adder,
"typ": "str",
},
),
),
gold=Expr(
Call(
args=[set_value("--byo")],
func=Attribute(
Name("argument_parser", Load()),
"add_argument",
Load(),
),
keywords=[
keyword(
arg="type",
value=Name("pickle.loads", Load()),
identifier=None,
),
keyword(
arg="default",
value=set_value(pickled_adder),
identifier=None,
),
],
expr=None,
expr_func=None,
)
),
test_case_instance=self,
)
def test_param2argparse_param_default_torch(self) -> None:
"""
Tests that param2argparse_param works to change the type based on the default
whence said default is a proxy for an internal PyTorch type
"""
class FakeTorch(object):
""" Not a real torch """
def __str__(self):
"""But a real str
:returns: An actual str
:rtype: ```Literal['<required parameter>']```
"""
return "<required parameter>"
# type("FakeTorch", tuple(), {"__str__": lambda _: "<required parameter>"})
run_ast_test(
gen_ast=param2argparse_param(
(
"byo",
{
"default": FakeTorch(),
},
),
),
gold=Expr(
Call(
args=[set_value("--byo")],
func=Attribute(
Name("argument_parser", Load()),
"add_argument",
Load(),
),
keywords=[
keyword(
arg="type",
value=Name(FakeTorch.__name__, Load()),
identifier=None,
),
keyword(arg="required", value=set_value(True), identifier=None),
],
expr=None,
expr_func=None,
)
),
test_case_instance=self,
)
def test_replace_in_ast_with_val_on_non_function(self) -> None:
"""
Tests that `RewriteAtQuery` can actually replace a node at given location
"""
parsed_ast = ast_parse(class_str)
rewrite_at_query = RewriteAtQuery(
search="ConfigClass.dataset_name".split("."),
replacement_node=AnnAssign(
annotation=Name("int", Load()),
simple=1,
target=Name("dataset_name", Store()),
value=set_value(15),
expr=None,
expr_target=None,
expr_annotation=None,
),
)
gen_ast = rewrite_at_query.visit(parsed_ast)
self.assertTrue(rewrite_at_query.replaced, True)
run_ast_test(
self,
gen_ast,
ast.parse(
class_str.replace('dataset_name: str = "mnist"',
"dataset_name: int = 15")),
)
def test_emit_ann_assign(self) -> None:
""" Tests that AnnAssign is emitted from `emit_ann_assign` """
self.assertIsInstance(class_ast.body[1], AnnAssign)
self.assertIsInstance(emit_ann_assign(class_ast.body[1]), AnnAssign)
self.assertIsInstance(emit_ann_assign(class_ast.body[1]), AnnAssign)
gen_ast = emit_ann_assign(
find_in_ast(
"C.function_name.dataset_name".split("."),
class_with_method_and_body_types_ast,
))
self.assertIsInstance(gen_ast, AnnAssign)
run_ast_test(
self,
gen_ast,
AnnAssign(
annotation=Name(
"str",
Load(),
),
simple=1,
target=Name("dataset_name", Store()),
value=set_value("~/tensorflow_datasets"),
expr=None,
expr_target=None,
expr_annotation=None,
),
)
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__set_name_and_type(self) -> None:
"""
Tests that `_set_name_and_type` parsed AST code into a code str.
Not working since I explicitly deleted the typ from ``` quoted defaults. Changed mock to match.
"""
self.assertTupleEqual(
_set_name_and_type(
(
"adder",
{
"default": BinOp(
set_value(5),
Mult(),
set_value(5),
),
},
),
infer_type=True,
word_wrap=True,
),
("adder", {
"default": "```(5 * 5)```"
}),
)
self.assertTupleEqual(
_set_name_and_type(
(
"adder",
{
"default": BinOp(
set_value(5),
Mult(),
set_value(5),
),
"doc": ["5", "b"],
},
),
infer_type=True,
word_wrap=True,
),
("adder", {
"default": "```(5 * 5)```",
"doc": "5b"
}),
)
def test_param2ast_with_wrapped_default(self) -> None:
""" Check that `param2ast` behaves correctly with a wrapped default """
run_ast_test(
self,
param2ast(
("zion", {"typ": None, "default": set_value(NoneStr)}),
),
gold=AnnAssign(
annotation=Name("object", Load()),
simple=1,
target=Name("zion", Store()),
value=set_value(None),
expr=None,
expr_target=None,
expr_annotation=None,
),
)
def _make_call_meth(body, return_type, param_names):
"""
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]]```
:return: Internal function for `__call__`
:rtype: ```FunctionDef```
"""
body_len = len(body)
return_ = (ast.fix_missing_locations(
RewriteName(param_names).visit(
Return(get_value(ast.parse(return_type[3:-3]).body[0]),
expr=None))) if return_type is not None
and len(return_type) > 6 and return_type.startswith("```")
and return_type.endswith("```") else None)
if body_len:
if isinstance(body[0], Expr):
doc_str = get_value(body[0].value)
if isinstance(doc_str, str) and body_len > 0:
body = (body[1:] if body_len > 1 else ([
set_value(doc_str.replace(":cvar", ":param"))
if return_ is None else return_
] if body_len == 1 else body))
# elif not isinstance(body[0], Return) and return_ is not None:
# body.append(return_)
# elif return_ is not None:
# body = [return_]
return 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)
def test_param2argparse_param_default_ast_binop(self) -> None:
"""
Tests that param2argparse_param works to change the type based on the default
whence said default is a non specially handled ast.AST
"""
run_ast_test(
gen_ast=param2argparse_param(
(
"byo",
{
"default": BinOp(
set_value(5),
Mult(),
set_value(5),
),
"typ": "str",
},
),
),
gold=Expr(
Call(
args=[set_value("--byo")],
func=Attribute(
Name("argument_parser", Load()),
"add_argument",
Load(),
),
keywords=[
keyword(arg="required", value=set_value(True), identifier=None),
keyword(
arg="default",
value=set_value("```(5 * 5)```"),
identifier=None,
),
],
expr=None,
expr_func=None,
)
),
test_case_instance=self,
)
def test_param2argparse_param_default_ast_tuple(self) -> None:
"""
Tests that param2argparse_param works to change the type based on the default
whence said default is an ast.Tuple
"""
run_ast_test(
gen_ast=param2argparse_param((
"byo",
{
"default": Tuple(
elts=[],
ctx=Load(),
expr=None,
),
"typ": "str",
},
), ),
gold=Expr(
Call(
args=[set_value("--byo")],
func=Attribute(
Name("argument_parser", Load()),
"add_argument",
Load(),
),
keywords=[
keyword(arg="type",
value=Name("loads", Load()),
identifier=None),
keyword(arg="required",
value=set_value(True),
identifier=None),
keyword(arg="default",
value=set_value("()"),
identifier=None),
],
expr=None,
expr_func=None,
)),
test_case_instance=self,
)
def _merge_name_to_column(assign):
"""
Merge `a = Column()` into `Column("a")`
:param assign: Of form `a = Column()`
:type assign: ```Assign```
:returns: Unwrapped Call with name prepended
:rtype: ```Call```
"""
assign.value.args.insert(0, set_value(assign.targets[0].id))
return assign.value
def test_gen_with_imports_from_file_and_prepended_import(self) -> None:
""" Tests `gen` with `imports_from_file` and `prepend` """
output_filename = os.path.join(
self.tempdir,
"test_gen_with_imports_from_file_and_prepended_import_output.py",
)
with patch("sys.stdout", new_callable=StringIO), patch(
"sys.stderr", new_callable=StringIO
):
self.assertIsNone(
gen(
name_tpl="{name}Config",
input_mapping="gen_test_module.input_map",
imports_from_file="gen_test_module",
type_="class",
prepend=_import_gen_test_module_str,
output_filename=output_filename,
emit_call=True,
emit_default_doc=False,
)
)
with open(output_filename, "rt") as f:
gen_ast = ast.parse(f.read())
gold = Module(
body=[
_import_gen_test_module_ast,
_import_star_from_input_ast,
self.expected_class_ast,
# self.input_module_ast.body[1],
Assign(
targets=[Name("__all__", Store())],
value=List(
ctx=Load(),
elts=[set_value("FooConfig")],
expr=None,
),
expr=None,
lineno=None,
**maybe_type_comment
),
],
type_ignores=[],
stmt=None,
)
run_ast_test(
self,
gen_ast=gen_ast,
gold=gold,
)
def test_param2argparse_param_default_code_quoted(self) -> None:
"""
Tests that param2argparse_param works to change the type based on the default
whence said default is a code quoted str
"""
run_ast_test(
gen_ast=(param2argparse_param((
"byo",
{
"default": "```(4)```",
"typ": "str",
},
), )),
gold=Expr(
Call(
args=[set_value("--byo")],
func=Attribute(
Name("argument_parser", Load()),
"add_argument",
Load(),
),
keywords=[
keyword(arg="type",
value=Name("int", Load()),
identifier=None),
keyword(arg="required",
value=set_value(True),
identifier=None),
keyword(arg="default",
value=set_value(4),
identifier=None),
],
expr=None,
expr_func=None,
)),
test_case_instance=self,
)
def test_param2ast_with_assign(self) -> None:
""" Check that `param2ast` behaves correctly with a non annotated (typeless) input """
run_ast_test(
self,
param2ast(("zion", {
"typ": None
}), ),
gold=AnnAssign(
annotation=Name("object", Load()),
simple=1,
target=Name("zion", Store()),
value=set_value(None),
expr=None,
expr_target=None,
expr_annotation=None,
),
)
def test_find_in_ast(self) -> None:
""" Tests that `find_in_ast` successfully finds nodes in AST """
run_ast_test(
self,
find_in_ast("ConfigClass.dataset_name".split("."), class_ast),
AnnAssign(
annotation=Name(
"str",
Load(),
),
simple=1,
target=Name("dataset_name", Store()),
value=set_value("mnist"),
expr=None,
expr_target=None,
expr_annotation=None,
),
)
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_param2ast_with_bad_default(self) -> None:
""" Check that `param2ast` behaves correctly with a bad default """
run_ast_test(
self,
param2ast(
(
"stateful_metrics",
{"typ": "NoneType", "default": "the `Model`'s metrics"},
),
),
gold=AnnAssign(
annotation=Name("NoneType", Load()),
simple=1,
target=Name("stateful_metrics", Store()),
value=set_value("```the `Model`'s metrics```"),
expr=None,
expr_annotation=None,
expr_target=None,
),
)
def sqlalchemy(class_def):
"""
Parse out a `class C(Base): __tablename__= 'tbl'; dataset_name = Column(String, doc="p", primary_key=True)`,
as constructed on an SQLalchemy declarative `Base`.
:param class_def: A class inheriting from declarative `Base`, where `Base = sqlalchemy.orm.declarative_base()`
:type class_def: ```Union[ClassDef]```
:returns: 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}),)]] }
:rtype: ```dict```
"""
assert isinstance(class_def, ClassDef)
# Parse into the same format that `sqlalchemy_table` can read, then return with a call to it
name = get_value(
next(
filter(
lambda assign: any(
filter(
partial(eq, "__tablename__"),
map(attrgetter("id"), assign.targets),
)),
filter(rpartial(isinstance, Assign), class_def.body),
)).value)
doc_string = get_docstring(class_def)
def _merge_name_to_column(assign):
"""
Merge `a = Column()` into `Column("a")`
:param assign: Of form `a = Column()`
:type assign: ```Assign```
:returns: Unwrapped Call with name prepended
:rtype: ```Call```
"""
assign.value.args.insert(0, set_value(assign.targets[0].id))
return assign.value
return sqlalchemy_table(
Call(
func=Name("Table", Load()),
args=list(
chain.from_iterable((
iter((set_value(name), Name("metadata", Load()))),
map(
_merge_name_to_column,
filterfalse(
lambda assign: any(
map(
lambda target: target.id == "__tablename__"
or hasattr(target, "value") and isinstance(
target.value, Call) and target.func.
rpartition(".")[2] == "Column",
assign.targets,
), ),
filter(rpartial(isinstance, Assign),
class_def.body),
),
),
))),
keywords=[] if doc_string is None else [
keyword(arg="comment",
value=set_value(doc_string),
identifier=None)
],
expr=None,
expr_func=None,
))
Name,
Return,
Store,
Tuple,
arguments,
fix_missing_locations,
keyword,
)
from doctrans.ast_utils import FALLBACK_TYP, maybe_type_comment, set_arg, set_value
from doctrans.pure_utils import tab
from doctrans.tests.mocks.docstrings import docstring_header_str
argparse_add_argument_ast = Expr(
Call(
args=[set_value("--num")],
func=Attribute(
Name("argument_parser", Load()),
"add_argument",
Load(),
),
keywords=[
keyword(arg="type", value=Name("int", Load()), identifier=None),
keyword(
arg="required",
value=set_value(True),
identifier=None,
),
],
expr=None,
expr_func=None,
"data_loader_kwargs",
JSON,
doc="pass this as arguments to data_loader function",
default=None,
nullable=True,
),
comment={comment!r},
)
""".format(comment=docstring_header_and_return_str)
config_tbl_ast = Assign(
targets=[Name("config_tbl", Store())],
value=Call(
func=Name("Table", Load()),
args=[
set_value("config_tbl"),
Name("metadata", Load()),
Call(
func=Name("Column", Load()),
args=[set_value("dataset_name"),
Name("String", Load())],
keywords=[
keyword(arg="doc",
value=set_value("name of dataset"),
identifier=None),
keyword(arg="default",
value=set_value("mnist"),
identifier=None),
keyword(arg="primary_key",
value=set_value(True),
identifier=None),
def argparse_function(
intermediate_repr,
emit_default_doc=False,
emit_default_doc_in_return=False,
function_name="set_cli_args",
function_type="static",
):
"""
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 emit_default_doc_in_return: Whether help/docstring in return should include 'With default' text
:type emit_default_doc_in_return: ```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']```
:return: 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(
"\n Set CLI arguments\n\n "
":param argument_parser: argument parser\n "
":type argument_parser: ```ArgumentParser```\n\n "
"{return_params}".format(return_params=(
lambda returns:
":return: argument_parser{return_doc}\n "
"{rtype}".format(
return_doc=", {}".format(returns["doc"])
if "doc" in returns else "",
rtype=":rtype: ```ArgumentParser```\n "
if intermediate_repr["returns"][
"return_type"].get("typ") in
(None, "None", NoneStr) else
":rtype: ```Tuple[ArgumentParser, {returns[typ]}]```\n{tab}"
"".format(returns=returns, tab=tab),
)
)(returns=set_default_doc(
next(iter(intermediate_repr["returns"].items())
),
emit_default_doc=emit_default_doc_in_return,
)[1]) if "return_type" in (
intermediate_repr.get("returns") or iter(())
) else ":return: argument_parser\n "
":rtype: ```ArgumentParser```\n "
))),
Assign(targets=[
Attribute(
Name("argument_parser", Load()),
"description",
Store(),
)
],
value=set_value(intermediate_repr["doc"]),
lineno=None,
expr=None,
**maybe_type_comment),
)),
filter(
None,
(
*((map(
partial(
param2argparse_param,
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
intermediate_repr["returns"]["return_type"]
["default"].startswith("```")
and intermediate_repr["returns"]
["return_type"]["default"].endswith("```")
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,
emit_default_doc=False,
docstring_format="rest",
indent_level=2,
emit_separating_tab=PY3_8,
inline_types=True,
emit_as_kwonlyargs=True,
):
"""
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', 'numpy', 'google']```
: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 inline_types: Whether the type should be inline or in docstring
:type inline_types: ```bool```
:param emit_as_kwonlyargs: Whether argument(s) emitted must be keyword only
:type emit_as_kwonlyargs: ```bool```
:return: 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 doctrans.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 inline_types 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, "None", NoneStr) 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(
to_docstring(
intermediate_repr,
emit_default_doc=emit_default_doc,
docstring_format=docstring_format,
emit_types=not inline_types,
indent_level=indent_level,
emit_separating_tab=emit_separating_tab,
))),
*(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 inline_types 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)
def class_(
intermediate_repr,
emit_call=False,
class_name="ConfigClass",
class_bases=("object", ),
decorator_list=None,
emit_default_doc=False,
):
"""
Construct a class
: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_call: Whether to emit a `__call__` method from the `_internal` IR subdict
:type emit_call: ```bool```
:param class_name: name of class
:type class_name: ```str```
:param class_bases: bases of class (the generated class will inherit these)
:type class_bases: ```Iterable[str]```
:param decorator_list: List of decorators
:type decorator_list: ```Optional[Union[List[Str], List[]]]```
:param emit_default_doc: Whether help/docstring should include 'With default' text
:type emit_default_doc: ```bool```
:return: Class AST of the docstring
:rtype: ```ClassDef```
"""
returns = (intermediate_repr["returns"] if "return_type"
in intermediate_repr.get("returns", {}) else OrderedDict())
param_names = frozenset(intermediate_repr["params"].keys())
if returns:
intermediate_repr["params"].update(returns)
del intermediate_repr["returns"]
internal_body = intermediate_repr.get("_internal", {}).get("body", [])
# TODO: Add correct classmethod/staticmethod to decorate function using `annotate_ancestry` and first-field checks
# Such that the `self.` or `cls.` rewrite only applies to non-staticmethods
# assert internal_body, "Expected `internal_body` to have contents"
if internal_body and param_names:
internal_body = list(
map(
ast.fix_missing_locations,
map(RewriteName(param_names).visit, internal_body),
))
return ClassDef(
bases=list(map(rpartial(Name, Load()), class_bases)),
body=list(
chain.from_iterable((
(Expr(
set_value(
to_docstring(
intermediate_repr,
indent_level=0,
emit_separating_tab=False,
emit_default_doc=emit_default_doc,
emit_types=False,
).replace("\n:param ",
"{tab}:cvar ".format(tab=tab)).replace(
"{tab}:cvar ".format(tab=tab),
"\n{tab}:cvar ".format(tab=tab),
1,
).rstrip())), ),
map(param2ast, intermediate_repr["params"].items()),
iter((_make_call_meth(
internal_body,
returns["return_type"]["default"] if "default" in (
(returns or {
"return_type": iter(())
}).get("return_type") or iter(())) else None,
param_names,
), ) if emit_call and internal_body else tuple()),
))),
decorator_list=list(map(rpartial(Name, Load()), decorator_list))
if decorator_list else [],
keywords=[],
name=class_name,
expr=None,
identifier_name=None,
)
def param_to_sqlalchemy_column_call(param, include_name):
"""
Turn a param into a `Column(…)`
:param param: Name, dict with keys: 'typ', 'doc', 'default'
:type param: ```Tuple[str, dict]```
:param include_name: Whether to include the name (exclude in declarative base)
:type include_name: ```bool```
:returns: Form of: `Column(…)`
:rtype: ```Call```
"""
print("param_to_sqlalchemy_column_call::include_name:", include_name, ";")
name, _param = param
del param
args, keywords, nullable = [], [], None
if _param["typ"].startswith("Optional["):
_param["typ"] = _param["typ"][len("Optional["):-1]
nullable = True
if include_name:
args.append(set_value(name))
if "Literal[" in _param["typ"]:
parsed_typ = get_value(ast.parse(_param["typ"]).body[0])
assert (parsed_typ.value.id == "Literal"
), "Only basic Literal support is implemented, not {}".format(
parsed_typ.value.id)
args.append(
Call(
func=Name("Enum", Load()),
args=get_value(parsed_typ.slice).elts,
keywords=[
keyword(arg="name", value=set_value(name), identifier=None)
],
expr=None,
expr_func=None,
))
else:
args.append(Name(typ2column_type[_param["typ"]], Load()))
has_default = _param.get("default", ast) is not ast
pk = _param.get("doc", "").startswith("[PK]")
if pk:
_param["doc"] = _param["doc"][4:].lstrip()
elif has_default and _param["default"] not in none_types:
nullable = False
keywords.append(
keyword(arg="doc",
value=set_value(_param["doc"].rstrip(".")),
identifier=None))
if has_default:
if _param["default"] == NoneStr:
_param["default"] = None
keywords.append(
keyword(
arg="default",
value=set_value(_param["default"]),
identifier=None,
))
# Sorting :\
if pk:
keywords.append(
keyword(arg="primary_key", value=set_value(True),
identifier=None), )
if isinstance(nullable, bool):
keywords.append(
keyword(arg="nullable", value=set_value(nullable),
identifier=None))
return Call(
func=Name("Column", Load()),
args=args,
keywords=keywords,
expr=None,
expr_func=None,
)
def generate_repr_method(params, cls_name, docstring_format):
"""
Generate a `__repr__` method with all params, using `str.format` syntax
:param params: an `OrderedDict` of form
OrderedDict[str, {'typ': str, 'doc': Optional[str], 'default': Any}]
:type params: ```OrderedDict```
:param cls_name: Name of class
:type cls_name: ```str```
:param docstring_format: Format of docstring
:type docstring_format: ```Literal['rest', 'numpydoc', 'google']```
:returns: `__repr__` method
:rtype: ```FunctionDef```
"""
keys = tuple(params.keys())
return FunctionDef(
name="__repr__",
args=arguments(
posonlyargs=[],
arg=None,
args=[
arg(arg="self",
annotation=None,
expr=None,
identifier_arg=None,
**maybe_type_comment)
],
kwonlyargs=[],
kw_defaults=[],
defaults=[],
vararg=None,
kwarg=None,
),
body=[
Expr(
set_value(docstring_repr_str if docstring_format ==
"rest" else docstring_repr_google_str)),
Return(
value=Call(
func=Attribute(
set_value("{cls_name}({format_args})".format(
cls_name=cls_name,
format_args=", ".join(
map("{0}={{{0}!r}}".format, keys)),
)),
"format",
Load(),
),
args=[],
keywords=list(
map(
lambda key: keyword(
arg=key,
value=Attribute(Name("self", Load()), key,
Load()),
identifier=None,
),
keys,
)),
expr=None,
expr_func=None,
),
expr=None,
),
],
decorator_list=[],
arguments_args=None,
identifier_name=None,
stmt=None,
lineno=None,
returns=None,
**maybe_type_comment)