def load_helpers(self) -> Iterable[stmt]:
helpers = []
if self.track_import_call:
helpers.append(
lineinfo(
make_assign(
[lineinfo(ast.Name("<strict-modules>", ast.Store()))],
lineinfo(
ast.Subscript(
lineinfo(ast.Name("<fixed-modules>", ast.Load())),
lineinfo(ast.Index(lineinfo(Constant("__strict__")))),
ast.Load(),
)
),
)
)
)
helpers.append(
lineinfo(
make_assign(
[lineinfo(ast.Name("<track-import-call>", ast.Store()))],
lineinfo(
ast.Subscript(
lineinfo(ast.Name("<strict-modules>", ast.Load())),
lineinfo(
ast.Index(lineinfo(Constant("track_import_call")))
),
ast.Load(),
)
),
)
),
)
return helpers
def test_Subscript(self):
sub = ast.Subscript(ast.Name('X', ast.Load()), [], ast.Load())
# Index
slice1 = ast.Index(ast.Num(42))
sub.slice = slice1
self.verify(sub, 'X[42]')
# Slice
slice2 = ast.Slice(None, None, ast.Num(2))
sub.slice = slice2
self.verify(sub, 'X[::2]')
# ExtSlice
sub.slice = ast.ExtSlice([slice1, slice2])
self.verify(sub, 'X[42,::2]')
# Issue #20
expect = ast.Expr(value=ast.Subscript(
value=ast.BinOp(left=ast.Name(id='p', ctx=ast.Load()),
op=ast.Add(),
right=ast.List(elts=[
ast.Num(n=1),
ast.Num(n=2),
ast.Num(n=3),
ast.Num(n=4)
],
ctx=ast.Load())),
slice=ast.Slice(lower=None, upper=ast.Num(n=4), step=None),
ctx=ast.Load()))
self.verify(expect, '(p+[1,2,3,4])[:4]')
def generate_io_stmt(input_idx, var_name, value, func_name):
if isinstance(value, list):
return [ast.Expr(
ast.Call(
ast.Attribute(
ast.Subscript(
ast.Name(var_name, ast.Load()),
ast.Index(
ast.Tuple([ast.Num(input_idx),
ast.Num(val_idx)],
ast.Load())),
ast.Load()),
func_name,
ast.Load()),
[ast.Num(val)],
[], None, None))
for val_idx, val in enumerate(value)]
else:
return [ast.Expr(
ast.Call(
ast.Attribute(
ast.Subscript(
ast.Name(var_name, ast.Load()),
ast.Index(ast.Num(input_idx)),
ast.Load()),
func_name,
ast.Load()),
[ast.Num(value)],
[], None, None))]
def construct_python_args(self):
new_args = []
for i in range(len(self.itypes)):
new_annotation = None
if (self.itypes[i] == 'VectorI'):
new_annotation = ast.Subscript(value=ast.Name(id="List",
ctx=ast.Load()),
slice=ast.Name(id='int',
ctx=ast.Load()))
if (self.itypes[i] == 'VectorF'):
new_annotation = ast.Subscript(value=ast.Name(id="List",
ctx=ast.Load()),
slice=ast.Name(id='float',
ctx=ast.Load()))
if (self.itypes[i] == 'int'):
new_annotation = ast.Name(id='int', cts=ast.Load())
if (self.itypes[i] == 'float'):
new_annotation = ast.Name(id='float', cts=ast.Load())
if (self.itypes[i] == 'MatrixI'):
pass
if (self.itypes[i] == 'MatrixF'):
pass
new_args.append(ast.arg(arg=f'arg_{i}', annotation=new_annotation))
return ast.arguments(args=new_args,
defaults=[],
vararg=None,
kwarg=None)
def _get_annotation_for_property(self, prop: raml_types.Property):
"""Create an node which represents an annotation for a property"""
if prop.type is None:
annotation_type = BUILTIN_TYPES["any"]
elif prop.type.name in BUILTIN_TYPES:
annotation_type = BUILTIN_TYPES[prop.type.name]
elif prop.type.base and prop.type.base.name == "string" and not prop.type.enum:
annotation_type = ast.Str(s="str")
else:
if self.resource.package_name != prop.type.package_name:
self.generator.import_resource_typing(
self.resource.package_name, prop.type.package_name,
prop.type.name)
annotation_type = ast.Str(s=prop.type.name)
# use typing.List[]. We make an hardcoded exception for
# resources ending on PagedQueryResponse and mark that as Sequence. The
# reason is to keep mypy happy when we pass subclasses of Resource to
# PagedQueryResponse object
if prop.many:
attr = "List"
if self.resource.name.endswith("PagedQueryResponse"):
attr = "Sequence"
annotation_type = ast.Subscript(
value=ast.Attribute(value=ast.Name(id="typing"), attr=attr),
slice=ast.Index(value=annotation_type),
)
# Wrap it in a typing.Optional[T]
annotation_type = ast.Subscript(
value=ast.Attribute(value=ast.Name(id="typing"), attr="Optional"),
slice=ast.Index(value=annotation_type),
)
return annotation_type
def visitSubscript(self, n, *args):
value = self.dispatch(n.value, *args)
slice = self.dispatch(n.slice, *args)
if isinstance(n.ctx, ast.Load):
responsible, assign = generateTemp(value)
sub = ast.Subscript(value=responsible,
slice=slice,
ctx=n.ctx,
lineno=n.lineno,
col_offset=n.col_offset)
return retic_ast.Flattened([assign],
retic_ast.BlameCheck(
value=sub,
type=n.retic_type,
responsible=responsible,
tag=retic_ast.GetItem(),
lineno=n.lineno,
col_offset=n.col_offset),
lineno=n.lineno,
col_offset=n.col_offset)
else:
return ast.Subscript(value=value,
slice=slice,
ctx=n.ctx,
lineno=n.lineno,
col_offset=n.col_offset)
def visit_comprehension(self, node):
self.generic_visit(node)
assignments = flat_map_assignments([node.target], depth=1)
declarations = []
elts = []
for name in assignments:
new_name = ast.Name(id=name.id, ctx=ast.Load())
dec_details = self.get_assignment_details(name, True)
assn_details = self.get_assignment_details(name, True)
declarations.append(
self.wrapper(ast.NameConstant(None), dec_details))
elts.append(self.wrapper(
new_name,
assn_details,
))
elts.append(ast.NameConstant(True))
declarations.append(node.iter)
dec_list_node = ast.List(elts=declarations)
node.iter = ast.Subscript(value=dec_list_node,
slice=ast.Index(value=ast.Num(-1)))
assign_list_node = ast.List(elts=elts)
node.ifs.insert(
0,
ast.Subscript(value=assign_list_node,
slice=ast.Index(value=ast.Num(-1))))
return node
def _generate_set_call_body(self, external_function_name, arg_length, strart_idx, target_ids):
arg_ids = np.arange(arg_length)
targets = []
args = []
for target_id in target_ids:
targets.append(ast.Subscript(value=GLOBAL_ARRAY,
slice=ast.Index(
value=ast.BinOp(
left=ast.Constant(value=target_id + strart_idx),
op=ast.Add(),
right=ast.BinOp(left=ast.Constant(value=arg_length, kind=None),
op=ast.Mult(),
right=ast.Name(id='i', ctx=ast.Load())))),
ctx=ast.Store()))
for arg_id in arg_ids:
args.append(
ast.Subscript(value=GLOBAL_ARRAY, slice=ast.Index(value=ast.BinOp(
left=ast.Constant(value=arg_id + strart_idx),
op=ast.Add(),
right=ast.BinOp(left=ast.Constant(value=arg_length, kind=None),
op=ast.Mult(),
right=ast.Name(id='i', ctx=ast.Load())))),
ctx=ast.Load))
if len(targets) > 1:
return ast.Assign(targets=[ast.Tuple(elts=targets)],
value=ast.Call(func=ast.Name(id=external_function_name, ctx=ast.Load()),
args=args, keywords=[]), lineno=0)
else:
return ast.Assign(targets=[targets[0]],
value=ast.Call(func=ast.Name(id=external_function_name, ctx=ast.Load()),
args=args, keywords=[]), lineno=0)
def add_call(self, external_function_name, args, target_ids):
arg_ids = np.arange(len(args))
start_idx = self.variable_names[args[0]]
targets = []
args = []
for target_id in target_ids:
targets.append(
ast.Subscript(value=GLOBAL_ARRAY,
slice=ast.Index(value=ast.Constant(
value=target_id + start_idx)),
ctx=ast.Store()))
for arg_id in arg_ids:
args.append(
ast.Subscript(value=GLOBAL_ARRAY,
slice=ast.Index(value=ast.Constant(value=arg_id +
start_idx)),
ctx=ast.Load))
if len(targets) > 1:
self.body.append(
ast.Assign(targets=[ast.Tuple(elts=targets)],
value=ast.Call(func=ast.Name(
id=external_function_name, ctx=ast.Load()),
args=args,
keywords=[]),
lineno=0))
else:
self.body.append(
ast.Assign(targets=[targets[0]],
value=ast.Call(func=ast.Name(
id=external_function_name, ctx=ast.Load()),
args=args,
keywords=[]),
lineno=0))
def add_mapping(self, args, target):
if len(target) > 1:
raise ValueError("Only mapping to single target is supported")
arg_idxs = []
for arg in args:
arg_idxs.append(self.variable_names[arg])
target_idx = self.variable_names[target[0]]
print(target_idx)
if len(args) == 1:
self.body.append(ast.Assign(targets=[ast.Subscript(value=GLOBAL_ARRAY,
slice=ast.Index(
value=ast.Constant(value=target_idx, kind=None)))],
value=ast.Subscript(value=GLOBAL_ARRAY,
slice=ast.Index(
value=ast.Constant(value=arg_idxs[0], kind=None)))
, lineno=0))
else:
self.body.append(ast.Assign(targets=[ast.Subscript(value=GLOBAL_ARRAY,
slice=ast.Index(
value=ast.Constant(value=target_idx, kind=None)))],
value=ast.BinOp(left=self._generate_sum_left(arg_idxs[1:]), op=ast.Add(),
right=ast.Subscript(value=GLOBAL_ARRAY,
slice=ast.Index(
value=ast.Constant(value=arg_idxs[0],
kind=None))))
, lineno=0))
def visit_For(self, node):
super().generic_visit(node)
# TODO: Support attributes
if isinstance(node.iter, ast.Name) or ((isinstance(node.iter, ast.List) or isinstance(node.iter, ast.Tuple)) and not (all(isinstance(item, ast.Num) for item in node.iter.elts) or all(isinstance(item, ast.Str) for item in node.iter.elts))):
targets = node.target.elts if hasattr(node.target, 'elts') else [node.target]
indexator = self.get_unique_iter_target()
set_targets = [
ast.Assign(
targets=[t],
value=ast.Subscript(
value=ast.Subscript(value=node.iter, slice=ast.Index(value=ast.Name(id=indexator))),
slice=ast.Index(value=ast.Num(n=i))
)
) for i, t in enumerate(targets)
]
return ast.For(
target=indexator,
iter=ast.Call(
func=ast.Name(id='range'),
args=[ast.Call(
func=ast.Name(id='len'),
args=node.iter
)]
),
body=set_targets + node.body
)
return node
def visit_For(self, node):
if not isinstance(node.target, ast.Name):
self.error(node.target,
"Only assignment to target names is supported.")
# NOTE: this would be easier to do during type inference, since you
# can rewrite the AST and run the type inferer on the rewrite.
# This way, you won't have to create variables and types manually.
# This will also take care of proper coercions.
if node.iter.type.is_range:
# make sure to analyse children, in case of nested loops
self.generic_visit(node)
node.index = ast.Name(id=node.target.id, ctx=ast.Load())
return node
elif node.iter.type.is_array and node.iter.type.ndim == 1:
# Convert 1D array iteration to for-range and indexing
logger.debug(ast.dump(node))
orig_target = node.target
orig_iter = node.iter
# replace node.target with a temporary
target_name = orig_target.id + '.idx'
target_temp = nodes.TempNode(minitypes.Py_ssize_t)
node.target = target_temp.store()
# replace node.iter
call_func = ast.Name(id='range', ctx=ast.Load())
call_func.variable = Variable(_types.RangeType())
shape_index = ast.Index(nodes.ConstNode(0, _types.Py_ssize_t))
stop = ast.Subscript(value=nodes.ShapeAttributeNode(orig_iter),
slice=shape_index,
ctx=ast.Load())
stop.type = _types.intp
call_args = [
nodes.ConstNode(0, _types.Py_ssize_t),
nodes.CoercionNode(stop, _types.Py_ssize_t),
nodes.ConstNode(1, _types.Py_ssize_t),
]
node.iter = ast.Call(func=call_func, args=call_args)
node.iter.type = call_func.variable.type
node.index = target_temp.load()
# add assignment to new target variable at the start of the body
index = ast.Index(value=node.index)
index.variable = target_temp.variable
subscript = ast.Subscript(value=orig_iter,
slice=index,
ctx=ast.Load())
coercion = nodes.CoercionNode(subscript, orig_target.type)
assign = ast.Assign(targets=[orig_target], value=subscript)
node.body = [assign] + node.body
return node
else:
raise error.NumbaError("Unsupported for loop pattern")
def p_list_range(p):
'''list : LBRACKET atom COMMA atom RANGE atom RBRACKET
| LBRACKET atom RANGE atom RBRACKET'''
if len(p) == 8:
p[0] = ast.Subscript(
None, ast.Slice(p[2], p[6], ast.BinOp(p[4], ast.Sub, p[2])),
ast.Load())
else:
p[0] = ast.Subscript(None, ast.Slice(p[2], p[4], ast.Num(1)),
ast.Load())
def visit_Call(self, call):
if len(call.args) > 0 and isinstance(call.args[-1], ast.Starred):
star_arg = call.args[-1].value
try:
star_obj = eval(a2s(star_arg), self.globls, self.globls)
except Exception:
# print('ERROR', a2s(call))
# TODO: log a warning?
# list comprehension issue: [f(*c) for c in whatever]
return call
call.args.pop()
self.change = True
for i in range(len(star_obj)):
call.args.append(
ast.Subscript(value=star_arg, slice=ast.Index(ast.Num(i))))
kwarg = [(i, kw.value) for i, kw in enumerate(call.keywords)
if kw.arg is None]
if len(kwarg) == 1:
i, kwarg = kwarg[0]
try:
kwarg_obj = eval(a2s(kwarg), self.globls, self.globls)
except Exception:
print('ERROR', a2s(call))
raise
if isinstance(kwarg, ast.Name):
# HACK: in case where kwarg is reused and modified,
# e.g. kw = {}; foo(**kw); kw['x'] = 1; bar(**kw);
# then our compilation strategy is unsound. For now,
# just check if **kw is used twice
if kwarg.id in self.inlined_kwarg:
raise CancelPass
self.inlined_kwarg.add(kwarg.id)
self.change = True
del call.keywords[i]
# Add a keyword argument for each key in kwargs, e.g.
# kwargs = {'x': 1}; f(**kwargs) => f(x=kwargs['x'])
for k in kwarg_obj.keys():
idx_expr = ast.Subscript(value=kwarg,
slice=ast.Index(ast.Str(k)))
call.keywords.append(ast.keyword(arg=k, value=idx_expr))
self.generic_visit(call)
return call
def abi_type_to_ast(atype):
if atype in ('int128', 'uint256', 'bool', 'address', 'bytes32'):
return ast.Name(atype, None)
elif atype == 'decimal':
return ast.Name('int128', None)
elif atype == 'bytes':
return ast.Subscript(value=ast.Name('bytes', None),
slice=ast.Index(256))
elif atype == 'string':
return ast.Subscript(value=ast.Name('string', None),
slice=ast.Index(256))
else:
raise ParserException('Type {} not supported by vyper.'.format(atype))
def parametrize(exp):
if isinstance(exp, ast.Index):
return lambda _: {exp}
elif isinstance(exp, ast.Name):
try:
w = node.args.args.index(exp)
def return_alias(args):
if w < len(args):
return {args[w]}
else:
return {node.args.defaults[w - len(args)]}
return return_alias
except ValueError:
return lambda _: {NewMem}
elif isinstance(exp, ast.Subscript):
values = parametrize(exp.value)
slices = parametrize(exp.slice)
return lambda args: {
ast.Subscript(value, slice, ast.Load())
for value in values(args) for slice in slices(args)
}
else:
return lambda _: {NewMem}
def _create_regex_call(self, field_name, method_name):
"""Generate `data = self.fields[{ field_name }].{ method_name }(data)`"""
self.generator.add_import_statement(self.resource.package_name,
"typing")
return ast.Assign(
targets=[ast.Name(id="data")],
value=ast.Call(
func=ast.Attribute(
value=ast.Call(
func=ast.Attribute(value=ast.Name(id="typing"),
attr="cast"),
args=[
ast.Attribute(value=ast.Name(id="helpers"),
attr="RegexField"),
ast.Subscript(
value=ast.Attribute(value=ast.Name(id="self"),
attr="fields"),
slice=ast.Index(
value=ast.Str(s=field_name, kind=None)),
),
],
keywords=[],
),
attr=method_name,
),
args=[ast.Name(id="data")],
keywords=[],
),
)
def expr__subscript(p):
src, _, subscript, _ = p
return ast.Subscript(
value=src,
slice=ast.Index(value=subscript, ctx=ast.Load()),
ctx=ast.Load(),
)
def copy_assignee(n, ctx):
if isinstance(n, ast.Name):
ret = ast.Name(id=n.id, ctx=ctx)
elif isinstance(n, ast.Attribute):
ret = ast.Attribute(value=n.value, attr=n.attr, ctx=ctx)
elif isinstance(n, ast.Subscript):
ret = ast.Subscript(value=n.value, slice=n.slice, ctx=ctx)
elif isinstance(n, ast.List):
elts = [copy_assignee(e, ctx) for e in n.elts]
ret = ast.List(elts=elts, ctx=ctx)
elif isinstance(n, ast.Tuple):
elts = [copy_assignee(e, ctx) for e in n.elts]
ret = ast.Tuple(elts=elts, ctx=ctx)
elif isinstance(n, ast.Starred):
ret = ast.Starred(value=copy_assignee(n.value, ctx), ctx=ctx)
elif isinstance(n, ast.Call):
args = [copy_assignee(e, ctx) for e in n.args]
ret = ast.Call(func=n.func,
args=args,
keywords=n.keywords,
starargs=n.starargs,
kwargs=n.kwargs)
else:
return n
ast.copy_location(ret, n)
return ret
def visit_Assign(self, node):
self.generic_visit(node)
extra_assign = [node]
for i, t in enumerate(node.targets):
if isinstance(t, ast.Tuple) or isinstance(t, ast.List):
renamings = dict()
self.traverse_tuples(t, (), renamings)
if renamings:
self.counter += 1
gtarget = "{0}{1}{2}".format(
NormalizeTuples.tuple_name,
self.counter,
i)
node.targets[i] = ast.Name(gtarget, node.targets[i].ctx)
for rename, state in sorted(renamings.iteritems()):
nnode = reduce(
lambda x, y: ast.Subscript(
x,
ast.Index(ast.Num(y)),
ast.Load()),
state,
ast.Name(gtarget, ast.Load()))
if isinstance(rename, str):
extra_assign.append(
ast.Assign(
[ast.Name(rename, ast.Store())],
nnode))
else:
extra_assign.append(ast.Assign([rename], nnode))
return (ast.If(ast.Num(1), extra_assign, [])
if len(extra_assign) > 1
else extra_assign)
def visit_For(self, node):
target = node.target
if isinstance(target, ast.Tuple) or isinstance(target, ast.List):
renamings = dict()
self.traverse_tuples(target, (), renamings)
if renamings:
self.counter += 1
gtarget = "{0}{1}".format(
NormalizeTuples.tuple_name,
self.counter
)
node.target = ast.Name(gtarget, node.target.ctx)
for rename, state in sorted(renamings.iteritems()):
nnode = reduce(
lambda x, y: ast.Subscript(
x,
ast.Index(ast.Num(y)),
ast.Load()),
state,
ast.Name(gtarget, ast.Load()))
if isinstance(rename, str):
node.body.insert(0,
ast.Assign(
[ast.Name(rename, ast.Store())],
nnode)
)
else:
node.body.insert(0, ast.Assign([rename], nnode))
self.generic_visit(node)
return node
def createAst(self, attrOfAst):
if not self.closed:
raise icon.IconExecException(self, "Unclosed temporary icon")
if self.sites.indexIcon.att is None:
if not self.hasSite('upperIcon'):
raise icon.IconExecException(self, "Missing subscript")
indexAst = None
else:
indexAst = self.sites.indexIcon.att.createAst()
if self.hasSite('upperIcon'):
if self.sites.upperIcon.att:
upperAst = self.sites.upperIcon.att.createAst()
else:
upperAst = None
if self.hasSite('stepIcon') and self.sites.stepIcon.att:
stepAst = self.sites.stepIcon.att.createAst()
else:
stepAst = None
slice = ast.Slice(indexAst, upperAst, stepAst)
else:
slice = ast.Index(value=indexAst)
return icon.composeAttrAst(
self,
ast.Subscript(value=attrOfAst,
slice=slice,
lineno=self.id,
col_offset=0,
ctx=nameicons.determineCtx(self)))
def format_target(self, var, read):
if read:
_ctx = ast.Load()
else:
_ctx = ast.Store()
if self.llvm:
return ast.Subscript(
slice=ast.Index(value=ast.Num(n=0, lineno=0, col_offset=0),
lineno=0,
col_offset=0),
value=ast.Call(args=[
ast.Name(id=var.replace('scope.', 's_'),
lineno=0,
col_offset=0,
ctx=ast.Load()),
ast.Tuple(ctx=ast.Load(),
elts=[ast.Num(n=1, lineno=0, col_offset=0)],
lineno=0,
col_offset=0)
],
func=ast.Name(id='carray',
lineno=0,
col_offset=0,
ctx=ast.Load()),
keywords=[],
lineno=0,
col_offset=0),
lineno=0,
col_offset=0,
ctx=_ctx)
else:
return ast.Name(id=var.replace('scope.', 's_'),
lineno=0,
col_offset=0,
ctx=_ctx)
def subscript_action(s, loc, tokens):
if 'slice' in tokens:
slice = tokens.slice
start = slice.start if 'start' in slice else None
stop = slice.stop if 'stop' in slice else None
if 'step' in slice:
step = slice.step[0]
else:
step = None
return ast.Subscript(slice=ast.Slice(lower=start,
upper=stop,
step=step),
ctx=ast.Load())
else:
return ast.Subscript(slice=ast.Index(value=tokens[0]), ctx=ast.Load())
def _annotate_literal(value):
if value is None:
return ast.Name(id="None")
literal = _typing("Literal")
literal_slice = ast.Index(value=ast.Constant(value=value, kind=None))
literal_subscript = ast.Subscript(value=literal, slice=literal_slice)
return literal_subscript
def visit_DictComp_Rec(
self, generators: List[Type[ast.AST]]) -> List[Type[ast.AST]]:
if not generators:
if self.dictCompReg[
0]: # bug if there is else statement in comprehension
return [
ast.Assign(targets=[
ast.Subscript(
value=ast.Name(id=self.dictCompReg[0],
ctx=ast.Load()),
slice=ast.Index(value=self.dictCompReg[1].key),
ctx=ast.Store())
],
value=self.dictCompReg[1].value)
]
# else: # not supported yet
# return [ast.Expr(value=self.dictCompReg[1].elt)]
else:
return [
ast.For(
target=generators[-1].target,
iter=generators[-1].iter,
body=[
ast.If(test=self.combine_conditions(
generators[-1].ifs),
body=self.visit_DictComp_Rec(generators[:-1]),
orelse=[])
] if generators[-1].ifs else self.visit_DictComp_Rec(
generators[:-1]),
orelse=[])
]
def generate(self, element:Element, GC:GenerationContext):
acode = element.code
sig = acode[1].code
mac_name = sig[0].code.full_name
param_names = [p.code.full_name for p in sig[1:]]
body = [s.code for s in acode[2:]]
params_code = ast.List(elts=[ast.Str(s) for s in param_names],
ctx=ast.Load())
# rcode_code =
# __aky__.Macro( ... )
mac_code = ast.Call(func=ast.Attribute(value=ast.Name(id="__aky__", ctx=ast.Load()),
attr="Macro", ctx=ast.Load()),
args=[],
keywords=[ast.keyword(arg="specs",
value=ast.Tuple(elts=[params_code,
rcode_code],
ctx=ast.Load()))])
target_code = ast.Subscript(value=ast.Name(id="__macros__", ctx=ast.Load()),
slice=ast.Index(ast.Str(mac_name)),
ctx=ast.Store())
#
assign_code = ast.Assign(targets=[target_code],
value=mac_code)
return assign_code
def _annotation_optional(self, type_, optional=False):
if optional:
return ast.Subscript(
value=ast.Name(id="Optional"), slice=ast.Index(value=ast.Name(id=type_))
)
else:
return ast.Name(id=type_)
def __getitem__(
self, expr: Union[Callable, DataFrame, Column, str,
int]) -> DataFrame:
'''A filtering operation of some sort or a branch look up or a slice'''
assert isinstance(expr, (DataFrame, Column, int, str)) or callable(expr), \
"Filtering a data frame must be done by a DataFrame expression " \
f"(type: DataFrame or Column or int) not '{type(expr).__name__}'"
# Index into an item
if isinstance(expr, int):
c_expr = ast.Subscript(value=ast_DataFrame(self),
slice=ast.Index(value=expr))
return DataFrame(expr=c_expr)
# A branch look up - like a ".pt" rather than ['pt']
if isinstance(expr, str):
return self.__getattr__(expr)
if callable(expr) and not (isinstance(expr, DataFrame)
or isinstance(expr, Column)):
c_expr = expr(self)
assert isinstance(c_expr, DataFrame) or isinstance(c_expr, Column), \
f"Filter function '{expr.__name__}'' did not return a DataFrame expression"
expr = c_expr
if isinstance(expr, DataFrame):
assert expr.filter is None
assert expr.child_expr is not None
expr = Column(bool, expr.child_expr)
# Redundant, but above too complex for type processor?
assert isinstance(
expr, Column), 'Internal error - filter must be a bool column!'
return DataFrame(ast_DataFrame(self), filter=expr)
def visit_Assign(self, node):
self.generic_visit(node)
extra_assign = [node]
for i, t in enumerate(node.targets):
if isinstance(t, ast.Tuple) or isinstance(t, ast.List):
renamings = dict()
self.traverse_tuples(t, (), renamings)
if renamings:
gtarget = self.get_new_id()
node.targets[i] = ast.Name(gtarget, node.targets[i].ctx)
for rename, state in sorted(renamings.iteritems()):
nnode = reduce(
lambda x, y: ast.Subscript(
x,
ast.Index(ast.Num(y)),
ast.Load()),
state,
ast.Name(gtarget, ast.Load()))
if isinstance(rename, str):
extra_assign.append(
ast.Assign(
[ast.Name(rename, ast.Store())],
nnode))
else:
extra_assign.append(ast.Assign([rename], nnode))
return extra_assign
