def test_gen(self) -> None:
""" Tests `gen` """
output_filename = os.path.join(self.tempdir, "test_gen_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",
type_="class",
output_filename=output_filename,
prepend="PREPENDED\n",
emit_call=True,
emit_default_doc=False,
)
)
with open(output_filename, "rt") as f:
gen_module_str = f.read()
gen_module_ast = ast.parse(gen_module_str)
run_ast_test(
self,
gen_ast=next(filter(rpartial(isinstance, ClassDef), gen_module_ast.body)),
gold=self.expected_class_ast,
)
def test_to_function_with_inline_types(self) -> None:
"""
Tests that `function` can generate a function_def with inline types
"""
function_def = deepcopy(
next(
filter(
rpartial(isinstance, FunctionDef), class_with_method_types_ast.body
)
)
)
function_name = function_def.name
function_type = get_function_type(function_def)
gen_ast = emit.function(
intermediate_repr=parse.function(
function_def=function_def,
function_name=function_name,
function_type=function_type,
),
function_name=function_name,
function_type=function_type,
emit_default_doc=False,
inline_types=True,
emit_separating_tab=True,
indent_level=1,
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,
)
def test_diff(self) -> None:
""" Tests that `diff` gives correct results """
lstrip_l = partial(diff, op=str.lstrip)
self.assertTupleEqual(lstrip_l("A"), (0, "A"))
self.assertTupleEqual(lstrip_l(""), (0, ""))
self.assertTupleEqual(lstrip_l(" A"), (1, "A"))
self.assertTupleEqual(diff("AB", op=rpartial(str.lstrip, "A")),
(1, "B"))
def test_parse_out_param(self) -> None:
""" Test that parse_out_param parses out the right dict """
self.assertDictEqual(
parse_out_param(
next(
filter(rpartial(isinstance, Expr),
argparse_func_ast.body[::-1])))[1],
# Last element:
intermediate_repr["params"]["data_loader_kwargs"],
)
def _merge_inner_function(
class_def, infer_type, intermediate_repr, merge_inner_function
):
"""
Merge the inner function if found within the class, with the class IR
:param class_def: Class AST
:type class_def: ```ClassDef```
:param infer_type: Whether to try inferring the typ (from the default)
:type infer_type: ```bool```
: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 merge_inner_function: Name of inner function to merge. If None, merge nothing.
:type merge_inner_function: ```Optional[str]```
:return: 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```
"""
function_def = next(
filter(
lambda func: func.name == merge_inner_function,
filter(rpartial(isinstance, FunctionDef), ast.walk(class_def)),
),
None,
)
if function_def is not None:
function_type = (
"static" if not function_def.args.args else function_def.args.args[0].arg
)
inner_ir = function(
function_def,
function_name=merge_inner_function,
function_type=function_type,
infer_type=infer_type,
)
ir_merge(other=inner_ir, target=intermediate_repr)
return intermediate_repr
def get_at_root(node, types):
"""
Get the imports from a node
:param node: AST node with .body, probably an `ast.Module`
:type node: ```AST```
:param types: The types to search for (uses in an `isinstance` check)
:type types: ```Tuple[str,...]````
:return: List of imports. Doesn't handle those within a try/except, condition, or not in root scope
:rtype: ```List[Union[]]```
"""
assert hasattr(node, "body") and isinstance(node.body, (list, tuple))
return list(filter(rpartial(isinstance, types), node.body))
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 _parse(scan, partitioned=None):
"""
Parse the scanned input (Google)
:param scan: Scanned input
:type scan: ```List[str]```
:param partitioned: Prep-partitioned `scan`, if given doesn't partition on `scan`, just uses this
:type partitioned: ```Optional[Tuple[str, str, str]]```
:returns: dict of shape {'name': ..., 'typ': ..., 'doc': ..., 'default': ..., 'required': ... }
:rtype: ```dict```
"""
offset = next(idx for idx, ch in enumerate(scan[0]) if ch == ":")
s = scan[0][:offset].lstrip()
name, delim, typ = partitioned or s.partition("(")
name = name.rstrip()
typ = (delim + typ).rstrip()
# if not name: return None
cur = {"name": name}
if typ:
assert typ.startswith("(") and typ.endswith(
")"
), "Expected third partition to be paren wrapped {!r}".format(
s)
cur["typ"] = typ[1:-1]
if " or " in cur["typ"]:
cur["typ"] = "Union[{}]".format(", ".join(
cur["typ"].split(" or ")))
end = scan[0][offset + 1:].lstrip()
if len(end) > 3 and end.startswith("{") and end.endswith("}"):
# PyTorch invented their own syntax for this I guess?
cur["typ"] = "Literal{}".format(
list(
map(rpartial(str.strip, "'"),
end[1:-1].split(", "))))
scan[0] = ""
# elif partitioned is None:
# return _parse(scan, " ".join((name, typ)).partition("="))
# else:
# elif name.endswith("kwargs"): cur["typ"] = "dict"
cur["doc"] = "\n".join([scan[0][offset + 1:].lstrip()] +
scan[1:]).strip()
return cur
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 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,
)
def find_ast_type(node, node_name=None, of_type=ClassDef):
"""
Finds first AST node of the given type and possibly name
:param node: Any AST node
:type node: ```AST```
:param node_name: Name of AST node. If None, gives first found.
:type node_name: ```Optional[str]```
:param of_type: Of which type to find
:type of_type: ```AST```
:return: Found AST node
:rtype: ```AST```
"""
if isinstance(node, Module):
it = filter(rpartial(isinstance, of_type), node.body)
if node_name is not None:
return next(
filter(
lambda e: hasattr(e, "name") and e.name == node_name,
it,
)
)
matching_nodes = tuple(it)
if len(matching_nodes) > 1: # We could convert every one I guess?
raise NotImplementedError()
elif matching_nodes:
return matching_nodes[0]
else:
raise TypeError("No {!r} in AST".format(type(of_type).__name__))
elif isinstance(node, AST):
assert node_name is None or not hasattr(node, "name") or node.name == node_name
return node
else:
raise NotImplementedError(type(node).__name__)
def class_(
class_def,
class_name=None,
merge_inner_function=None,
infer_type=False,
word_wrap=True,
):
"""
Converts an AST to our IR
:param class_def: Class AST or Module AST with a ClassDef inside
:type class_def: ```Union[Module, ClassDef]```
:param class_name: Name of `class`. If None, gives first found.
:type class_name: ```Optional[str]```
:param merge_inner_function: Name of inner function to merge. If None, merge nothing.
:type merge_inner_function: ```Optional[str]```
:param infer_type: Whether to try inferring the typ (from the default)
:type infer_type: ```bool```
:param word_wrap: Whether to word-wrap. Set `DOCTRANS_LINE_LENGTH` to configure length.
:type word_wrap: ```bool```
: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 not isinstance(class_def, FunctionDef)
is_supported_ast_node = isinstance(class_def, (Module, ClassDef))
if not is_supported_ast_node and isinstance(class_def, type):
ir = _inspect(class_def, class_name, word_wrap)
parsed_body = ast.parse(getsource(class_def).lstrip()).body[0]
parsed_body.body = (parsed_body.body
if ast.get_docstring(parsed_body) is None else
parsed_body.body[1:])
if merge_inner_function is not None:
_merge_inner_function(
parsed_body,
infer_type=infer_type,
intermediate_repr=ir,
merge_inner_function=merge_inner_function,
)
return ir
ir["_internal"] = {
"body":
list(
filterfalse(
rpartial(isinstance, AnnAssign),
parsed_body.body,
)),
"from_name":
class_name,
"from_type":
"cls",
}
body_ir = class_(
class_def=parsed_body,
class_name=class_name,
merge_inner_function=merge_inner_function,
)
ir_merge(ir, body_ir)
return ir
assert (is_supported_ast_node
), "Expected 'Union[Module, ClassDef]' got `{!r}`".format(
type(class_def).__name__)
class_def = find_ast_type(class_def, class_name)
doc_str = get_docstring(class_def)
intermediate_repr = ({
"name": class_name,
"type": "static",
"doc": "",
"params": OrderedDict(),
"returns": None,
} if doc_str is None else docstring(get_docstring(class_def).replace(
":cvar", ":param"),
emit_default_doc=False))
if "return_type" in intermediate_repr["params"]:
intermediate_repr["returns"] = OrderedDict(
(("return_type",
intermediate_repr["params"].pop("return_type")), ))
body = class_def.body if doc_str is None else class_def.body[1:]
for e in body:
if isinstance(e, AnnAssign):
typ = to_code(e.annotation).rstrip("\n")
val = (lambda v: {
"default": NoneStr
} if v is None else {
"default":
v if type(v).__name__ in simple_types else (lambda value: {
"{}": {} if isinstance(v, Dict) else set(),
"[]": [],
"()": (),
}.get(value, parse_to_scalar(value)))(to_code(v).rstrip("\n"))
})(get_value(get_value(e)))
# if 'str' in typ and val: val["default"] = val["default"].strip("'") # Unquote?
typ_default = dict(typ=typ, **val)
for key in "params", "returns":
if e.target.id in (intermediate_repr[key] or iter(())):
intermediate_repr[key][e.target.id].update(typ_default)
typ_default = False
break
if typ_default:
k = "returns" if e.target.id == "return_type" else "params"
if intermediate_repr.get(k) is None:
intermediate_repr[k] = OrderedDict()
intermediate_repr[k][e.target.id] = typ_default
elif isinstance(e, Assign):
val = get_value(e)
if val is not None:
val = get_value(val)
deque(
map(
lambda target: setitem(*(
(intermediate_repr["params"][target.id], "default",
val)
if target.id in intermediate_repr["params"] else (
intermediate_repr["params"],
target.id,
{
"default": val
},
))),
e.targets,
),
maxlen=0,
)
intermediate_repr.update({
"params":
OrderedDict(
map(
partial(_set_name_and_type,
infer_type=infer_type,
word_wrap=word_wrap),
intermediate_repr["params"].items(),
)),
"_internal": {
"body":
list(filterfalse(rpartial(isinstance, (AnnAssign, Assign)), body)),
"from_name":
class_def.name,
"from_type":
"cls",
},
})
if merge_inner_function is not None:
assert isinstance(class_def, ClassDef)
_merge_inner_function(
class_def,
infer_type=infer_type,
intermediate_repr=intermediate_repr,
merge_inner_function=merge_inner_function,
)
# intermediate_repr['_internal']["body"]= list(filterfalse(rpartial(isinstance,(AnnAssign,Assign)),class_def.body))
return intermediate_repr
def function(
function_def,
infer_type=False,
word_wrap=True,
function_type=None,
function_name=None,
):
"""
Converts a method to our IR
:param function_def: AST node for function definition
:type function_def: ```Union[FunctionDef, FunctionType]```
:param infer_type: Whether to try inferring the typ (from the default)
:type infer_type: ```bool```
:param word_wrap: Whether to word-wrap. Set `DOCTRANS_LINE_LENGTH` to configure length.
:type word_wrap: ```bool```
: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 function_name: name of function_def
:type function_name: ```str```
: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```
"""
if isinstance(function_def, FunctionType):
# Dynamic function, i.e., this isn't source code; and is in your memory
ir = _inspect(function_def, function_name, word_wrap)
parsed_source = ast.parse(getsource(function_def).lstrip()).body[0]
body = (parsed_source.body if ast.get_docstring(parsed_source) is None
else parsed_source.body[1:])
ir["_internal"] = {
"body": list(filterfalse(rpartial(isinstance, AnnAssign), body)),
"from_name": parsed_source.name,
"from_type": "cls",
}
return ir
assert isinstance(function_def,
FunctionDef), "Expected 'FunctionDef' got `{!r}`".format(
type(function_def).__name__)
assert (function_name is None or function_def.name
== function_name), "Expected {!r} got {!r}".format(
function_name, function_def.name)
found_type = get_function_type(function_def)
# Read docstring
doc_str = (get_docstring(function_def) if isinstance(
function_def, FunctionDef) else None)
function_def = deepcopy(function_def)
function_def.args.args = (function_def.args.args if found_type == "static"
else function_def.args.args[1:])
if doc_str is None:
intermediate_repr = {
"name": function_name or function_def.name,
"params": OrderedDict(),
"returns": None,
}
else:
intermediate_repr = docstring(
doc_str.replace(":cvar", ":param"),
infer_type=infer_type,
)
intermediate_repr.update({
"name": function_name or function_def.name,
"type": function_type or found_type,
})
function_def.body = function_def.body if doc_str is None else function_def.body[
1:]
if function_def.body:
intermediate_repr["_internal"] = {
"body": function_def.body,
"from_name": function_def.name,
"from_type": found_type,
}
params_to_append = OrderedDict()
if (hasattr(function_def.args, "kwarg") and function_def.args.kwarg
and function_def.args.kwarg.arg in intermediate_repr["params"]):
_param = intermediate_repr["params"].pop(function_def.args.kwarg.arg)
assert "typ" in _param
_param["default"] = NoneStr
# if "typ" not in _param:
# _param["typ"] = (
# "Optional[dict]"
# if function_arguments.kwarg.annotation is None
# else to_code(function_arguments.kwarg.annotation).rstrip("\n")
# )
params_to_append[function_def.args.kwarg.arg] = _param
del _param
# Set defaults
# Fill with `None`s when no default is given to make the `zip` below it work cleanly
for args, defaults in (
("args", "defaults"),
("kwonlyargs", "kw_defaults"),
):
diff = len(getattr(function_def.args, args)) - len(
getattr(function_def.args, defaults))
if diff:
setattr(
function_def.args,
defaults,
list(islice(cycle(
(None, )), 10)) + getattr(function_def.args, defaults),
)
ir_merge(
intermediate_repr,
{
"params":
OrderedDict((func_arg2param(
getattr(function_def.args, args)[idx],
default=getattr(function_def.args, defaults)[idx],
) for args, defaults in (
("args", "defaults"),
("kwonlyargs", "kw_defaults"),
) for idx in range(len(getattr(function_def.args, args))))),
"returns":
None,
},
)
intermediate_repr["params"].update(params_to_append)
intermediate_repr["params"] = OrderedDict(
map(
partial(_set_name_and_type,
infer_type=infer_type,
word_wrap=word_wrap),
intermediate_repr["params"].items(),
))
# Convention - the final top-level `return` is the default
intermediate_repr = _interpolate_return(function_def, intermediate_repr)
if "return_type" in (intermediate_repr.get("returns") or iter(())):
intermediate_repr["returns"] = OrderedDict(
map(
partial(_set_name_and_type,
infer_type=infer_type,
word_wrap=word_wrap),
intermediate_repr["returns"].items(),
))
return intermediate_repr
def function(function_def, infer_type=False, function_type=None, function_name=None):
"""
Converts a method to our IR
:param function_def: AST node for function definition
:type function_def: ```Union[FunctionDef, FunctionType]```
:param infer_type: Whether to try inferring the typ (from the default)
:type infer_type: ```bool```
: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 function_name: name of function_def
:type function_name: ```str```
:return: 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```
"""
if isinstance(function_def, FunctionType):
# Dynamic function, i.e., this isn't source code; and is in your memory
ir = _inspect(function_def, function_name)
parsed_source = ast.parse(getsource(function_def).lstrip()).body[0]
ir["_internal"] = {
"body": list(
filterfalse(rpartial(isinstance, AnnAssign), parsed_source.body)
),
"from_name": parsed_source.name,
"from_type": "cls",
}
return ir
assert isinstance(
function_def, FunctionDef
), "Expected 'FunctionDef' got `{!r}`".format(type(function_def).__name__)
assert (
function_name is None or function_def.name == function_name
), "Expected {!r} got {!r}".format(function_name, function_def.name)
found_type = get_function_type(function_def)
# Read docstring
intermediate_repr_docstring = (
get_docstring(function_def) if isinstance(function_def, FunctionDef) else None
)
if intermediate_repr_docstring is None:
intermediate_repr = {
"name": function_name or function_def.name,
"params": OrderedDict(
map(
func_arg2param,
function_def.args.args
if found_type == "static"
else function_def.args.args[1:],
)
),
"returns": None,
}
else:
intermediate_repr = docstring(
intermediate_repr_docstring.replace(":cvar", ":param"),
infer_type=infer_type,
)
intermediate_repr.update(
{
"name": function_name or function_def.name,
"type": function_type or found_type,
}
)
if function_def.body:
intermediate_repr["_internal"] = {
"body": function_def.body
if intermediate_repr_docstring is None
else function_def.body[1:],
"from_name": function_def.name,
"from_type": found_type,
}
params_to_append = OrderedDict()
if (
hasattr(function_def.args, "kwarg")
and function_def.args.kwarg
and function_def.args.kwarg.arg in intermediate_repr["params"]
):
_param = intermediate_repr["params"].pop(function_def.args.kwarg.arg)
assert "typ" in _param
_param["default"] = NoneStr
# if "typ" not in _param:
# _param["typ"] = (
# "Optional[dict]"
# if function_def.args.kwarg.annotation is None
# else to_code(function_def.args.kwarg.annotation).rstrip("\n")
# )
params_to_append[function_def.args.kwarg.arg] = _param
del _param
idx = count()
# Set defaults
if intermediate_repr["params"]:
deque(
map(
lambda args_defaults: deque(
map(
lambda idxparam_idx_arg: intermediate_repr["params"][
idxparam_idx_arg[2].arg
].update(
dict(
(
{}
if getattr(idxparam_idx_arg[2], "annotation", None)
is None
else dict(
typ=to_code(
idxparam_idx_arg[2].annotation
).rstrip("\n")
)
),
**(
lambda _defaults: dict(
default=(
lambda v: (
_defaults[idxparam_idx_arg[1]]
if isinstance(
_defaults[idxparam_idx_arg[1]],
(NameConstant, Constant),
)
else v
)
if v is None
else v
)(get_value(_defaults[idxparam_idx_arg[1]]))
)
if idxparam_idx_arg[1] < len(_defaults)
and _defaults[idxparam_idx_arg[1]] is not None
else {}
)(getattr(function_def.args, args_defaults[1])),
)
),
(
lambda _args: map(
lambda idx_arg: (next(idx), idx_arg[0], idx_arg[1]),
enumerate(
filterfalse(
lambda _arg: _arg.arg[0] == "*",
(
_args
if found_type == "static"
or args_defaults[0] == "kwonlyargs"
else _args[1:]
),
)
),
)
)(getattr(function_def.args, args_defaults[0])),
),
maxlen=0,
),
(("args", "defaults"), ("kwonlyargs", "kw_defaults")),
),
maxlen=0,
)
intermediate_repr["params"].update(params_to_append)
intermediate_repr["params"] = OrderedDict(
map(
partial(_set_name_and_type, infer_type=infer_type),
intermediate_repr["params"].items(),
)
)
# Convention - the final top-level `return` is the default
intermediate_repr = _interpolate_return(function_def, intermediate_repr)
if "return_type" in (intermediate_repr.get("returns") or iter(())):
intermediate_repr["returns"] = OrderedDict(
map(
partial(_set_name_and_type, infer_type=infer_type),
intermediate_repr["returns"].items(),
)
)
return intermediate_repr
def _scan_phase_numpydoc_and_google(docstring, arg_tokens, return_tokens,
style):
"""
numpydoc and google scanner phase. Lexical analysis; to some degree…
:param docstring: the docstring
:type docstring: ```str```
:param arg_tokens: Valid tokens like `"Parameters\n----------"`
:type arg_tokens: ```Tuple[str]```
:param return_tokens: Valid tokens like `"Returns\n-------"`
:type return_tokens: ```Tuple[str]```
:param style: the style of docstring
:type style: ```Style```
:returns: List with each element a tuple of (whether value is a token, value)
:rtype: ```Dict[str, str]```
"""
scanned: Dict[str, List[List[str]]] = {
token: []
for token in ("doc", ) + arg_tokens + return_tokens
}
# ^ Dict[Union[Literal["doc"], arg_tokens, return_tokens], List[dict]]
# First doc, if present
_start_idx, _end_idx, _found = location_within(docstring, arg_tokens)
if _start_idx == -1:
# Return type no args?
_start_idx, _end_idx, _found = location_within(docstring,
return_tokens)
if _start_idx > -1:
namespace = _found
scanned["doc"] = docstring[:_start_idx].strip()
docstring = docstring[_end_idx + 1:] # .strip()
else:
scanned["doc"] = docstring.strip()
return scanned
# Scan all lines so that that each element in `stacker` refers to one 'unit'
stacker, docstring_lines = [], docstring.splitlines()
first_indent = count_iter_items(takewhile(str.isspace, docstring_lines[0]))
for line_no, line in enumerate(docstring_lines):
indent = count_iter_items(takewhile(str.isspace, line))
if indent == first_indent:
stacker.append([line])
else:
if indent < first_indent:
scanned[namespace] = scanned.get(namespace,
[]) + deepcopy(stacker)
stacker.clear()
if len(docstring_lines) > line_no + 3 and any(
filter(rpartial(eq, docstring_lines[line_no + 1]),
return_tokens)
# return_token = return_tokens[0].splitlines()
# filter(rpartial(eq, docstring_lines[line_no + 1]), return_tokens)
):
return_indent = count_iter_items(
takewhile(str.isspace, docstring_lines[line_no + 3]))
next_smallest_indent = count_iter_items(
takewhile(
partial(le, return_indent),
map(
lambda l: count_iter_items(
takewhile(str.isspace, l)),
docstring_lines[line_no + 3:],
),
))
scanned[return_tokens[0]] = docstring_lines[
line_no + 2:line_no + 3 + next_smallest_indent]
scanned_afterward = docstring_lines[line_no + 3 +
next_smallest_indent:]
else:
scanned_afterward = docstring_lines[line_no + 1:]
if (len(scanned_afterward) > 1
and scanned_afterward[0] == return_tokens[0]):
return_indent = count_iter_items(
takewhile(str.isspace, scanned_afterward[1]))
next_smallest_indent = count_iter_items(
takewhile(
partial(le, return_indent),
map(
lambda l: count_iter_items(
takewhile(str.isspace, l)),
scanned_afterward[2:],
),
))
scanned[return_tokens[0]] = scanned_afterward[
1:next_smallest_indent + 2]
scanned_afterward = (
None if next_smallest_indent == 0 else
scanned_afterward[next_smallest_indent + 2:])
if scanned_afterward:
scanned["scanned_afterward"] = scanned_afterward
break
else:
stacker[-1].append(line)
# Split out return, if present and not already set
if not scanned.get(return_tokens[0], False):
stacker = _return_parse_phase_numpydoc_and_google(
return_tokens, scanned, stacker, style)
if stacker:
scanned[namespace] = stacker
# # Join ["a", "b"] into ["ab"] as each has been resolved to a scanned symbol now
# scanned.update(
# {
# token: [
# line_list
# if len(line_list) < 3 or isinstance(line_list, str)
# else [
# "{} {}".format(
# line_list[0], " ".join(map(str.lstrip, line_list[1:]))
# )
# ]
# for line_list in scanned[token]
# ]
# for token in arg_tokens + return_tokens
# if scanned.get(token)
# }
# )
return scanned
def test_rpartial(self) -> None:
""" Test that rpartial works as advertised """
self.assertTrue(rpartial(isinstance, str)(""))
self.assertFalse(rpartial(isinstance, str)(0))
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 annotate_ancestry(node):
"""
Look to your roots. Find the child; find the parent.
Sets _location attribute to every child node.
:param node: AST node. Will be annotated in-place.
:type node: ```AST```
:return: Annotated AST node; also `node` arg will be annotated in-place.
:rtype: ```AST```
"""
# print("annotating", getattr(node, "name", None))
node._location = [node.name] if hasattr(node, "name") else []
parent_location = []
for _node in walk(node):
name = [_node.name] if hasattr(_node, "name") else []
for child_node in iter_child_nodes(_node):
if hasattr(child_node, "name") and not isinstance(child_node, alias):
child_node._location = name + [child_node.name]
parent_location = child_node._location
elif isinstance(child_node, (Constant, Str)):
child_node._location = parent_location + [get_value(child_node)]
elif isinstance(child_node, Assign) and all(
map(rpartial(isinstance, Name), child_node.targets)
):
for target in child_node.targets:
child_node._location = name + [target.id]
elif isinstance(child_node, AnnAssign) and isinstance(
child_node.target, Name
):
child_node._location = name + [child_node.target.id]
if isinstance(child_node, FunctionDef):
def set_index_and_location(idx_arg):
"""
:param idx_arg: Index and Any; probably out of `enumerate`
:type idx_arg: ```Tuple[int, Any]```
:return: Second element, with _idx set with value of first
:rtype: ```Any```
"""
idx_arg[1]._idx = idx_arg[0]
idx_arg[1]._location = child_node._location + [idx_arg[1].arg]
return idx_arg[1]
child_node.args.args = list(
map(
set_index_and_location,
enumerate(
child_node.args.args,
-1
if len(child_node.args.args) > 0
and child_node.args.args[0].arg
in frozenset(("self", "cls"))
else 0,
),
)
)
child_node.args.kwonlyargs = list(
map(
set_index_and_location,
enumerate(
child_node.args.kwonlyargs,
0,
),
)
)
return node
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,
))
