overlay[local]

module;

import python

/** The base class for list, set and dictionary comprehensions, and generator expressions. */

abstract class Comp extends Expr {

abstract Function getFunction();

/** Gets the iterable of this set comprehension. */

abstract Expr getIterable();

/** Gets the iteration variable for the nth innermost generator of this comprehension. */

Variable getIterationVariable(int n) {

result.getAnAccess() = this.getNthInnerLoop(n).getTarget()

}

/** Gets the nth innermost For expression of this comprehension. */

For getNthInnerLoop(int n) {

n = 0 and result = this.getFunction().getStmt(0)

or

result = this.getNthInnerLoop(n - 1).getStmt(0)

}

/** Gets the iteration variable for a generator of this list comprehension. */

Variable getAnIterationVariable() { result = this.getIterationVariable(_) }

/** Gets the scope in which the body of this list comprehension evaluates. */

Scope getEvaluatingScope() { result = this.getFunction() }

/** Gets the expression for elements of this comprehension. */

Expr getElt() {

exists(Yield yield, Stmt body |

result = yield.getValue() and

body = this.getNthInnerLoop(_).getAStmt()

|

yield = body.(ExprStmt).getValue()

or

yield = body.(If).getStmt(0).(ExprStmt).getValue()

)

}

}

/** A list comprehension, such as `[ chr(x) for x in range(ord('A'), ord('Z')+1) ]` */

class ListComp extends ListComp_, Comp {

override Expr getASubExpression() {

result = this.getAGenerator().getASubExpression() or

result = this.getElt() or

result = this.getIterable()

}

override AstNode getAChildNode() {

result = this.getAGenerator() or

result = this.getIterable() or

result = this.getFunction()

}

override predicate hasSideEffects() { any() }

/** Gets the scope in which the body of this list comprehension evaluates. */

override Scope getEvaluatingScope() {

major_version() = 2 and result = this.getScope()

or

major_version() = 3 and result = this.getFunction()

}

/** Gets the iteration variable for the nth innermost generator of this list comprehension */

override Variable getIterationVariable(int n) { result = Comp.super.getIterationVariable(n) }

override Function getFunction() { result = ListComp_.super.getFunction() }

override Expr getIterable() { result = ListComp_.super.getIterable() }

override Expr getElt() { result = Comp.super.getElt() }

}

/** A set comprehension such as `{ v for v in "0123456789" }` */

class SetComp extends SetComp_, Comp {

override Expr getASubExpression() { result = this.getIterable() }

override AstNode getAChildNode() {

result = this.getASubExpression() or

result = this.getFunction()

}

override predicate hasSideEffects() { any() }

override Function getFunction() { result = SetComp_.super.getFunction() }

override Expr getIterable() { result = SetComp_.super.getIterable() }

}

/** A dictionary comprehension, such as `{ k:v for k, v in enumerate("0123456789") }` */

class DictComp extends DictComp_, Comp {

override Expr getASubExpression() { result = this.getIterable() }

override AstNode getAChildNode() {

result = this.getASubExpression() or

result = this.getFunction()

}

override predicate hasSideEffects() { any() }

override Function getFunction() { result = DictComp_.super.getFunction() }

override Expr getIterable() { result = DictComp_.super.getIterable() }

}

/** A generator expression, such as `(var for var in iterable)` */

class GeneratorExp extends GeneratorExp_, Comp {

override Expr getASubExpression() { result = this.getIterable() }

override AstNode getAChildNode() {

result = this.getASubExpression() or

result = this.getFunction()

}

override predicate hasSideEffects() { any() }

override Function getFunction() { result = GeneratorExp_.super.getFunction() }

override Expr getIterable() { result = GeneratorExp_.super.getIterable() }

}