support dict comprehension as a pass

polystat · Aug 19, 2022 · eb2c056 · eb2c056
1 parent c987908
commit eb2c056
Showing 1 changed file with 51 additions and 23 deletions.
diff --git a/transpiler/src/main/scala/org/polystat/py2eo/transpiler/ExpressionPasses.scala b/transpiler/src/main/scala/org/polystat/py2eo/transpiler/ExpressionPasses.scala
@@ -14,33 +14,55 @@ import org.polystat.py2eo.transpiler.StatementPasses.{EAfterPass, Names, NamesU,
 import scala.collection.immutable.HashMap
 
 object ExpressionPasses {
-  def simplifyCollectionComprehension(lhs : Boolean, e : T, ns : NamesU) : (EAfterPass, NamesU) = {
+
+  def simplifyComprehensionList(inner : Statement.T, l : List[Comprehension], ann : GeneralAnnotation) : Statement.T = {
+    l.foldRight(inner : Statement.T)((comprehension, accum) => {
+      comprehension match {
+        case IfComprehension(cond) =>
+          IfSimple(cond, accum, Pass(ann.pos), ann.pos)
+        case ForComprehension(what, in, isAsync) =>
+          For(what, in, accum, None, isAsync, ann.pos)
+      }
+    })
+  }
+
+  def simplifyComprehension(lhs : Boolean, e : T, ns : NamesU) : (EAfterPass, NamesU) = {
     if (!lhs) {
       e match {
         case CollectionComprehension(kind, base, l, ann) => {
+          val inner = Assign(List(CallIndex(
+            true,
+            Field(Ident("collectionAccum", ann.pos), "append", ann.pos),
+            List((None, base)),
+            ann.pos
+          )),
+            ann.pos
+          )
           val st = Suite(
-              List(
-                Assign(List(Ident("collectionAccum", ann.pos), CollectionCons(kind, List(), ann.pos)), ann.pos),
-                l.foldRight(
-                  Assign(List(CallIndex(
-                    true,
-                    Field(Ident("collectionAccum", ann.pos), "append", ann.pos),
-                    List((None, base)),
-                    ann.pos
-                  )),
-                    ann.pos
-                  ): Statement.T
-                )((comprehension, accum) => {
-                  comprehension match {
-                    case IfComprehension(cond) =>
-                      IfSimple(cond, accum, Pass(ann.pos), ann.pos)
-                    case ForComprehension(what, in, isAsync) =>
-                      For(what, in, accum, None, isAsync, ann.pos)
-                  }
-                })
-              ),
-              ann.pos
-            )
+            List(
+              Assign(List(Ident("collectionAccum", ann.pos), CollectionCons(kind, List(), ann.pos)), ann.pos),
+              simplifyComprehensionList(inner, l, ann)
+            ),
+            ann.pos
+          )
+          (Right((st : Statement.T, Ident("collectionAccum", ann.pos))), ns)
+        }
+        case DictComprehension(Left((k, v)), l, ann) => {
+          val inner = Assign(List(CallIndex(
+            true,
+            Field(Ident("collectionAccum", ann.pos), "add", ann.pos),
+            List((None, k),(None, v)),
+            ann.pos
+          )),
+            ann.pos
+          )
+          val st = Suite(
+            List(
+              Assign(List(Ident("collectionAccum", ann.pos), DictCons(List(), ann.pos)), ann.pos),
+              simplifyComprehensionList(inner, l, ann)
+            ),
+            ann.pos
+          )
           (Right((st : Statement.T, Ident("collectionAccum", ann.pos))), ns)
         }
         case x : Any => (Left(x), ns)
@@ -341,6 +363,12 @@ object ExpressionPasses {
           val l3 = call2comprehensions(l.zip(l2))
           CollectionComprehension(kind, base, l3, ann.pos)
         }, l1, ns)
+      case DictComprehension(Left((k, v)), l, ann) =>
+        val l1 = k :: v :: comprehensions2calls(l, ann)
+        reconstruct(lhs, { case k :: v :: l2 =>
+          val l3 = call2comprehensions(l.zip(l2))
+          DictComprehension(Left((k, v)), l3, ann.pos)
+        }, l1, ns)
       case GeneratorComprehension(base, l, ann) =>
         val l1 = base :: comprehensions2calls(l, ann)
         reconstruct(lhs, { case base :: l2 =>