@@ -249,6 +249,60 @@ module API {
249249 */
250250 Node getASubscript ( ) { result = this .getASuccessor ( Label:: subscript ( ) ) }
251251
252+ /**
253+ * Gets a node representing an index of a subscript of this node.
254+ * For example, in `obj[x]`, `x` is an index of `obj`.
255+ */
256+ Node getIndex ( ) { result = this .getASuccessor ( Label:: index ( ) ) }
257+
258+ /**
259+ * Gets a node representing a subscript of this node at (string) index `key`.
260+ * This requires that the index can be statically determined.
261+ *
262+ * For example, the subscripts of `a` and `b` below would be found using
263+ * the index `foo`:
264+ * ```py
265+ * a["foo"]
266+ * x = "foo" if cond else "bar"
267+ * b[x]
268+ * ```
269+ */
270+ Node getSubscript ( string key ) {
271+ exists ( API:: Node index | result = this .getSubscriptAt ( index ) |
272+ key = index .getAValueReachingSink ( ) .asExpr ( ) .( PY:: StrConst ) .getText ( )
273+ )
274+ }
275+
276+ /**
277+ * Gets a node representing a subscript of this node at index `index`.
278+ */
279+ Node getSubscriptAt ( API:: Node index ) {
280+ result = this .getASubscript ( ) and
281+ index = this .getIndex ( ) and
282+ (
283+ // subscripting
284+ exists ( PY:: SubscriptNode subscript |
285+ subscript .getObject ( ) = this .getAValueReachableFromSource ( ) .asCfgNode ( ) and
286+ subscript .getIndex ( ) = index .asSink ( ) .asCfgNode ( )
287+ |
288+ // reading
289+ subscript = result .asSource ( ) .asCfgNode ( )
290+ or
291+ // writing
292+ subscript .( PY:: DefinitionNode ) .getValue ( ) = result .asSink ( ) .asCfgNode ( )
293+ )
294+ or
295+ // dictionary literals
296+ exists ( PY:: Dict dict , PY:: KeyValuePair item |
297+ dict = this .getAValueReachingSink ( ) .asExpr ( ) and
298+ dict .getItem ( _) = item and
299+ item .getKey ( ) = index .asSink ( ) .asExpr ( )
300+ |
301+ item .getValue ( ) = result .asSink ( ) .asExpr ( )
302+ )
303+ )
304+ }
305+
252306 /**
253307 * Gets a string representation of the lexicographically least among all shortest access paths
254308 * from the root to this node.
@@ -405,7 +459,7 @@ module API {
405459 Node builtin ( string n ) { result = moduleImport ( "builtins" ) .getMember ( n ) }
406460
407461 /**
408- * An `CallCfgNode` that is connected to the API graph.
462+ * A `CallCfgNode` that is connected to the API graph.
409463 *
410464 * Can be used to reason about calls to an external API in which the correlation between
411465 * parameters and/or return values must be retained.
@@ -694,12 +748,31 @@ module API {
694748 rhs = aw .getValue ( )
695749 )
696750 or
697- // TODO: I had expected `DataFlow::AttrWrite` to contain the attribute writes from a dict, that's how JS works.
751+ // dictionary literals
698752 exists ( PY:: Dict dict , PY:: KeyValuePair item |
699753 dict = pred .( DataFlow:: ExprNode ) .getNode ( ) .getNode ( ) and
700- dict .getItem ( _) = item and
701- lbl = Label:: member ( item .getKey ( ) .( PY:: StrConst ) .getS ( ) ) and
702- rhs .( DataFlow:: ExprNode ) .getNode ( ) .getNode ( ) = item .getValue ( )
754+ dict .getItem ( _) = item
755+ |
756+ // from `x` to `{ "key": x }`
757+ // TODO: once convenient, this should be done at a higher level than the AST,
758+ // at least at the CFG layer, to take splitting into account.
759+ rhs .( DataFlow:: ExprNode ) .getNode ( ) .getNode ( ) = item .getValue ( ) and
760+ lbl = Label:: subscript ( )
761+ or
762+ // from `"key"` to `{ "key": x }`
763+ // TODO: once convenient, this should be done at a higher level than the AST,
764+ // at least at the CFG layer, to take splitting into account.
765+ rhs .( DataFlow:: ExprNode ) .getNode ( ) .getNode ( ) = item .getKey ( ) and
766+ lbl = Label:: index ( )
767+ )
768+ or
769+ // list literals, from `x` to `[x]`
770+ // TODO: once convenient, this should be done at a higher level than the AST,
771+ // at least at the CFG layer, to take splitting into account.
772+ // Also consider `SequenceNode for generality.
773+ exists ( PY:: List list | list = pred .( DataFlow:: ExprNode ) .getNode ( ) .getNode ( ) |
774+ rhs .( DataFlow:: ExprNode ) .getNode ( ) .getNode ( ) = list .getAnElt ( ) and
775+ lbl = Label:: subscript ( )
703776 )
704777 or
705778 exists ( PY:: CallableExpr fn | fn = pred .( DataFlow:: ExprNode ) .getNode ( ) .getNode ( ) |
@@ -720,6 +793,20 @@ module API {
720793 lbl = Label:: memberFromRef ( aw )
721794 )
722795 or
796+ // subscripting
797+ exists ( DataFlow:: LocalSourceNode src , DataFlow:: Node subscript , DataFlow:: Node index |
798+ use ( base , src ) and
799+ subscript = trackUseNode ( src ) .getSubscript ( index )
800+ |
801+ // from `x` to a definition of `x[...]`
802+ rhs .asCfgNode ( ) = subscript .asCfgNode ( ) .( PY:: DefinitionNode ) .getValue ( ) and
803+ lbl = Label:: subscript ( )
804+ or
805+ // from `x` to `"key"` in `x["key"]`
806+ rhs = index and
807+ lbl = Label:: index ( )
808+ )
809+ or
723810 exists ( EntryPoint entry |
724811 base = root ( ) and
725812 lbl = Label:: entryPoint ( entry ) and
@@ -757,7 +844,8 @@ module API {
757844 or
758845 // Subscripting a node that is a use of `base`
759846 lbl = Label:: subscript ( ) and
760- ref = pred .getASubscript ( )
847+ ref = pred .getSubscript ( _) and
848+ ref .asCfgNode ( ) .isLoad ( )
761849 or
762850 // Subclassing a node
763851 lbl = Label:: subclass ( ) and
@@ -973,8 +1061,7 @@ module API {
9731061 member = any ( DataFlow:: AttrRef pr ) .getAttributeName ( ) or
9741062 exists ( Builtins:: likelyBuiltin ( member ) ) or
9751063 ImportStar:: namePossiblyDefinedInImportStar ( _, member , _) or
976- Impl:: prefix_member ( _, member , _) or
977- member = any ( PY:: Dict d ) .getAnItem ( ) .( PY:: KeyValuePair ) .getKey ( ) .( PY:: StrConst ) .getS ( )
1064+ Impl:: prefix_member ( _, member , _)
9781065 } or
9791066 MkLabelUnknownMember ( ) or
9801067 MkLabelParameter ( int i ) {
@@ -992,6 +1079,7 @@ module API {
9921079 MkLabelSubclass ( ) or
9931080 MkLabelAwait ( ) or
9941081 MkLabelSubscript ( ) or
1082+ MkLabelIndex ( ) or
9951083 MkLabelEntryPoint ( EntryPoint ep )
9961084
9971085 /** A label for a module. */
@@ -1072,6 +1160,11 @@ module API {
10721160 override string toString ( ) { result = "getASubscript()" }
10731161 }
10741162
1163+ /** A label that gets the index of a subscript. */
1164+ class LabelIndex extends ApiLabel , MkLabelIndex {
1165+ override string toString ( ) { result = "getIndex()" }
1166+ }
1167+
10751168 /** A label for entry points. */
10761169 class LabelEntryPoint extends ApiLabel , MkLabelEntryPoint {
10771170 private EntryPoint entry ;
@@ -1120,6 +1213,9 @@ module API {
11201213 /** Gets the `subscript` edge label. */
11211214 LabelSubscript subscript ( ) { any ( ) }
11221215
1216+ /** Gets the `subscript` edge label. */
1217+ LabelIndex index ( ) { any ( ) }
1218+
11231219 /** Gets the label going from the root node to the nodes associated with the given entry point. */
11241220 LabelEntryPoint entryPoint ( EntryPoint ep ) { result = MkLabelEntryPoint ( ep ) }
11251221 }
0 commit comments