def ast_default_field(value): return Call( func=Name(id='field', ctx=Load()), args=[], keywords=[keyword(arg='default', value=Constant(value=value, kind=None))] )
def visit_Name(self, name_node: ast.Name): 'Visiting a name - which should represent something' id = self.resolve_id(name_node.id) name_node.rep = self.get_rep(id)
def include_var(self, var: Expr, assign=False): """Returns a variable AST node for the given Relay var depending on whether it must appear in an assignment or not.""" name = self.get_var_name(var) return Name(name, Store() if assign else Load())
def NameConstant(value): return Name(id=repr(value))
def visit_global_var(self, gvar: Expr): # we don't need to add numbers to global var names because # the *names* are checked for uniqueness in the mod return (Name(gvar.name_hint, Load()), [])
def global_action_points_class(project: CachedProject) -> str: tree = Module(body=[], type_ignores=[]) tree.body.append( ImportFrom( module=arcor2.data.common.__name__, names=[ alias(name=ActionPoint.__name__, asname=None), alias(name=Position.__name__, asname=None), alias(name=Pose.__name__, asname=None), alias(name=ProjectRobotJoints.__name__, asname=None), ], level=0, )) tree.body.append( ImportFrom( module=copy.__name__, names=[alias(name=copy.deepcopy.__name__, asname=None)], level=0, )) tree.body.append( ImportFrom( module=RES_MODULE, names=[alias(name=RES_CLS, asname=None)], level=0, )) aps_init_body: list[Union[Assign, Pass]] = [] for ap in project.action_points: ap_cls_body: list[Assign] = [ Assign( targets=[ Attribute(value=Name(id="self", ctx=Load()), attr="_position", ctx=Store()) ], value=Attribute( value=Call( func=Attribute( value=Attribute(value=Name(id="res", ctx=Load()), attr="project", ctx=Load()), attr=CachedProject.bare_action_point.__name__, ctx=Load(), ), args=[Str(s=ap.id, kind="")], keywords=[], ), attr="position", ctx=Load(), ), type_comment=None, ) ] ap_type_name = humps.pascalize(ap.name) ap_joints_init_body: list[Assign] = [] for joints in project.ap_joints(ap.id): ap_joints_init_body.append( Assign( targets=[ Attribute(value=Name(id="self", ctx=Load()), attr=f"_{joints.name}", ctx=Store()) ], value=Call( func=Attribute( value=Attribute(value=Name(id="res", ctx=Load()), attr="project", ctx=Load()), attr="joints", ctx=Load(), ), args=[Str(s=joints.id, kind="")], keywords=[], ), type_comment=None, )) if ap_joints_init_body: ap_joints_cls_def = ClassDef( name=f"{ap_type_name}Joints", bases=[], keywords=[], body=[ FunctionDef( name="__init__", args=arguments( posonlyargs=[], args=[ arg(arg="self", annotation=None, type_comment=None), arg(arg="res", annotation=Name(id=RES_CLS, ctx=Load()), type_comment=None), ], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[], ), body=ap_joints_init_body, decorator_list=[], returns=None, type_comment=None, ) ], decorator_list=[], ) for joints in project.ap_joints(ap.id): ap_joints_cls_def.body.append( FunctionDef( name=joints.name, args=arguments( posonlyargs=[], args=[ arg(arg="self", annotation=None, type_comment=None) ], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[], ), body=[ Return(value=Call( func=Name(id=copy.deepcopy.__name__, ctx=Load()), args=[ Attribute(value=Name(id="self", ctx=Load()), attr=f"_{joints.name}", ctx=Load()) ], keywords=[], )) ], decorator_list=[Name(id="property", ctx=Load())], returns=Name(id=ProjectRobotJoints.__name__, ctx=Load()), type_comment=None, )) tree.body.append(ap_joints_cls_def) ap_cls_body.append( Assign( targets=[ Attribute(value=Name(id="self", ctx=Load()), attr="joints", ctx=Store()) ], value=Call( func=Name(id=f"{ap_type_name}Joints", ctx=Load()), args=[Name(id="res", ctx=Load())], keywords=[], ), type_comment=None, )) ap_orientations_init_body: list[Assign] = [] for ori in project.ap_orientations(ap.id): ap_orientations_init_body.append( Assign( targets=[ Attribute(value=Name(id="self", ctx=Load()), attr=f"_{ori.name}", ctx=Store()) ], value=Call( func=Attribute( value=Attribute(value=Name(id="res", ctx=Load()), attr="project", ctx=Load()), attr="pose", ctx=Load(), ), args=[Str(s=ori.id, kind="")], keywords=[], ), type_comment=None, )) if ap_orientations_init_body: ap_orientations_cls_def = ClassDef( name=f"{ap_type_name}Poses", bases=[], keywords=[], body=[ FunctionDef( name="__init__", args=arguments( posonlyargs=[], args=[ arg(arg="self", annotation=None, type_comment=None), arg(arg="res", annotation=Name(id=RES_CLS, ctx=Load()), type_comment=None), ], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[], ), body=ap_orientations_init_body, decorator_list=[], returns=None, type_comment=None, ) ], decorator_list=[], ) for ori in project.ap_orientations(ap.id): ap_orientations_cls_def.body.append( FunctionDef( name=ori.name, args=arguments( posonlyargs=[], args=[ arg(arg="self", annotation=None, type_comment=None) ], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[], ), body=[ Return(value=Call( func=Name(id=copy.deepcopy.__name__, ctx=Load()), args=[ Attribute(value=Name(id="self", ctx=Load()), attr=f"_{ori.name}", ctx=Load()) ], keywords=[], )) ], decorator_list=[Name(id="property", ctx=Load())], returns=Name(id=Pose.__name__, ctx=Load()), type_comment=None, )) tree.body.append(ap_orientations_cls_def) ap_cls_body.append( Assign( targets=[ Attribute(value=Name(id="self", ctx=Load()), attr="poses", ctx=Store()) ], value=Call( func=Name(id=f"{ap_type_name}Poses", ctx=Load()), args=[Name(id="res", ctx=Load())], keywords=[], ), type_comment=None, )) ap_cls_def = ClassDef( name=ap_type_name, bases=[], keywords=[], body=[ FunctionDef( name="__init__", args=arguments( posonlyargs=[], args=[ arg(arg="self", annotation=None, type_comment=None), arg(arg="res", annotation=Name(id=RES_CLS, ctx=Load()), type_comment=None), ], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[], ), body=ap_cls_body, decorator_list=[], returns=None, type_comment=None, ) ], decorator_list=[], ) # add copy property for position ap_cls_def.body.append( FunctionDef( name="position", args=arguments( posonlyargs=[], args=[arg(arg="self", annotation=None, type_comment=None)], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[], ), body=[ Return(value=Call( func=Name(id=copy.deepcopy.__name__, ctx=Load()), args=[ Attribute(value=Name(id="self", ctx=Load()), attr="_position", ctx=Load()) ], keywords=[], )) ], decorator_list=[Name(id="property", ctx=Load())], returns=Name(id=Position.__name__, ctx=Load()), type_comment=None, )) tree.body.append(ap_cls_def) aps_init_body.append( Assign( targets=[ Attribute(value=Name(id="self", ctx=Load()), attr=ap.name, ctx=Store()) ], value=Call(func=Name(id=ap_type_name, ctx=Load()), args=[Name(id="res", ctx=Load())], keywords=[]), type_comment=None, )) if not aps_init_body: # there are no action points aps_init_body.append(Pass()) aps_cls_def = ClassDef( name="ActionPoints", bases=[], keywords=[], body=[ FunctionDef( name="__init__", args=arguments( posonlyargs=[], args=[ arg(arg="self", annotation=None, type_comment=None), arg(arg="res", annotation=Name(id=RES_CLS, ctx=Load()), type_comment=None), ], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[], ), body=aps_init_body, decorator_list=[], returns=None, type_comment=None, ) ], decorator_list=[], ) tree.body.append(aps_cls_def) return tree_to_str(tree)
def F_6(x: int, y, *args: [str]) -> bool: return x == 6 expr = """ def F_6(x: int, y, *args: [str]) -> bool: return x == 6 """ if __name__ == "__main__": mod = Module(body=[ FunctionDef(name='F_6', args=arguments(args=[ arg(arg='x', annotation=Name(id='int', ctx=Load())), arg(arg='y', annotation=None) ], vararg=arg(arg='args', annotation=List(elts=[ Name(id='str', ctx=Load()) ], ctx=Load())), kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=[ Return(value=Compare(left=Name(id='x', ctx=Load()), ops=[Eq()], comparators=[Num(n=6)]))
def document(sentences, **kw): """ This macro takes literal strings and converts them into: _help_ID = type_hint+STRING where: ID is the first target of the last assignment. type_hint is the assigned type and default value (only works for a few types) STRING is the literal string """ for n in range(len(sentences)): s = sentences[n] if not n: prev = s continue # The whole sentence is a string? if (isinstance(s, Expr) and isinstance(s.value, Str) and # and the previous is an assign isinstance(prev, Assign)): # noqa: E128 # Apply it to the first target target = prev.targets[0] value = prev.value # Extract its name # variables and attributes are supported if isinstance(target, Name): # Note: The support for variables isn't currently used name = target.id is_attr = False elif isinstance(target, Attribute): name = target.attr is_attr = True # Remove starting underscore if name[0] == '_': name = name[1:] # Create a _help_ID doc_id = '_help_' + name # Create the type hint for numbers, strings and booleans type_hint = '' post_hint = '' if isinstance(value, Num): type_hint = '[number={}]'.format(value.n) elif isinstance(value, UnaryOp) and isinstance( value.operand, Num) and isinstance(value.op, USub): # -Num type_hint = '[number={}]'.format(-value.operand.n) elif isinstance(value, Str): type_hint = "[string='{}']".format(value.s) elif isinstance(value, NameConstant) and isinstance( value.value, bool): type_hint = '[boolean={}]'.format(str(value.value).lower()) elif isinstance(value, Attribute): # Used for the default options. I.e. GS.def_global_option val = eval(unparse(value)) if isinstance(val, bool): # Not used yet type_hint = '[boolean={}]'.format(str(val).lower()) elif isinstance(val, (int, float)): # Not used yet type_hint = '[number={}]'.format(val) elif isinstance(val, str): type_hint = "[string='{}']".format(val) post_hint += '. Affected by global options' # Transform the string into an assign for _help_ID if is_attr: target = Attribute(value=Name(id='self', ctx=Load()), attr=doc_id, ctx=Store()) else: target = Name(id=doc_id, ctx=Store()) # Reuse the s.value Str help_str = s.value help_str.s = type_hint + s.value.s.rstrip() + post_hint sentences[n] = Assign(targets=[target], value=help_str) # Copy the line number from the original docstring copy_location(target, s) copy_location(sentences[n], s) prev = s # Return the modified AST return sentences
def visit_Name(self, node: ast.Name): if node.id in self.repldict: node.id = self.repldict[node.id] return node
def visit_Attribute(self, node): """Flatten one level of attribute access.""" new_node = Name("%s.%s" % (node.value.id, node.attr), node.ctx) return copy_location(new_node, node)
"History", "LambdaCallback", "LearningRateScheduler", "ModelCheckpoint", "ProgbarLogger", "ReduceLROnPlateau", "RemoteMonitor", "TensorBoard", "TerminateOnNaN", ] ] ] = None ''' class_ast = ClassDef( bases=[Name("object", Load())], body=[ Expr( set_value( "\n Acquire from the official tensorflow_datasets model zoo," " or the ophthalmology focussed ml-prepare library\n\n " ':cvar dataset_name: name of dataset. Defaults to "mnist"\n ' ':cvar tfds_dir: directory to look for models in. Defaults to "~/tensorflow_datasets"\n ' ':cvar K: backend engine, e.g., `np` or `tf`. Defaults to "np"\n ' ":cvar as_numpy: Convert to numpy ndarrays. Defaults to None\n " ":cvar data_loader_kwargs: pass this as arguments to data_loader function\n " ":cvar return_type: Train and tests dataset splits. Defaults to (np.empty(0), np.empty(0))", ) ), AnnAssign( annotation=Name(
def visit_Name(self, node: ast.Name) -> Any: # Rename connector if node.id in self.repl_dict: node.id = self.repl_dict[node.id] return self.generic_visit(node)
import ast import os from ast import AnnAssign, Assign, ClassDef, Constant, ImportFrom, Index, List, Load, Module, Name, Store, Subscript, \ alias from ast_decompiler import decompile dataclass_decorator = Name(id='dataclass', ctx=Load()) text_data_value = Subscript( value=Name(id='Optional', ctx=Load()), slice=Index(value=Name(id='TextData', ctx=Load())), ctx=Load() ) chart_xy_data_value = Subscript( value=Name(id='Optional', ctx=Load()), slice=Index(value=Name(id='ChartXYData', ctx=Load())), ctx=Load() ) chart_category_data_value = Subscript( value=Name(id='Optional', ctx=Load()), slice=Index(value=Name(id='ChartCategoryData', ctx=Load())), ctx=Load() ) chart_bubble_data_value = Subscript( value=Name(id='Optional', ctx=Load()), slice=Index(value=Name(id='ChartBubbleData', ctx=Load())), ctx=Load() ) table_data_value = Subscript( value=Name(id='Optional', ctx=Load()), slice=Index(value=Name(id='TableData', ctx=Load())),
def compile_function_ast(expressions, symbols, arg_names, output_names=None, funname='anonymous', return_ast=False, print_code=False, definitions=None, vectorize=True, use_file=False): ''' expressions: list of equations as string ''' # TODO: definitions should be used only if necessary from collections import OrderedDict table = OrderedDict() aa = arg_names if output_names is not None: aa = arg_names + [output_names] for a in aa: symbol_group = a[0] date = a[1] an = a[2] for b in symbols[symbol_group]: index = symbols[symbol_group].index(b) table[(b, date)] = (an, index) table_symbols = {k: (std_date_symbol(*k)) for k in table.keys()} # standard assignment: i.e. k = s[0] def index(x): Index(Num(x)) # declare symbols aux_short_names = [e[2] for e in arg_names if e[0] == 'auxiliaries'] preamble = [] for k in table: # order it # k : var, date arg, pos = table[k] if not (arg in aux_short_names): std_name = table_symbols[k] val = Subscript(value=Name(id=arg, ctx=Load()), slice=index(pos), ctx=Load()) line = Assign(targets=[Name(id=std_name, ctx=Store())], value=val) if arg != 'out': preamble.append(line) body = [] std_dates = StandardizeDates(symbols, aa) if definitions is not None: for k, v in definitions.items(): if isinstance(k, str): lhs = ast.parse(k).body[0].value if isinstance(v, str): rhs = ast.parse(v).body[0].value else: rhs = v lhs = std_dates.visit(lhs) rhs = std_dates.visit(rhs) vname = lhs.id line = Assign(targets=[Name(id=vname, ctx=Store())], value=rhs) preamble.append(line) outs = [] for i, expr in enumerate(expressions): expr = ast.parse(expr).body[0].value # if definitions is not None: # expr = ReplaceName(defs).visit(expr) rexpr = std_dates.visit(expr) rhs = rexpr if output_names is not None: varname = symbols[output_names[0]][i] date = output_names[1] out_name = table_symbols[(varname, date)] else: out_name = 'out_{}'.format(i) line = Assign(targets=[Name(id=out_name, ctx=Store())], value=rhs) body.append(line) line = Assign(targets=[ Subscript(value=Name(id='out', ctx=Load()), slice=index(i), ctx=Store()) ], value=Name(id=out_name, ctx=Load())) body.append(line) arg_names = [e for e in arg_names if e[0] != "auxiliaries"] args = [e[2] for e in arg_names] + ['out'] if is_python_3: from ast import arg f = FunctionDef(name=funname, args=arguments(args=[arg(arg=a) for a in args], vararg=None, kwarg=None, kwonlyargs=[], kw_defaults=[], defaults=[]), body=preamble + body, decorator_list=[]) else: f = FunctionDef(name=funname, args=arguments( args=[Name(id=a, ctx=Param()) for a in args], vararg=None, kwarg=None, kwonlyargs=[], kw_defaults=[], defaults=[]), body=preamble + body, decorator_list=[]) mod = Module(body=[f]) mod = ast.fix_missing_locations(mod) import dolo.config if dolo.config.debug: print_code = True if print_code: s = "Function {}".format(mod.body[0].name) print("-" * len(s)) print(s) print("-" * len(s)) print(to_source(mod)) if vectorize: from numba import float64, void coredims = [len(symbols[an[0]]) for an in arg_names] signature = str.join(',', ['(n_{})'.format(d) for d in coredims]) n_out = len(expressions) if n_out in coredims: signature += '->(n_{})'.format(n_out) # ftylist = float64[:](*([float64[:]] * len(coredims))) fty = "void(*[float64[:]]*{})".format(len(coredims) + 1) else: signature += ',(n_{})'.format(n_out) fty = "void(*[float64[:]]*{})".format(len(coredims) + 1) ftylist = [fty] else: signature = None ftylist = None if use_file: fun = eval_ast_with_file(mod, print_code=True) else: fun = eval_ast(mod) jitted = njit(fun) if vectorize: gufun = guvectorize([fty], signature, target='parallel', nopython=True)(fun) return jitted, gufun else: return jitted
fix_missing_locations, ) from functools import partial from operator import add from doctrans.ast_utils import maybe_type_comment, set_arg, set_slice, set_value from doctrans.tests.mocks.docstrings import docstring_google_tf_adadelta_str return_ast = Return( value=Tuple( ctx=Load(), elts=[ Call( args=[set_value(0)], func=Attribute( Name("np", Load()), "empty", Load(), ), keywords=[], expr=None, expr_func=None, ) ] * 2, expr=None, ), expr=None, ) class_with_method_str = ''' class C(object):
from shutil import rmtree from tempfile import mkdtemp from unittest import TestCase from unittest.mock import patch from doctrans import emit, parse from doctrans.ast_utils import maybe_type_comment, set_value from doctrans.gen import gen from doctrans.pure_utils import rpartial from doctrans.source_transformer import to_code from doctrans.tests.mocks.methods import function_adder_ast from doctrans.tests.utils_for_tests import run_ast_test method_adder_ast = deepcopy(function_adder_ast) method_adder_ast.body[0] = Expr(set_value(" C class (mocked!) ")) method_adder_ast.decorator_list = [Name("staticmethod", Load())] del function_adder_ast def populate_files(tempdir, input_module_str=None): """ Populate files in the tempdir :param tempdir: Temporary directory :type tempdir: ```str``` :param input_module_str: Input string to write to the input_filename. If None, uses preset mock module. :type input_module_str: ```Optional[str]``` :returns: input filename, input str, expected_output :rtype: ```Tuple[str, str, str, Module]```
def xonsh_call(name, args, lineno=None, col=None): """Creates the AST node for calling a function of a given name.""" return Call(func=Name(id=name, ctx=Load(), lineno=lineno, col_offset=col), args=args, keywords=[], starargs=None, kwargs=None, lineno=lineno, col_offset=col)
def populate_files(tempdir, input_module_str=None): """ Populate files in the tempdir :param tempdir: Temporary directory :type tempdir: ```str``` :param input_module_str: Input string to write to the input_filename. If None, uses preset mock module. :type input_module_str: ```Optional[str]``` :returns: input filename, input str, expected_output :rtype: ```Tuple[str, str, str, Module]``` """ input_filename = os.path.join(tempdir, "input.py") input_class_name = "Foo" input_class_ast = emit.class_( parse.function(deepcopy(method_adder_ast)), emit_call=False, class_name=input_class_name, ) input_module_ast = Module( body=[ input_class_ast, Assign(targets=[Name("input_map", Store())], value=Dict( keys=[set_value(input_class_name)], values=[Name(input_class_name, Load())], expr=None, ), expr=None, lineno=None, **maybe_type_comment), Assign( targets=[Name("__all__", Store())], value=List( ctx=Load(), elts=[set_value(input_class_name), set_value("input_map")], expr=None, ), expr=None, lineno=None, **maybe_type_comment), ], type_ignores=[], stmt=None, ) input_module_str = input_module_str or to_code(input_module_ast) # expected_output_class_str = ( # "class FooConfig(object):\n" # ' """\n' # " The amazing Foo\n\n" # " :cvar a: An a. Defaults to 5\n" # ' :cvar b: A b. Defaults to 16"""\n' # " a = 5\n" # " b = 16\n\n" # " def __call__(self):\n" # " self.a = 5\n" # " self.b = 16\n" # ) expected_class_ast = emit.class_( parse.function(deepcopy(method_adder_ast)), emit_call=True, class_name="{input_class_name}Config".format( input_class_name=input_class_name), ) with open(input_filename, "wt") as f: f.write(input_module_str) return input_filename, input_module_ast, input_class_ast, expected_class_ast
def empty_script_tree(project_id: str, add_main_loop: bool = True) -> Module: """Creates barebones of the script (empty 'main' function). Returns ------- """ main_body: list[stmt] = [ Assign( targets=[Name(id="aps", ctx=Store())], value=Call(func=Name(id="ActionPoints", ctx=Load()), args=[Name(id="res", ctx=Load())], keywords=[]), type_comment=None, ) ] if add_main_loop: main_body.append( While(test=NameConstant(value=True, kind=None), body=[Pass()], orelse=[])) else: """put there "pass" in order to make code valid even if there is no other statement (e.g. no object from resources)""" main_body.append(Pass()) # TODO helper function for try ... except tree = Module( body=[ FunctionDef( name="main", args=arguments( posonlyargs=[], args=[ arg(arg="res", annotation=Name(id=RES_CLS, ctx=Load()), type_comment=None) ], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[], ), body=main_body, decorator_list=[], returns=NameConstant(value=None, kind=None), type_comment=None, ), If( test=Compare(left=Name(id="__name__", ctx=Load()), ops=[Eq()], comparators=[Str(s="__main__", kind="")]), body=[ Try( body=[ With( items=[ withitem( context_expr=Call( func=Name(id=RES_CLS, ctx=Load()), args=[], keywords=[], ), optional_vars=Name(id="res", ctx=Store()), ) ], body=[ Expr(value=Call( func=Name(id="main", ctx=Load()), args=[Name(id="res", ctx=Load())], keywords=[], )) ], type_comment=None, ) ], handlers=[ ExceptHandler( type=Name(id=Exception.__name__, ctx=Load()), name="e", body=[ Expr(value=Call( func=Name(id="print_exception", ctx=Load()), args=[Name(id="e", ctx=Load())], keywords=[], )) ], ) ], orelse=[], finalbody=[], ) ], orelse=[], ), ], type_ignores=[], ) add_import(tree, "arcor2_runtime.exceptions", "print_exception", try_to_import=False) add_import(tree, RES_MODULE, RES_CLS, try_to_import=False) add_import(tree, "action_points", "ActionPoints", try_to_import=False) return tree
def visit_Call(self, node): if isinstance(node.func, Name) and node.func.id.lower() == 'col': newName = self.getName(self.sc[node.args[0].s]) return Name(newName, Load())
def symbol_type(name: str): return Name(id=name, ctx=Load())
def visit_Name(self, name_node: ast.Name): 'Visiting a name - which should represent something' id = self.resolve_id(name_node.id) if isinstance(id, ast.AST): name_node.rep = self.get_rep(id) # type: ignore
def _add_logic(container: Container, current_action: Action, super_container: Optional[Container] = None) -> None: # more paths could lead to the same action, so it might be already added # ...this is easier than searching the tree if current_action.id in added_actions: logger.debug(f"Action {current_action.name} already added, skipping.") return inputs, outputs = project.action_io(current_action.id) logger.debug(f"Adding action {current_action.name}, with {len(inputs)} input(s) and {len(outputs)} output(s).") act = current_action.parse_type() ac_obj = scene.object(act.obj_id).name args: List[AST] = [] # TODO make sure that the order of parameters is correct / re-order for param in current_action.parameters: if param.type == ActionParameter.TypeEnum.LINK: parsed_link = param.parse_link() parent_action = project.action(parsed_link.action_id) # TODO add support for tuples assert len(parent_action.flow(FlowTypes.DEFAULT).outputs) == 1, "Only one result is supported atm." assert parsed_link.output_index == 0 res_name = parent_action.flow(FlowTypes.DEFAULT).outputs[0] # make sure that the result already exists if parent_action.id not in added_actions: raise SourceException( f"Action {current_action.name} attempts to use result {res_name} " f"of subsequent action {parent_action.name}." ) args.append(Name(id=res_name, ctx=Load())) elif param.type == ActionParameter.TypeEnum.CONSTANT: args.append(Name(id=project.constant(param.str_from_value()).name, ctx=Load())) else: plugin = plugin_from_type_name(param.type) args.append(plugin.parameter_ast(type_defs, scene, project, current_action.id, param.name)) list_of_imp_tup = plugin.need_to_be_imported(type_defs, scene, project, current_action.id, param.name) if list_of_imp_tup: # TODO what if there are two same names? for imp_tup in list_of_imp_tup: add_import(tree, imp_tup.module_name, imp_tup.class_name, try_to_import=False) add_method_call( container.body, ac_obj, act.action_type, args, [keyword(arg="an", value=Str(s=current_action.name, kind=""))], current_action.flow(FlowTypes.DEFAULT).outputs, ) added_actions.add(current_action.id) if not outputs: raise SourceException(f"Action {current_action.name} has no outputs.") elif len(outputs) == 1: output = outputs[0] if output.end == output.END: # TODO this is just temporary (while there is while loop), should be rather Return() container.body.append(Continue()) return seq_act = project.action(output.end) seq_act_inputs, _ = project.action_io(seq_act.id) if len(seq_act_inputs) > 1: # the action belongs to a different block if seq_act.id in added_actions: return logger.debug(f"Action {seq_act.name} going to be added to super_container.") # test if this is the correct super_container -> count distance (number of blocks) to the START blocks_to_start: Dict[str, int] = {} for inp in seq_act_inputs: parsed_start = inp.parse_start() pact = project.action(parsed_start.start_action_id) blocks_to_start[pact.id] = _blocks_to_start(pact) winner = min(blocks_to_start, key=blocks_to_start.get) # type: ignore # TODO what is wrong with it? # TODO if blocks_to_start is cached somewhere, the second part of the condition is useless # it might happen that there are two different ways with the same distance if winner == current_action.id or all( value == list(blocks_to_start.values())[0] for value in blocks_to_start.values() ): assert super_container is not None _add_logic(super_container, seq_act) return logger.debug(f"Sequential action: {seq_act.name}") _add_logic(container, seq_act, super_container) else: root_if: Optional[If] = None # action has more outputs - each output should have condition for idx, output in enumerate(outputs): if not output.condition: raise SourceException("Missing condition.") # TODO use parameter plugin (action metadata will be needed - to get the return types) # TODO support for other countable types # ...this will only work for booleans import json condition_value = json.loads(output.condition.value) comp = NameConstant(value=condition_value, kind=None) what = output.condition.parse_what() output_name = project.action(what.action_id).flow(what.flow_name).outputs[what.output_index] cond = If( test=Compare(left=Name(id=output_name, ctx=Load()), ops=[Eq()], comparators=[comp]), body=[], orelse=[], ) if idx == 0: root_if = cond container.body.append(root_if) logger.debug(f"Adding branch for: {condition_value}") else: assert isinstance(root_if, If) root_if.orelse.append(cond) if output.end == output.END: cond.body.append(Continue()) # TODO should be rather return continue _add_logic(cond, project.action(output.end), container)
def generate_param_name(name): return Name(id=name, ctx=Param())
def visit_function(self, func: Expr): # Python's lambdas are very restrictive, so we do "name" inline functions converted_func, func_name = self.convert_func_node(func) return (Name(func_name, Load()), [converted_func])
def visit_FunctionDef(self, node): """ Instrument a function definition by creating a new report builder for this stack frame and putting it in a local variable. The local variable has the same name as the global variable so all calls can use the same CONTEXT_NAME symbol, but it means that I had to use this: x = globals()['x'].start_frame() Kind of ugly, but I think it was worth it to handle recursive calls. """ new_node = self.generic_visit(node) line_numbers = set() find_line_numbers(new_node, line_numbers) first_line_number = min(line_numbers) last_line_number = max(line_numbers) args = [Num(n=first_line_number), Num(n=last_line_number)] try_body = new_node.body globals_call = Call(func=Name(id='globals', ctx=Load()), args=[], keywords=[], starargs=None, kwargs=None) global_context = Subscript(value=globals_call, slice=Index(value=Str(s=CONTEXT_NAME)), ctx=Load()) start_frame_call = Call(func=Attribute(value=global_context, attr='start_frame', ctx=Load()), args=args, keywords=[], starargs=None, kwargs=None) context_assign = Assign(targets=[Name(id=CONTEXT_NAME, ctx=Store())], value=start_frame_call) new_node.body = [context_assign] if isinstance(try_body[0], Expr) and isinstance( try_body[0].value, Str): # Move docstring back to top of function. # noinspection PyUnresolvedReferences new_node.body.insert(0, try_body.pop(0)) # trace function parameter values arg_nodes = [] arg_nodes.extend(getattr(new_node.args, 'posonlyargs', [])) arg_nodes.extend(new_node.args.args) arg_nodes.append(new_node.args.kwarg) arg_nodes.append(new_node.args.vararg) arg_nodes.extend(new_node.args.kwonlyargs) for target in arg_nodes: if target is None: continue if isinstance(target, Name) and target.id == 'self': continue if isinstance(target, arg) and target.arg == 'self': continue new_node.body.append(self._trace_assignment(target, node.lineno)) if try_body: handler_body = [self._create_context_call('exception'), Raise()] new_node.body.append( Try(body=try_body, handlers=[ExceptHandler(body=handler_body)], orelse=[], finalbody=[])) self._set_statement_line_numbers(try_body, first_line_number) self._set_statement_line_numbers(handler_body, last_line_number) return new_node
def visit_Lambda(self, node): args = [self.visit(arg) for arg in node.args.args] body = self.visit(node.body) n = Call(Name('Lambda', Load()), [Tuple(args, Load()), body], [], None, None) return fix_missing_locations(n)
def mkdir_and_emit_file( name_orig_ir, emit_name, module_name, new_module_name, filesystem_layout, output_directory, dry_run, ): """ Generate Java-package—or match input—style file hierarchy from fully-qualified module name :param name_orig_ir: FQ module name, original filename path, IR :type name_orig_ir: ```Tuple[str, str, dict]``` :param emit_name: What type(s) to generate. :type emit_name: ```List[Literal["argparse", "class", "function", "sqlalchemy", "sqlalchemy_table"]]``` :param module_name: Name of [original] module :type module_name: ```str``` :param new_module_name: Name of [new] module :type new_module_name: ```str``` :param filesystem_layout: Hierarchy of folder and file names generated. "java" is file per package per name. :type filesystem_layout: ```Literal["java", "as_input"]``` :param output_directory: Where to place the generated exposed interfaces to the given `--module`. :type output_directory: ```str``` :param dry_run: Show what would be created; don't actually write to the filesystem :type dry_run: ```bool``` :returns: Import to generated module :rtype: ```ImportFrom``` """ mod_name, _, name = name_orig_ir[0].rpartition(".") original_relative_filename_path, ir = name_orig_ir[1], name_orig_ir[2] mod_path = path.join( output_directory, new_module_name, mod_name.replace(".", path.sep), ) if not path.isdir(mod_path): if dry_run: print("mkdir\t{mod_path!r}".format(mod_path=mod_path)) else: makedirs(mod_path) init_filepath = path.join( path.dirname(mod_path), "__init__{extsep}py".format(extsep=extsep) ) if dry_run: print("touch\t{init_filepath!r}".format(init_filepath=init_filepath)) else: open(init_filepath, "a").close() gen_node = getattr(emit, emit_name.replace("class", "class_"))( ir, **dict( **{"{emit_name}_name".format(emit_name=emit_name): name}, **{} if emit_name == "class" else {"function_type": "static"} ) ) __all___node = Assign( targets=[Name("__all__", Store())], value=List( ctx=Load(), elts=[set_value(name)], expr=None, ), expr=None, lineno=None, **maybe_type_comment ) if not isinstance(gen_node, Module): gen_node = Module( body=list( chain.from_iterable( ( ( Expr( set_value( "\nGenerated from {module_name}.{name}\n".format( module_name=module_name, name=name_orig_ir[0], ) ) ), ), ast.parse(imports_header).body, (gen_node, __all___node), ) ) ), stmt=None, type_ignores=[], ) emit_filename, init_filepath = ( map( partial(path.join, output_directory, new_module_name), ( original_relative_filename_path, path.join( path.dirname(original_relative_filename_path), "__init__{extsep}py".format(extsep=extsep), ), ), ) if filesystem_layout == "as_input" else map( partial(path.join, mod_path), ( "{name}{extsep}py".format(name=name, extsep=extsep), "__init__{extsep}py".format(extsep=extsep), ), ) ) if path.isfile(emit_filename): with open(emit_filename, "rt") as f: mod = ast.parse(f.read()) gen_node = merge_modules(mod, gen_node) merge_assignment_lists(gen_node, "__all__") if dry_run: print("write\t{emit_filename!r}".format(emit_filename=emit_filename)) else: emit.file(gen_node, filename=emit_filename, mode="wt") if name != "__init__" and not path.isfile(init_filepath): if dry_run: print("write\t{emit_filename!r}".format(emit_filename=emit_filename)) else: emit.file( Module( body=[ Expr( set_value("\n__init__ to expose internals of this module\n") ), ImportFrom( module=name, names=[ alias( name=name, asname=None, identifier=None, identifier_name=None, ), ], level=1, identifier=None, ), __all___node, ], stmt=None, type_ignores=[], ), filename=init_filepath, mode="wt", ) return ( mod_name, original_relative_filename_path, ImportFrom( module=name, names=[ alias( name=name, asname=None, identifier=None, identifier_name=None, ), ], level=1, identifier=None, ), )
def create_op_call(self, op: Function, relay_args, py_args): """Lowers the passed primitive function, registers it in TVM's global compiler, and produces a call to the lowered function in the generated Python code.""" # compile the function and register globally cc_key = compile_engine.CCacheKey(op, self.tgt) func_hash = relay.analysis.structural_hash(op) op_name = '_lowered_op_{}'.format(func_hash) if not tvm.get_global_func(op_name, allow_missing=True): jitted = self.engine.jit(cc_key, self.tgt) tvm.register_func(op_name, jitted) def convert_input(py_input, arg_type): """Use the types of the function arguments to determine whether we expect a tensor or tuple (returns list of inputs to the lowered op call)""" # equivalent: input.data if isinstance(arg_type, relay.TensorType): return [py_input] assert isinstance(arg_type, relay.TupleType) # convert each input.fields[i] ret = [] for i in range(len(arg_type.fields)): ret += convert_input( ast.Subscript( py_input, ast.Index(Num(i)), Load()), arg_type.fields[i]) return ret def convert_output(ret_type): """Use the function return type to produce auxiliary variables to store outputs. Returns ([assignments of output vars], [extra arguments to pass to op call], expression collecting output)""" if isinstance(ret_type, relay.TensorType): output_var_name = self.generate_var_name('_out') output_var = Name(output_var_name, Load()) shape = ast.Tuple([Num(dim) for dim in ret_type.concrete_shape], Load()) # create a new NDArray of the right shape and dtype assign_output = Assign( [Name(output_var_name, Store())], self.create_call('nd.array', [ self.create_call('numpy.empty', [shape, Str(ret_type.dtype)]) ])) return ([assign_output], [output_var], output_var) assert isinstance(ret_type, relay.TupleType) assignments = [] extra_args = [] fields = [] for t in ret_type.fields: inner_assignments, inner_args, inner_output = convert_output(t) assignments += inner_assignments extra_args += inner_args fields.append(inner_output) fields = [ast.List(fields, Load())] return (assignments, extra_args, self.create_call('_container.tuple_object', fields)) # create a function to wrap the call of the lowered op and return # a call to that function wrap_name = self.generate_function_name('_{}_wrapper'.format(op_name)) wrap_args = [self.generate_var_name('_arg_{}'.format(i)) for i in range(len(py_args))] inner_call_args = [] for i in range(len(py_args)): inner_call_args += convert_input(Name(wrap_args[i], Load()), relay_args[i].checked_type) output_assignments, aux_args, output = convert_output(op.checked_type.ret_type) # equiv: _op = tvm.get_global_func(op_name) op_var = self.generate_var_name('_op') op_call = self.create_call('tvm.get_global_func', [Str(op_name)]) op_assign = Assign([Name(op_var, Store())], op_call) # equiv: _op(args) inner_call = self.create_call(op_var, inner_call_args + aux_args) body = output_assignments + [op_assign, ast.Expr(inner_call), Return(output)] wrap_def = self.create_def(wrap_name, wrap_args, body) return wrap_def, self.create_call(wrap_name, py_args)
def ast_index(name: str): return Index(value=Name(id=name, ctx=Load()))