def _visit_ListComprehensionNode(self, stmt):
import numpy as np
result = self._visit(stmt.result)
generators = list(self._visit(stmt.generators))
lhs = self._visit(stmt.parent.target)
index = create_variable(lhs)
if isinstance(result, (Tuple, list, tuple)):
rank = len(np.shape(result))
else:
rank = 0
args = [Slice(None, None)] * rank
args.append(index)
target = IndexedBase(lhs)[args]
target = Assign(target, result)
assign1 = Assign(index, Integer(0))
assign1.set_fst(stmt)
target.set_fst(stmt)
generators[-1].insert2body(target)
assign2 = Assign(index, index + 1)
assign2.set_fst(stmt)
generators[-1].insert2body(assign2)
indices = [generators[-1].target]
while len(generators) > 1:
F = generators.pop()
generators[-1].insert2body(F)
indices.append(generators[-1].target)
indices = indices[::-1]
return FunctionalFor([assign1, generators[-1]], target.rhs, target.lhs,
indices, index)
def _visit_ListComp(self, stmt):
result = self._visit(stmt.elt)
generators = list(self._visit(stmt.generators))
if not isinstance(self._scope[-2], ast.Assign):
errors.report(PYCCEL_RESTRICTION_LIST_COMPREHENSION_ASSIGN,
symbol=stmt,
severity='error')
lhs = self.get_new_variable()
else:
lhs = self._visit(self._scope[-2].targets)
if len(lhs) == 1:
lhs = lhs[0]
else:
raise NotImplementedError(
"A list comprehension cannot be unpacked")
index = self.get_new_variable()
args = [index]
target = IndexedBase(lhs)[args]
target = Assign(target, result)
assign1 = Assign(index, LiteralInteger(0))
assign1.set_fst(stmt)
target.set_fst(stmt)
generators[-1].insert2body(target)
assign2 = Assign(index, PyccelAdd(index, LiteralInteger(1)))
assign2.set_fst(stmt)
generators[-1].insert2body(assign2)
indices = [generators[-1].target]
while len(generators) > 1:
F = generators.pop()
generators[-1].insert2body(F)
indices.append(generators[-1].target)
indices = indices[::-1]
return FunctionalFor([assign1, generators[-1]], target.rhs, target.lhs,
indices, index)
