Merge 5936735 into 458beb8

chanzuckerberg · Sep 10, 2020 · fbb9acb · fbb9acb
2 parents 458beb8 + 5936735
commit fbb9acb
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 5 deletions.
diff --git a/WDL/StdLib.py b/WDL/StdLib.py
@@ -127,6 +127,7 @@ def sep(sep: Value.String, iterable: Value.Array) -> Value.String:
         self.transpose = _Transpose()
         self.quote = _Quote()
         self.squote = _Quote(squote=True)
+        self.keys = _Keys()
 
     def _read(self, parse: Callable[[str], Value.Base]) -> Callable[[Value.File], Value.Base]:
         "generate read_* function implementation based on parse"
@@ -954,3 +955,23 @@ def _call_eager(self, expr: "Expr.Apply", arguments: List[Value.Base]) -> Value.
                 for s in arguments[0].value
             ],
         )
+
+
+class _Keys(EagerFunction):
+    def infer_type(self, expr: "Expr.Apply") -> Type.Base:
+        if len(expr.arguments) != 1:
+            raise Error.WrongArity(expr, 1)
+        arg0ty = expr.arguments[0].type
+        if not isinstance(arg0ty, Type.Map):
+            raise Error.StaticTypeMismatch(
+                expr.arguments[0], Type.Map((Type.Any(), Type.Any())), arg0ty
+            )
+        return Type.Array(arg0ty.item_type[0].copy())
+
+    def _call_eager(self, expr: "Expr.Apply", arguments: List[Value.Base]) -> Value.Base:
+        if isinstance(arguments[0], Value.Null):
+            return Value.Array(Type.Any(), [])
+        assert isinstance(arguments[0], Value.Map)
+        mapty = arguments[0].type
+        assert isinstance(mapty, Type.Map)
+        return Value.Array(mapty.item_type[0], [p[0] for p in arguments[0].value])
diff --git a/WDL/Tree.py b/WDL/Tree.py
@@ -400,8 +400,8 @@ def typecheck(
                     type_env = type_env2
             errors.maybe_raise()
             # Typecheck the output expressions
+            stdlib = StdLib.TaskOutputs()
             for decl in self.outputs:
-                stdlib = StdLib.TaskOutputs()
                 errors.try1(
                     lambda: decl.typecheck(type_env, stdlib=stdlib, check_quant=check_quant)
                 )

diff --git a/WDL/Value.py b/WDL/Value.py
@@ -266,8 +266,10 @@ def json(self) -> Any:
         ""
         ans = {}
         for k, v in self.value:
-            assert isinstance(k, String)  # TODO
-            ans[k.value] = v.json
+            assert k.type.coerces(Type.String())
+            kstr = k.coerce(Type.String()).value
+            if kstr not in ans:
+                ans[kstr] = v.json
         return ans
 
     @property

diff --git a/WDL/runtime/workflow.py b/WDL/runtime/workflow.py
@@ -753,7 +753,8 @@ def _workflow_main_loop(
                 if terminating():
                     raise Terminated()
                 # schedule all runnable calls
-                next_call = state.step(cfg, _StdLib(cfg, state, cache))
+                stdlib = _StdLib(cfg, state, cache)
+                next_call = state.step(cfg, stdlib)
                 while next_call:
                     call_dir = os.path.join(run_dir, next_call.id)
                     if os.path.exists(call_dir):
@@ -779,7 +780,7 @@ def _workflow_main_loop(
                     else:
                         assert False
                     call_futures[future] = next_call.id
-                    next_call = state.step(cfg, _StdLib(cfg, state, cache))
+                    next_call = state.step(cfg, stdlib)
                 # no more calls to launch right now; wait for an outstanding call to finish
                 future = next(futures.as_completed(call_futures), None)
                 if future:

diff --git a/tests/test_5stdlib.py b/tests/test_5stdlib.py
@@ -921,3 +921,38 @@ def test_squote(self):
             }
         }
         """,expected_exception=WDL.Error.StaticTypeMismatch)
+
+    def test_keys(self):
+        outputs = self._test_task(R"""
+        version development
+        task test_keys {
+            input {
+                Map[String,String] m1 = {"a": "b", "c": "d"}
+                Map[Int,Boolean] m2 = {1: true, -1: false}
+                Map[Int,Float]? m3
+            }
+            command {}
+            output {
+                Array[String] k1 = keys(m1)
+                Array[Int] k2 = keys(m2)
+                Array[Int] k3 = keys(m3)
+                Array[Boolean] k4 = keys({})
+            }
+        }
+        """)
+        self.assertEqual(outputs["k1"], ["a", "c"])
+        self.assertEqual(outputs["k2"], [1,-1])
+        self.assertEqual(outputs["k3"], [])
+        self.assertEqual(outputs["k4"], [])
+
+        with self.assertRaises(WDL.Error.StaticTypeMismatch):
+            self._test_task(R"""
+            version development
+            task test_keys {
+                input {
+                    Array[Int] a = keys([1,2,3])
+                }
+                command {}
+                output {}
+            }
+            """)