add substring (with an alias to subs)

heyLu · heyLu · commit 7e29c9819230 · 2014-11-08T16:43:17.000+01:00
diff --git a/pixie/stdlib.lisp b/pixie/stdlib.lisp
@@ -485,6 +485,8 @@
          (assoc m k (assoc-in (get m k) ks v))
          (assoc m k v)))))
 
+(def subs pixie.string/substring)
+
 (defmacro assert
   ([test]
      `(if ~test
diff --git a/pixie/vm/libs/string.py b/pixie/vm/libs/string.py
@@ -54,6 +54,23 @@ def index_of4(a, sep, start, end):
 index_of.set_root(MultiArityFn({2: wrap_fn(index_of2), 3: wrap_fn(index_of3), 4: wrap_fn(index_of4)},
                                required_arity = 2))
 
+def substring2(a, start):
+    return substring3(a, start, rt._count(a))
+
+def substring3(a, start, end):
+    affirm(isinstance(a, String), u"First argument must be a string")
+    affirm(isinstance(start, Integer) and isinstance(end, Integer), u"Second and third argument must be integers")
+    start = start.int_val()
+    end = end.int_val()
+    if start >= 0 and end >= 0:
+        return rt.wrap(rt.name(a)[start:end])
+    else:
+        runtime_error(u"Second and third argument must be non-negative integers")
+
+substring = intern_var(u"pixie.string", u"substring")
+substring.set_root(MultiArityFn({2: wrap_fn(substring2), 3: wrap_fn(substring3)},
+                                required_arity = 2))
+
 @as_var("pixie.string", "upper-case")
 def upper_case(a):
     a = rt.name(a)
diff --git a/tests/test-strings.lisp b/tests/test-strings.lisp
@@ -39,6 +39,16 @@
 
     (t/assert= (s/index-of s "h" 1 2) -1)))
 
+(t/deftest test-substring
+  (let [s "heyhohuh"]
+    (t/assert= (s/substring s 0) s)
+    (t/assert= (s/substring s 3) (s/substring s 3 (count s)))
+    (t/assert= (s/substring s 0 0) "")
+    (t/assert= (s/substring s 0 3) "hey")
+    (t/assert= (s/substring s 3 5) "ho")
+    (t/assert= (s/substring s 5 8) "huh")
+    (t/assert= (s/substring s 3 10000) "hohuh")))
+
 (t/deftest test-upper-case
   (t/assert= (s/lower-case "") "")
   (t/assert= (s/upper-case "hey") "HEY")
