implement take-while and drop-while

1 files changed, 40 insertions, 0 deletions

diff --git a/neb/std/core.py b/neb/std/core.py
index 50a0c3e..9a9e56d 100644
--- a/neb/std/core.py
+++ b/neb/std/core.py
@@ -98,6 +98,46 @@ def interpretForEach(symbol, args, env, ns):
 for_each_arg = Arg("list", TypeEnum.LIST)
 CORE.register("for-each", NebSyntax("for-each", interpretForEach, [for_each_arg, for_body_arg], for_body_arg))
 
+def interpretTakeWhile(symbol, args, env, ns):
+    coll = evaluate(args[0], env, ns)
+    if not isinstance(coll, List):
+        raise InterpretPanic(symbol, "coll must be a :list", coll)
+    new_env = Environment(env)
+    ret = []
+    ev = None
+    for item in coll.args:
+        new_env.register("_item_", evaluate(item, env, ns))
+        for arg in args[1:]:
+            ev = evaluate(arg, new_env, ns)
+        if not isinstance(ev, Bool):
+            raise InterpretPanic(symbol, "condition must evaluate to a :bool", ev)
+        if ev.value is False:
+            break
+        ret.append(item)
+    return List(ret)
+
+CORE.register("take-while", NebSyntax("take-while", interpretTakeWhile, [for_each_arg, for_body_arg], for_body_arg))
+
+def interpretDropWhile(symbol, args, env, ns):
+    coll = evaluate(args[0], env, ns)
+    if not isinstance(coll, List):
+        raise InterpretPanic(symbol, "coll must be a :list", coll)
+    new_env = Environment(env)
+    ev = None
+    which_idx = None
+    for idx, item in enumerate(coll.args):
+        which_idx = idx
+        new_env.register("_item_", evaluate(item, env, ns))
+        for arg in args[1:]:
+            ev = evaluate(arg, new_env, ns)
+        if not isinstance(ev, Bool):
+            raise InterpretPanic(symbol, "condition must evaluate to a :bool", ev)
+        if ev.value is False:
+            break
+    return List(coll.args[which_idx:])
+
+CORE.register("drop-while", NebSyntax("drop-while", interpretDropWhile, [for_each_arg, for_body_arg], for_body_arg))
+
 def interpretBranch(symbol, args, env, ns):
     for arg in args:
         if len(arg.args) != 2: