From fef4a81296bbd72e3a69c2fdc549acaaa835e0e1 Mon Sep 17 00:00:00 2001
From: Mike Lin <dna@mlin.net>
Date: Sat, 21 Dec 2019 13:13:19 -1000
Subject: [PATCH] codelab example standard library functions, factorial
 word_count & choose_random

---
 WDL/StdLib.py | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/WDL/StdLib.py b/WDL/StdLib.py
index 4ece08d4..0d0d7cac 100644
--- a/WDL/StdLib.py
+++ b/WDL/StdLib.py
@@ -85,6 +85,18 @@ def defined(v: Value.Base):
         def sep(sep: Value.String, iterable: Value.Array) -> Value.String:
             return Value.String(sep.value.join(v.value for v in iterable.value))
 
+        @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)
@@ -136,6 +148,7 @@ def sep(sep: Value.String, iterable: Value.Array) -> Value.String:
         self.cross = _Cross()
         self.flatten = _Flatten()
         self.transpose = _Transpose()
+        self.choose_random = _ChooseRandom()
 
         if self.wdl_version not in ["draft-2", "1.0"]:
             self.min = _ArithmeticOperator("min", lambda l, r: min(l, r))
@@ -1161,3 +1174,29 @@ def _call_eager(self, expr: "Expr.Apply", arguments: List[Value.Base]) -> Value.
                 raise Error.EvalError(expr, "duplicate keys supplied to as_map(): " + str(k))
             singletons.append((k, vs.value[0]))
         return Value.Map((collectedty.item_type[0], arrayty.item_type), singletons, expr)
+
+
+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))