Python ast.DictComp() Examples

def _collect_comprehension_children(self, node):
        # type: (ComprehensionNode) -> List[ast.expr]
        if isinstance(node, ast.DictComp):
            # dict comprehensions have two values to be checked
            child_nodes = [node.key, node.value]
        else:
            child_nodes = [node.elt]

        if node.generators:
            first_generator = node.generators[0]
            child_nodes.append(first_generator.target)
            for if_expr in first_generator.ifs:
                child_nodes.append(if_expr)

        for generator in node.generators[1:]:
            # rest need to be visited in the child scope
            child_nodes.append(generator.iter)
            child_nodes.append(generator.target)
            for if_expr in generator.ifs:
                child_nodes.append(if_expr)
        return child_nodes 

def ast_names(code):
    """Iterator that yields all the (ast) names in a Python expression.

    :arg code: A string containing a Python expression.
    """
    # Syntax that allows new name bindings to be introduced is tricky to
    # handle here, so we just refuse to do so.
    disallowed_ast_nodes = (ast.Lambda, ast.ListComp, ast.GeneratorExp)
    if sys.version_info >= (2, 7):
        disallowed_ast_nodes += (ast.DictComp, ast.SetComp)

    for node in ast.walk(ast.parse(code)):
        if isinstance(node, disallowed_ast_nodes):
            raise PatsyError("Lambda, list/dict/set comprehension, generator "
                             "expression in patsy formula not currently supported.")
        if isinstance(node, ast.Name):
            yield node.id 

def _should_instrument_as_expression(self, node):
        return (
            isinstance(node, _ast.expr)
            and hasattr(node, "end_lineno")
            and hasattr(node, "end_col_offset")
            and not getattr(node, "incorrect_range", False)
            and "ignore" not in node.tags
            and (not hasattr(node, "ctx") or isinstance(node.ctx, ast.Load))
            # TODO: repeatedly evaluated subexpressions of comprehensions
            # can be supported (but it requires some redesign both in backend and GUI)
            and "ListComp.elt" not in node.tags
            and "SetComp.elt" not in node.tags
            and "DictComp.key" not in node.tags
            and "DictComp.value" not in node.tags
            and "comprehension.if" not in node.tags
        ) 

def ast_names(code):
    """Iterator that yields all the (ast) names in a Python expression.

    :arg code: A string containing a Python expression.
    """
    # Syntax that allows new name bindings to be introduced is tricky to
    # handle here, so we just refuse to do so.
    disallowed_ast_nodes = (ast.Lambda, ast.ListComp, ast.GeneratorExp)
    if sys.version_info >= (2, 7):
        disallowed_ast_nodes += (ast.DictComp, ast.SetComp)

    for node in ast.walk(ast.parse(code)):
        if isinstance(node, disallowed_ast_nodes):
            raise PatsyError("Lambda, list/dict/set comprehension, generator "
                             "expression in patsy formula not currently supported.")
        if isinstance(node, ast.Name):
            yield node.id 

def _visit_dict_comp(self: 'ASTTagger', node: ast.DictComp):
    new = self.symtable.enter_new()
    new.entered.add('.0')
    new_tagger = ASTTagger(new)
    node.key = new_tagger.visit(node.key)
    node.value = new_tagger.visit(node.value)

    head, *tail = node.generators
    head.iter = self.visit(head.iter)
    head.target = new_tagger.visit(head.target)
    if head.ifs:
        head.ifs = [new_tagger.visit(each) for each in head.ifs]

    if any(each.is_async for each in node.generators):
        new.cts.add(ContextType.Coroutine)
    node.generators = [head, *[new_tagger.visit(each) for each in tail]]

    return Tag(node, new) 

def visit_comprehension(self, node):

        # elt is an expression generating elements
        # generators is a comprehension object

        if len(node.generators) != 1:
            raise NotSupported("Only one generator supported",
                               line=node.lineno)

        gen = node.generators[0]

        iter = self.visit(gen.iter)

        bounds = self.BOUND_VARS
        self.BOUND_VARS = []
        target = self.listofnames(gen.target)
        tlen = Const(str(len(target)))
        self.BOUND_VARS = bounds

        bound_len = len(target)
        self.BOUND_VARS = list(target) + self.BOUND_VARS

        if gen.ifs is None or len(gen.ifs) == 0:
            ifs = Const('True')
        elif len(gen.ifs) == 1:
            ifs = self.visit_expr(gen.ifs[0])
        else:
            raise NotSupported("Comprehension multiple ifs")

        if isinstance(node, ast.DictComp):
            key = self.visit_expr(node.key)
            value = self.visit_expr(node.value)
            op = Op(node.__class__.__name__, tlen, key, value, iter, ifs)
        else:
            elt = self.visit_expr(node.elt)
            op = Op(node.__class__.__name__, tlen, elt, iter, ifs)

        self.BOUND_VARS = self.BOUND_VARS[bound_len:]

        return op 

def visit_DictComp(self, node):
        # type: (ast.DictComp) -> None
        self._handle_comprehension(node, 'dictcomp') 

def test_dictcomp(self):
        g = ast.comprehension(ast.Name("y", ast.Store()),
                              ast.Name("p", ast.Load()), [])
        c = ast.DictComp(ast.Name("x", ast.Store()),
                         ast.Name("y", ast.Load()), [g])
        self.expr(c, "must have Load context")
        c = ast.DictComp(ast.Name("x", ast.Load()),
                         ast.Name("y", ast.Store()), [g])
        self.expr(c, "must have Load context")
        def factory(comps):
            k = ast.Name("x", ast.Load())
            v = ast.Name("y", ast.Load())
            return ast.DictComp(k, v, comps)
        self._check_comprehension(factory) 

def p_dictorsetmaker_1(p):
    '''dictorsetmaker : test COLON test comp_for'''
    #                      1     2    3        4
    p[0] = ast.DictComp(p[1], p[3], p[4], rule=inspect.currentframe().f_code.co_name)
    inherit_lineno(p[0], p[1]) 

def p_atom_7(p):
    '''atom : LBRACE dictorsetmaker RBRACE'''
    #              1              2      3
    if isinstance(p[2], (ast.SetComp, ast.DictComp)):
        p[0] = p[2]
        p[0].alt = p[1][1]
    else:
        keys, values = p[2]
        if keys is None:
            p[0] = ast.Set(values, rule=inspect.currentframe().f_code.co_name, **p[1][1])
        else:
            p[0] = ast.Dict(keys, values, rule=inspect.currentframe().f_code.co_name, **p[1][1]) 

def visit_DictComp(self, node):
        return self.visit_comprehension(node)
    
        # key = self.visit_expr(node.key)
        # value = self.visit_expr(node.value)
        
        # if len(node.generators) != 1:
        #     raise NotSupported("Only one generator supported",
        #                        line=node.lineno)
        
        # gen = self.visit_expr(node.generators[0])

        # return Op('DictComp', key, value, gen, line=node.lineno) 

def visit_DictComp(self, node: ast.DictComp) -> Any:
        """Compile the dictionary comprehension as a function and call it."""
        result = self._execute_comprehension(node=node)

        for generator in node.generators:
            self.visit(generator.iter)

        self.recomputed_values[node] = result
        return result 

def _execute_comprehension(self, node: Union[ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp]) -> Any:
        """Compile the generator or comprehension from the node and execute the compiled code."""
        args = [ast.arg(arg=name) for name in sorted(self._name_to_value.keys())]

        if platform.python_version_tuple() < ('3', ):
            raise NotImplementedError("Python versions below not supported, got: {}".format(platform.python_version()))

        if platform.python_version_tuple() < ('3', '8'):
            func_def_node = ast.FunctionDef(
                name="generator_expr",
                args=ast.arguments(args=args, kwonlyargs=[], kw_defaults=[], defaults=[]),
                decorator_list=[],
                body=[ast.Return(node)])

            module_node = ast.Module(body=[func_def_node])
        else:
            func_def_node = ast.FunctionDef(
                name="generator_expr",
                args=ast.arguments(args=args, posonlyargs=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
                decorator_list=[],
                body=[ast.Return(node)])

            module_node = ast.Module(body=[func_def_node], type_ignores=[])

        ast.fix_missing_locations(module_node)

        code = compile(source=module_node, filename='<ast>', mode='exec')

        module_locals = {}  # type: Dict[str, Any]
        module_globals = {}  # type: Dict[str, Any]
        exec(code, module_globals, module_locals)  # pylint: disable=exec-used

        generator_expr_func = module_locals["generator_expr"]

        return generator_expr_func(**self._name_to_value) 

def visit_DictComp(self, node: ast.DictComp) -> None:
        """Represent the dictionary comprehension by dumping its source code."""
        if node in self._recomputed_values:
            value = self._recomputed_values[node]
            text = self._atok.get_text(node)

            self.reprs[text] = value

        self.generic_visit(node=node) 

def test_dictcomp(self):
        g = ast.comprehension(ast.Name("y", ast.Store()),
                              ast.Name("p", ast.Load()), [])
        c = ast.DictComp(ast.Name("x", ast.Store()),
                         ast.Name("y", ast.Load()), [g])
        self.expr(c, "must have Load context")
        c = ast.DictComp(ast.Name("x", ast.Load()),
                         ast.Name("y", ast.Store()), [g])
        self.expr(c, "must have Load context")
        def factory(comps):
            k = ast.Name("x", ast.Load())
            v = ast.Name("y", ast.Load())
            return ast.DictComp(k, v, comps)
        self._check_comprehension(factory) 

def visit_DictComp(self, node: ast.DictComp) -> None:
        # This is mostly because there is no reason to use it at the moment. It might also be an attack vector.
        raise IllegalMethodError("Dict Comprehension is not allowed") 

def test_dictcomp(self):
        g = ast.comprehension(ast.Name("y", ast.Store()),
                              ast.Name("p", ast.Load()), [])
        c = ast.DictComp(ast.Name("x", ast.Store()),
                         ast.Name("y", ast.Load()), [g])
        self.expr(c, "must have Load context")
        c = ast.DictComp(ast.Name("x", ast.Load()),
                         ast.Name("y", ast.Store()), [g])
        self.expr(c, "must have Load context")
        def factory(comps):
            k = ast.Name("x", ast.Load())
            v = ast.Name("y", ast.Load())
            return ast.DictComp(k, v, comps)
        self._check_comprehension(factory) 

def loops(node):
    # type: (ast.AST) -> Tuple[Loop, ...]
    """
    Return all the 'enclosing loops' of a node, up to the innermost class or
    function definition. This also includes the 'for in' clauses in list/dict/set/generator
    comprehensions. So for example, in this code:

      for x in ...:
          def foo():
              while True:
                  print([z for y in ...])

    The loops enclosing the node 'z' are 'while True' and 'for y in ...', in that order.
    """
    result = []
    while True:
        try:
            parent = node.parent
        except AttributeError:
            break
        if isinstance(parent, ast.FunctionDef):
            break

        is_containing_loop = (((isinstance(parent, ast.For) and parent.iter is not node or
                                isinstance(parent, ast.While))
                               and node not in parent.orelse) or
                              (isinstance(parent, ast.comprehension) and node in parent.ifs))
        if is_containing_loop:
            result.append(parent)

        elif isinstance(parent, (ast.ListComp,
                                 ast.GeneratorExp,
                                 ast.DictComp,
                                 ast.SetComp)):
            generators = parent.generators
            if node in generators:
                generators = list(takewhile(lambda n: n != node, generators))
            result.extend(reversed(generators))

        node = parent

    result.reverse()
    return tuple(result) 

def _process_args(self, func_ast, code_lines, args, kwargs) -> 'Generator[DebugArgument, None, None]':  # noqa: C901
        import ast

        complex_nodes = (
            ast.Call,
            ast.Attribute,
            ast.Subscript,
            ast.IfExp,
            ast.BoolOp,
            ast.BinOp,
            ast.Compare,
            ast.DictComp,
            ast.ListComp,
            ast.SetComp,
            ast.GeneratorExp,
        )

        arg_offsets = list(self._get_offsets(func_ast))
        for i, arg in enumerate(args):
            try:
                ast_node = func_ast.args[i]
            except IndexError:  # pragma: no cover
                # happens when code has been commented out and there are fewer func_ast args than real args
                yield self.output_class.arg_class(arg)
                continue

            if isinstance(ast_node, ast.Name):
                yield self.output_class.arg_class(arg, name=ast_node.id)
            elif isinstance(ast_node, complex_nodes):
                # TODO replace this hack with astor when it get's round to a new release
                start_line, start_col = arg_offsets[i]

                if i + 1 < len(arg_offsets):
                    end_line, end_col = arg_offsets[i + 1]
                else:
                    end_line, end_col = len(code_lines) - 1, None

                name_lines = []
                for l_ in range(start_line, end_line + 1):
                    start_ = start_col if l_ == start_line else 0
                    end_ = end_col if l_ == end_line else None
                    name_lines.append(code_lines[l_][start_:end_].strip(' '))
                yield self.output_class.arg_class(arg, name=' '.join(name_lines).strip(' ,'))
            else:
                yield self.output_class.arg_class(arg)

        kw_arg_names = {}
        for kw in func_ast.keywords:
            if isinstance(kw.value, ast.Name):
                kw_arg_names[kw.arg] = kw.value.id
        for name, value in kwargs.items():
            yield self.output_class.arg_class(value, name=name, variable=kw_arg_names.get(name))