def visit_dictcomp(self, node: _ast.DictComp): # key, value, generators
result: Dict[Any, Any] = {}
current_gen = node.generators[0]
if current_gen.__class__ == _ast.comprehension:
for val in self._run(current_gen.iter):
self.assign(current_gen.target, val)
add = True
for cond in current_gen.ifs:
add = add and self._run(cond)
if add:
if len(node.generators) > 1:
r = self.visit_dictcomp(
_ast.DictComp(
key=node.key,
value=node.value,
generators=node.generators[1:],
)
)
result.update(r)
else:
key = self._run(node.key)
value = self._run(node.value)
result[key] = value
self.delete(current_gen.target)
return result
def extract_dictcomp(self, function, sequence):
assert len(function.body) == 1
assert isinstance(function.body[0], _ast.Return)
value = function.body[0].value
assert isinstance(value, _ast.ListComp)
generators = list(value.generators)
for generator in generators:
if generator.iter.id == '.0':
generator.iter = sequence
setcomp = _ast.DictComp(key=value.elt[0], value=value.elt[1], generators=generators, lineno=value.lineno, col_offset=0)
self.push_ast_item(setcomp)
def DictComp(left_side_key, left_side_value, for_part, in_part, *ifs):
"""Creates _ast.DictComp nodes.
'left_side', 'left_side_value' for 'for_part' in 'in_part' if 'ifs'
Args:
left_side_key: key in leftmost side of the expression.
left_side_value: value in leftmost side of the expression.
for_part: The part after '[left_side] for '
in_part: The part after '[left_side] for [for_part] in '
*ifs: Any if statements that come at the end.
Returns:
{_ast.DictComp}
"""
left_side_key = _WrapWithName(left_side_key, ctx_type=CtxEnum.LOAD)
left_side_value = _WrapWithName(left_side_value, ctx_type=CtxEnum.LOAD)
for_part = _WrapWithName(for_part, ctx_type=CtxEnum.STORE)
in_part = _WrapWithName(in_part, ctx_type=CtxEnum.LOAD)
return _ast.DictComp(key=left_side_key,
value=left_side_value,
generators=[comprehension(for_part, in_part, *ifs)])
