codelab example standard library functions, factorial word_count & ch…

…oose_random
chanzuckerberg · Apr 30, 2020 · 167010a · 167010a
1 parent 50b4132
commit 167010a
Showing 1 changed file with 39 additions and 0 deletions.
diff --git a/WDL/StdLib.py b/WDL/StdLib.py
@@ -77,6 +77,18 @@ def sub(input: Value.String, pattern: Value.String, replace: Value.String) -> Va
         def defined(v: Value.Base):
             return Value.Boolean(not isinstance(v, Value.Null))
 
+        @static([Type.Int()], Type.Int())
+        def factorial(v: Value.Int) -> Value.Int:
+            def f(n: int) -> int:
+                return 1 if n <= 1 else n * f(n - 1)
+
+            return Value.Int(f(v.value))
+
+        @static([Type.File()], Type.Int())
+        def word_count(v: Value.File) -> Value.Int:
+            with open(self._devirtualize_filename(v.value), "r") as infile:
+                return Value.Int(len(infile.read().split(" ")))
+
         # write_*
         static([Type.Array(Type.String())], Type.File(), "write_lines")(
             self._write(_serialize_lines)
@@ -120,6 +132,7 @@ def defined(v: Value.Base):
         self.cross = _Cross()
         self.flatten = _Flatten()
         self.transpose = _Transpose()
+        self.choose_random = _ChooseRandom()
 
     def _read(self, parse: Callable[[str], Value.Base]) -> Callable[[Value.File], Value.Base]:
         "generate read_* function implementation based on parse"
@@ -883,3 +896,29 @@ def _call_eager(self, expr: "Expr.Apply", arguments: List[Value.Base]) -> Value.
             Type.String(),
             [Value.String(pfx + s.coerce(Type.String()).value) for s in arguments[1].value],
         )
+
+
+class _ChooseRandom(EagerFunction):
+    def infer_type(self, expr: "Expr.Apply") -> Type.Base:
+        if len(expr.arguments) not in [1, 2]:
+            raise Error.WrongArity(expr, 1)
+        arg0ty = expr.arguments[0].type
+        if not isinstance(arg0ty, Type.Array):
+            raise Error.StaticTypeMismatch(expr.arguments[0], Type.Array(Type.Any()), arg0ty)
+        if len(expr.arguments) == 1:
+            return arg0ty.item_type
+        arg1ty = expr.arguments[1].type
+        if not isinstance(arg1ty, Type.Array):
+            raise Error.StaticTypeMismatch(expr.arguments[1], Type.Array(Type.Any()), arg1ty)
+        return Type.Pair(arg0ty.item_type, arg1ty.item_type)
+
+    def _call_eager(self, expr: "Expr.Apply", arguments: List[Value.Base]) -> Value.Base:
+        import random
+
+        if not arguments[0].value or (len(arguments) > 1 and not arguments[1].value):
+            raise Error.RuntimeError("empty array passed to choose_random()")
+        item0 = random.choice(arguments[0].value)
+        if len(arguments) == 1:
+            return item0
+        item1 = random.choice(arguments[1].value)
+        return Value.Pair(item0.type, item1.type, (item0, item1))