add for in dict support (#22006)

wanchaol · facebook-github-bot · commit d96ce9b9fee9 · 2019-06-23T20:49:35.000-07:00
Summary: Pull Request resolved: #22006 ghimport-source-id: d9686c0 Test Plan: Imported from OSS Differential Revision: D15948548 Pulled By: wanchaol fbshipit-source-id: 4227502ca050099085ad481aef725ac2cab06d74
diff --git a/aten/src/ATen/core/interned_strings.h b/aten/src/ATen/core/interned_strings.h
@@ -142,6 +142,7 @@ namespace c10 {
   _(aten, list)                    \
   _(aten, wait)                    \
   _(aten, save)                    \
+  _(aten, keys)                     \
   _(aten, ord)                     \
   _(aten, chr)                     \
   _(aten, hex)                     \
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -9384,6 +9384,28 @@ def test_list_strings(x):
         self.checkScript(test_list_strings, (["hello", "world"],))
         self.checkScript(test_list_strings, (["hello", " ", "world", ""],))
 
+    def test_for_in_dict(self):
+        def test_dicts(x):
+            # type: (Dict[str, int]) -> int
+            sum = 0
+            for key in x:
+                sum += x[key]
+            return sum
+
+        self.checkScript(test_dicts, ({"a": 1, "b": 2, "c": 3},))
+
+        def test_dict_keys_values(x):
+            # type: (Dict[str, int]) -> Tuple[str, int]
+            key_str = ""
+            sum = 0
+            for key in x.keys():
+                key_str += key
+            for val in x.values():
+                sum += val
+            return key_str, sum
+
+        self.checkScript(test_dicts, ({"a": 1, "b": 2, "c": 3},))
+
     def test_for_tuple_unpack(self):
         def for_tuple_unpack(x, y):
             for i, j in [[3, 4], [5, 6], [7, 8]]:
diff --git a/torch/csrc/jit/script/compiler.cpp b/torch/csrc/jit/script/compiler.cpp
@@ -1341,7 +1341,7 @@ struct to_ir {
     }
     // Emit loop information for builtinFunction values like range(), zip(), 
     // enumerate() or SimpleValue like List, Tensor, Dict, etc.
-    auto sv = emitSugaredExpr(itrs[0], 1);
+    SugaredValuePtr sv = emitSugaredExpr(itrs[0], 1);
 
     // We will get IterableTree for builtinFunctions zip() and enumerate(),
     // RangeValue for range(), and SimpleValue for types like List/Tensor/Dict/String.
@@ -1352,6 +1352,11 @@ struct to_ir {
     // For SimpleValue(except Tuple) or RanveValue/IterableTree, emit common loop
     if ((siv && !siv->getValue()->type()->cast<TupleType>())
         || range_val || iterable_tree) {
+      // looping over a dict defaults to looping over the keys in python
+      if (siv && siv->getValue()->type()->cast<DictType>()) {
+        sv = std::make_shared<SimpleValue>(
+          graph->insert(aten::keys, {siv->getValue()}, {}, stmt.range()));
+      }
       emitLoopCommon(stmt.range(), body, sv, targets, {}); 
       return;
     }
