add for in string support (#21990)

wanchaol · facebook-github-bot · commit c9344fc9c43e · 2019-06-23T20:49:30.000-07:00
Summary: Pull Request resolved: #21990 ghimport-source-id: 69b4882 Test Plan: Imported from OSS Differential Revision: D15948547 Pulled By: wanchaol fbshipit-source-id: 057e7f4fb67c6dca98458ceb14414368e1a86260
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -9363,6 +9363,27 @@ def list_iterables(x):
                 return x
             ''')
 
+    def test_for_in_string(self):
+        def test_strings(x):
+            # type: (str) -> str
+            reverse = ""
+            for c in x:
+                reverse = c + reverse
+            return reverse
+
+        self.checkScript(test_strings, ("hello",))
+        self.checkScript(test_strings, ("",))
+
+        def test_list_strings(x):
+            # type: (List[str]) -> str
+            result = ""
+            for sub_str in x:
+                result += sub_str
+            return result
+
+        self.checkScript(test_list_strings, (["hello", "world"],))
+        self.checkScript(test_list_strings, (["hello", " ", "world", ""],))
+
     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
@@ -1344,7 +1344,7 @@ struct to_ir {
     auto 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.
+    // RangeValue for range(), and SimpleValue for types like List/Tensor/Dict/String.
     auto range_val = std::dynamic_pointer_cast<RangeValue>(sv);
     auto siv = std::dynamic_pointer_cast<SimpleValue>(sv);
     auto iterable_tree = std::dynamic_pointer_cast<IterableTree>(sv);
diff --git a/torch/csrc/jit/script/sugared_value.cpp b/torch/csrc/jit/script/sugared_value.cpp
@@ -265,6 +265,7 @@ Value* SimpleValue::len(const SourceRange& loc, Function& m) {
   TypePtr val_type = val->type();
   Graph& g = *m.graph();
   if (val_type->cast<ListType>() ||
+      val_type->cast<StringType>() ||
       val_type->isSubtypeOf(TensorType::get())) {
     return g.insert(aten::len, {val}, {}, loc);
   } else {
@@ -280,6 +281,8 @@ Value* SimpleValue::getelem(const SourceRange&loc, Function& m, Value* i) {
   Value* cur_elem = nullptr;
   if (val_type->cast<ListType>()) {
     cur_elem = g.insert(aten::select, {val, i}, {}, loc);
+  } else if (val_type->cast<StringType>()) {
+    cur_elem = g.insert(prim::StringIndex, {val, i}, {}, loc);
   } else if (val_type->isSubtypeOf(TensorType::get())) {
     cur_elem = g.insert(aten::select, {val, 0, i}, {}, loc);
   } else {
