Python ast.ListComp() Examples

def _visit_list_set_gen_comp(self: 'ASTTagger', node: ast.ListComp):
    new = self.symtable.enter_new()
    new.entered.add('.0')
    new_tagger = ASTTagger(new)
    node.elt = new_tagger.visit(node.elt)
    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 _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_Call(self, node: ast.Call) -> None:
        if (
                isinstance(node.func, ast.Name) and
                node.func.id == 'set' and
                len(node.args) == 1 and
                not node.keywords and
                isinstance(node.args[0], SET_TRANSFORM)
        ):
            arg, = node.args
            key = _ast_to_offset(node.func)
            if isinstance(arg, (ast.List, ast.Tuple)) and not arg.elts:
                self.set_empty_literals[key] = arg
            else:
                self.sets[key] = arg
        elif (
                isinstance(node.func, ast.Name) and
                node.func.id == 'dict' and
                len(node.args) == 1 and
                not node.keywords and
                isinstance(node.args[0], (ast.ListComp, ast.GeneratorExp)) and
                isinstance(node.args[0].elt, (ast.Tuple, ast.List)) and
                len(node.args[0].elt.elts) == 2
        ):
            self.dicts[_ast_to_offset(node.func)] = node.args[0]
        self.generic_visit(node) 

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_SetComp(self, node: ast.ListComp) -> None:
        """Represent the set 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 visit_ListComp(self, list_comp: ast.ListComp) -> VisitExprReturnT:
        """
        Desugars a list comprehension into nested for-loops and if-statements before flattening.

        For example, this expression::

            l = [y for x in it if cond for y in x]

        gets desugared into::

            l = []
            for x in it:
                if cond:
                    for y in x:
                        l.append(y)

        This statement block then gets flattened.
        """
        result_id = self.next_symbol_id()

        # First, desugar the comprehension into nested for-loops and if-statements.
        # Iteratively wrap `body` in for-loops and if-statements.
        body: ast.stmt = ast.Expr(clone_node(
            parse_ast_expr(f"{result_id}.append()"),
            args=[list_comp.elt]
        ))

        for comp in reversed(list_comp.generators):
            if comp.is_async:  # type: ignore # Mypy doesn't recognize the `is_async` attribute of `comprehension`.
                raise NodeNotSupportedError(comp, "Asynchronous comprehension not supported")

            for if_test in reversed(comp.ifs):
                body = ast.If(test=if_test, body=[body], orelse=[])

            body = ast.For(target=comp.target, iter=comp.iter, body=[body], orelse=[])

        # Now that we've gotten rid of the comprehension, flatten the resulting action.
        define_list = parse_ast_stmt(f"{result_id} = []")
        return load(result_id), [define_list] + self.visit_stmt(body) 

def visit_ListComp(self, node):
        # type: (ast.ListComp) -> None
        # 'listcomp' is the string literal used by python
        # to creating the SymbolTable for the corresponding
        # list comp function.
        self._handle_comprehension(node, 'listcomp') 

def test_listcomp(self):
        self._simple_comp(ast.ListComp) 

def p_listmaker_1(p):
    '''listmaker : test list_for'''
    #                 1        2
    p[0] = ast.ListComp(p[1], p[2], rule=inspect.currentframe().f_code.co_name)
    inherit_lineno(p[0], p[1]) 

def p_atom_5(p):
    '''atom : LSQB listmaker RSQB'''
    #            1         2    3
    if isinstance(p[2], ast.ListComp):
        p[0] = p[2]
        p[0].alt = p[1][1]
    else:
        p[0] = ast.List(p[2], ast.Load(), rule=inspect.currentframe().f_code.co_name, **p[1][1]) 

def visit_ListComp(self, node: ast.ListComp) -> Any:
        """Compile the list 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_ListComp(self, node: ast.ListComp) -> None:
        """Represent the list 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_listcomp(self):
        self._simple_comp(ast.ListComp) 

def test_listcomp(self):
        self._simple_comp(ast.ListComp) 

def _get_offsets(func_ast):
        import ast

        for arg in func_ast.args:
            start_line, start_col = arg.lineno - 2, arg.col_offset - 1

            # horrible hack for http://bugs.python.org/issue31241
            if isinstance(arg, (ast.ListComp, ast.GeneratorExp)):
                start_col -= 1
            yield start_line, start_col
        for kw in func_ast.keywords:
            yield kw.value.lineno - 2, kw.value.col_offset - 2 - (len(kw.arg) if kw.arg else 0) 

def visit_Assign(self, node):
        # shortcut for expansion of ROM in case statement
        if isinstance(node.value, ast.Subscript) and \
                isinstance(node.value.slice, ast.Index) and\
                isinstance(node.value.value.obj, _Rom):
            rom = node.value.value.obj.rom
#            self.write("// synthesis parallel_case full_case")
#            self.writeline()
            self.write("case (")
            self.visit(node.value.slice)
            self.write(")")
            self.indent()
            for i, n in enumerate(rom):
                self.writeline()
                if i == len(rom) - 1:
                    self.write("default: ")
                else:
                    self.write("%s: " % i)
                self.visit(node.targets[0])
                if self.isSigAss:
                    self.write(' <= ')
                    self.isSigAss = False
                else:
                    self.write(' = ')
                s = self.IntRepr(n)
                self.write("%s;" % s)
            self.dedent()
            self.writeline()
            self.write("endcase")
            return
        elif isinstance(node.value, ast.ListComp):
            # skip list comprehension assigns for now
            return
        # default behavior
        self.visit(node.targets[0])
        if self.isSigAss:
            self.write(' <= ')
            self.isSigAss = False
        else:
            self.write(' = ')
        self.visit(node.value)
        self.write(';') 

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 _find_true_position(self, node):
    """Return correct line number and column offset for a given node.
    This is necessary mainly because ListComp's location reporting reports
    the next token after the list comprehension list opening.
    Args:
      node: Node for which we wish to know the lineno and col_offset
    """
    import re
    find_open = re.compile("^\s*(\\[).*$")
    find_string_chars = re.compile("['\"]")

    if isinstance(node, ast.ListComp):
      # Strangely, ast.ListComp returns the col_offset of the first token
      # after the '[' token which appears to be a bug. Workaround by
      # explicitly finding the real start of the list comprehension.
      line = node.lineno
      col = node.col_offset
      # loop over lines
      while 1:
        # Reverse the text to and regular expression search for whitespace
        text = self._lines[line-1]
        reversed_preceding_text = text[:col][::-1]
        # First find if a [ can be found with only whitespace between it and
        # col.
        m = find_open.match(reversed_preceding_text)
        if m:
          new_col_offset = col - m.start(1) - 1
          return line, new_col_offset
        else:
          if (reversed_preceding_text=="" or
             reversed_preceding_text.isspace()):
            line = line - 1
            prev_line = self._lines[line - 1]
            # TODO(aselle):
            # this is poor comment detection, but it is good enough for
            # cases where the comment does not contain string literal starting/
            # ending characters. If ast gave us start and end locations of the
            # ast nodes rather than just start, we could use string literal
            # node ranges to filter out spurious #'s that appear in string
            # literals.
            comment_start = prev_line.find("#")
            if comment_start ==  -1:
              col = len(prev_line) -1
            elif find_string_chars.search(prev_line[comment_start:]) is None:
              col = comment_start
            else:
              return None, None
          else:
            return None, None
    # Most other nodes return proper locations (with notably does not), but
    # it is not possible to use that in an argument.
    return node.lineno, node.col_offset 

def _find_true_position(self, node):
    """Return correct line number and column offset for a given node.

    This is necessary mainly because ListComp's location reporting reports
    the next token after the list comprehension list opening.

    Args:
      node: Node for which we wish to know the lineno and col_offset
    """
    import re
    find_open = re.compile("^\s*(\\[).*$")
    find_string_chars = re.compile("['\"]")

    if isinstance(node, ast.ListComp):
      # Strangely, ast.ListComp returns the col_offset of the first token
      # after the '[' token which appears to be a bug. Workaround by
      # explicitly finding the real start of the list comprehension.
      line = node.lineno
      col = node.col_offset
      # loop over lines
      while 1:
        # Reverse the text to and regular expression search for whitespace
        text = self._lines[line-1]
        reversed_preceding_text = text[:col][::-1]
        # First find if a [ can be found with only whitespace between it and
        # col.
        m = find_open.match(reversed_preceding_text)
        if m:
          new_col_offset = col - m.start(1) - 1
          return line, new_col_offset
        else:
          if (reversed_preceding_text=="" or
             reversed_preceding_text.isspace()):
            line = line - 1
            prev_line = self._lines[line - 1]
            # TODO(aselle):
            # this is poor comment detection, but it is good enough for
            # cases where the comment does not contain string literal starting/
            # ending characters. If ast gave us start and end locations of the
            # ast nodes rather than just start, we could use string literal
            # node ranges to filter out spurious #'s that appear in string
            # literals.
            comment_start = prev_line.find("#")
            if comment_start ==  -1:
              col = len(prev_line) -1
            elif find_string_chars.search(prev_line[comment_start:]) is None:
              col = comment_start
            else:
              return None, None
          else:
            return None, None
    # Most other nodes return proper locations (with notably does not), but
    # it is not possible to use that in an argument.
    return node.lineno, node.col_offset 

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)) 

def visit_Assign(self, node):
        lhs = node.targets[0]
        rhs = node.value
        # shortcut for expansion of ROM in case statement
        if isinstance(node.value, ast.Subscript) and \
                isinstance(node.value.slice, ast.Index) and \
                isinstance(node.value.value.obj, _Rom):
            rom = node.value.value.obj.rom
            self.write("case ")
            self.visit(node.value.slice)
            self.write(" is")
            self.indent()
            size = lhs.vhd.size
            for i, n in enumerate(rom):
                self.writeline()
                if i == len(rom) - 1:
                    self.write("when others => ")
                else:
                    self.write("when %s => " % i)
                self.visit(lhs)
                if self.SigAss:
                    self.write(' <= ')
                    self.SigAss = False
                else:
                    self.write(' := ')
                if isinstance(lhs.vhd, vhd_std_logic):
                    self.write("'%s';" % n)
                elif isinstance(lhs.vhd, vhd_int):
                    self.write("%s;" % n)
                else:
                    self.write('"%s";' % tobin(n, size))
            self.dedent()
            self.writeline()
            self.write("end case;")
            return
        elif isinstance(node.value, ast.ListComp):
            # skip list comprehension assigns for now
            return
        # default behavior
        convOpen, convClose = "", ""
        if isinstance(lhs.vhd, vhd_type):
            rhs.vhd = lhs.vhd
        self.isLhs = True
        self.visit(lhs)
        self.isLhs = False
        if self.SigAss:
            if isinstance(lhs.value, ast.Name):
                sig = self.tree.symdict[lhs.value.id]
            self.write(' <= ')
            self.SigAss = False
        else:
            self.write(' := ')
        self.write(convOpen)
        # node.expr.target = obj = self.getObj(node.nodes[0])
        self.visit(rhs)
        self.write(convClose)
        self.write(';')