def test_visit_match_multi_pattern(self):
fake_pos = self.fake_pos
fake_span = Span(fake_pos, fake_pos)
e = ast.Number(fake_span, u'0xf00')
p0 = ast.NameDefTree(fake_span, e)
p1 = ast.NameDefTree(fake_span, e)
arm = ast.MatchArm(patterns=(p0, p1), expr=e)
c = _Collector()
arm.accept(c)
self.assertEqual(c.collected, [e])
def _generate_expr(self, env: Env,
level: int) -> Tuple[ast.Expr, ast.TypeAnnotation]:
"""Generates an expression AST node and returns it."""
if self.should_nest(level):
identifier = self.gensym()
name_def = ast.NameDef(self.m, self.fake_span, identifier)
choices = collections.OrderedDict()
if self.options.emit_loops:
choices[self._generate_counted_for] = 1.0
choices[self._generate_tuple_or_index] = 1.0
choices[self._generate_concat] = 1.0
choices[self._generate_binop] = 1.0
choices[self._generate_unop] = 1.0
choices[self._generate_unop_builtin] = 1.0
choices[self._generate_one_hot_select_builtin] = 1.0
choices[self._generate_number] = 1.0
choices[self._generate_bit_slice] = 1.0
if not self._codegen_ops_only:
choices[self._generate_bitwise_reduction] = 1.0
choices[self._generate_cast_bits_to_array] = 1.0
# If maps recurse with equal probability, then the output will grow
# exponentially with level, so we need to scale inversely.
choices[lambda env: self._generate_map(level, env)] = 1.0 / (10**level)
rng_choices = getattr(self.rng, 'choices')
while True:
expr_generator = rng_choices(
population=list(choices.keys()),
weights=list(choices.values()),
k=1)[0]
try:
rhs, rhs_type = expr_generator(env)
except EmptyEnvError:
# Try a different generator that may be more accepting of env values.
del choices[expr_generator]
if not choices: # If we ran out of things to try, bail.
return self._generate_retval(env)
continue
else:
break
new_env = collections.OrderedDict(env)
new_env[identifier] = (
lambda: self._make_name_ref(name_def)), rhs_type, False
body, body_type = self._generate_expr(new_env, level + 1)
name_def_tree = ast.NameDefTree(self.m, self.fake_span, name_def)
let = ast.Let(
self.m,
name_def_tree,
rhs_type,
rhs,
body,
self.fake_span,
const=None)
return let, body_type
else: # Should not nest any more -- select return values.
return self._generate_retval(env)
def _generate_counted_for(self,
env: Env) -> Tuple[ast.For, ast.TypeAnnotation]:
"""Generates a counted for loop."""
# Right now just generates the 'identity' for loop.
ivar_type = self._make_type_annotation('u4')
zero = self._make_number(0, type_=ivar_type)
trips = self._make_number(self.randrange(8), type_=ivar_type)
iterable = self._make_range(zero, trips)
x_def = self._make_name_def('x')
name_def_tree = ast.NameDefTree(
self.fake_span,
tree=(ast.NameDefTree(self.fake_span, self._make_name_def('i')),
ast.NameDefTree(self.fake_span, x_def)))
make_expr, type_ = self._choose_env_value(env, self._not_array)
init = make_expr()
body = self._make_name_ref(x_def)
tree_type = self._make_tuple_type((ivar_type, type_))
return ast.For(self.fake_span, name_def_tree, tree_type, iterable,
body, init), type_
def test_ndt_preorder(self):
fake_pos = self.fake_pos
fake_span = Span(fake_pos, fake_pos)
t = ast.NameDef(fake_span, 't')
u = ast.NameDef(fake_span, 'u')
wrapped_t = ast.NameDefTree(fake_span, t)
wrapped_u = ast.NameDefTree(fake_span, u)
interior = ast.NameDefTree(fake_span, (wrapped_t, wrapped_u))
outer = ast.NameDefTree(fake_span, (interior, ))
walk_data = []
def walk(item: ast.NameDefTree, level: int, i: int):
walk_data.append((item, level, i))
outer.do_preorder(walk)
self.assertLen(walk_data, 3)
self.assertEqual(walk_data[0], (interior, 1, 0))
self.assertEqual(walk_data[1], (wrapped_t, 2, 0))
self.assertEqual(walk_data[2], (wrapped_u, 2, 1))