Fix handling for some wrappers + add test case

This commit is contained in:
Joe Farebrother
2025-05-30 11:24:06 +01:00
parent f3a5608b06
commit 73f2770acb
3 changed files with 14 additions and 6 deletions

View File

@@ -50,7 +50,7 @@ predicate iterWrapperMethods(Function iter, Function next) {
exists(string field |
exists(Return r, DataFlow::Node self, DataFlow::AttrRead read |
r.getScope() = iter and
r.getValue() = iterCall(read).asExpr() and
r.getValue() = [iterCall(read).asExpr(), read.asExpr()] and
read.accesses(self, field) and
isSelfVar(iter, self.asExpr())
) and
@@ -63,15 +63,13 @@ predicate iterWrapperMethods(Function iter, Function next) {
)
}
/** Gets a call to `iter(arg)`, `arg.__iter__()`, or `arg` itself (which we assume may already be an iterator). */
/** Gets a call to `iter(arg)` or `arg.__iter__()`. */
private DataFlow::CallCfgNode iterCall(DataFlow::Node arg) {
result.(DataFlow::MethodCallNode).calls(arg, "__iter__")
or
result = API::builtin("iter").getACall() and
arg = result.getArg(0) and
not exists(result.getArg(1))
or
result = arg // assume the wrapping field is already an iterator
}
/** Gets a call to `next(arg)` or `arg.__next__()`. */