Python: Fix missing DictionaryElementContent

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll

@@ -809,6 +809,8 @@ predicate dictStoreStep(CfgNode nodeFrom, DictionaryElementContent c, Node nodeT
  * TODO: Once TaintTracking no longer uses `dictStoreStep`, unify the two predicates.
  */
 private predicate moreDictStoreSteps(CfgNode nodeFrom, DictionaryElementContent c, Node nodeTo) {
+  // NOTE: It's important to add logic to the newtype definition of
+  // DictionaryElementContent if you add new cases here.
   exists(SubscriptNode subscript |
     nodeTo.(PostUpdateNode).getPreUpdateNode().asCfgNode() = subscript.getObject() and
     nodeFrom.asCfgNode() = subscript.(DefinitionNode).getValue() and

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPublic.qll

@@ -605,9 +605,19 @@ newtype TContent =
   } or
   /** An element of a dictionary under a specific key. */
   TDictionaryElementContent(string key) {
-    key = any(KeyValuePair kvp).getKey().(StrConst).getS()
+    // {"key": ...}
+    key = any(KeyValuePair kvp).getKey().(StrConst).getText()
     or
+    // func(key=...)
     key = any(Keyword kw).getArg()
+    or
+    // d["key"] = ...
+    key = any(SubscriptNode sub | sub.isStore() | sub.getIndex().getNode().(StrConst).getText())
+    or
+    // d.setdefault("key", ...)
+    exists(CallNode call | call.getFunction().(AttrNode).getName() = "setdefault" |
+      key = call.getArg(0).getNode().(StrConst).getText()
+    )
   } or
   /** An element of a dictionary under any key. */
   TDictionaryElementAnyContent() or

python/ql/test/experimental/dataflow/fieldflow/test_dict.py

@@ -52,7 +52,7 @@ def test_dict_update_fresh_key():
     # for keys used in "inline update" like this
     d = {}
     d["fresh_key"] = SOURCE
-    SINK(d["fresh_key"])  # $ MISSING: flow="SOURCE, l:-1 -> d['fresh_key']"
+    SINK(d["fresh_key"])  # $ flow="SOURCE, l:-1 -> d['fresh_key']"
 
 
 @expects(3) # $ unresolved_call=expects(..) unresolved_call=expects(..)(..)