Merge pull request #2445 from AkshayWarrier/akshaywarrier/implement-str-funcs

Implement ``str.split()``

Author: certik

integration_tests/test_str_01.py

@@ -123,6 +123,28 @@ def test_constant_str_subscript():
     assert "abc"[2] == "c"
     assert "abc"[:2] == "ab"
 
+def test_str_split():
+    a: str = "1,2,3"
+    b: str = "1,2,,3,"
+    c: str = "1and2and3"
+    d: str = "1 2 3"
+    e: str = "   1   2   3   "
+    f: str = "123"
+    res: list[str] = a.split(",")
+    res1: list[str] = b.split(",")
+    res2: list[str] = c.split("and")
+    res3: list[str] = d.split()
+    res4: list[str] = e.split()
+    res5: list[str] = f.split(" ")
+    # res6: list[str] = "".split(" ")
+    assert res == ["1", "2", "3"]
+    assert res1 == ["1", "2", "", "3", ""]
+    assert res2 == ["1", "2", "3"]
+    assert res3 == ["1", "2", "3"]
+    assert res4 == ["1", "2", "3"]
+    assert res5 == ["123"]
+    # assert res6 == [""]
+
 def check():
     f()
     test_str_concat()
@@ -137,5 +159,6 @@ def check():
     test_str_title()
     test_str_istitle()
     test_str_isalpha()
+    test_str_split()
 
 check()

src/lpython/semantics/python_ast_to_asr.cpp

@@ -6761,6 +6761,29 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
             // Push string and substring argument on top of Vector (or Function Arguments Stack basically)
             fn_args.push_back(al, str);
             fn_args.push_back(al, value);
+        } else if(attr_name == "split") {
+            if(args.size() > 1) {
+                throw SemanticError("str.split() takes at most one argument for now.", loc);
+            }
+            fn_call_name = "_lpython_str_split";
+            ASR::call_arg_t str;
+            str.loc = loc;
+            str.m_value = s_var;
+
+            if (args.size() == 1) {
+                ASR::expr_t *arg_value = args[0].m_value;
+                ASR::ttype_t *arg_value_type = ASRUtils::expr_type(arg_value);
+                if (!ASRUtils::is_character(*arg_value_type)) {
+                    throw SemanticError("str.split() takes one argument of type: str", loc);
+                }
+                ASR::call_arg_t value;
+                value.loc = loc;
+                value.m_value = args[0].m_value;
+                fn_args.push_back(al, str);
+                fn_args.push_back(al, value);
+            } else {
+                fn_args.push_back(al, str);
+            }
         } else if(attr_name.size() > 2 && attr_name[0] == 'i' && attr_name[1] == 's') {
             /*
                 String Validation Methods i.e all "is" based functions are handled here

src/lpython/semantics/python_comptime_eval.h

@@ -88,6 +88,7 @@ struct PythonIntrinsicProcedures {
             {"_lpython_str_rstrip", {m_builtin, &not_implemented}},
             {"_lpython_str_lstrip", {m_builtin, &not_implemented}},
             {"_lpython_str_strip", {m_builtin, &not_implemented}},
+            {"_lpython_str_split", {m_builtin, &not_implemented}},
             {"_lpython_str_swapcase", {m_builtin, &not_implemented}},
             {"_lpython_str_startswith", {m_builtin, &not_implemented}},
             {"_lpython_str_endswith", {m_builtin, &not_implemented}},

src/runtime/lpython_builtin.py

@@ -821,6 +821,44 @@ def _lpython_str_strip(x: str) -> str:
     res = _lpython_str_rstrip(res)
     return res
 
+@overload
+def _lpython_str_split(x: str) -> list[str]:
+    sep: str = ' '
+    res: list[str] = []
+    start:i32 = 0
+    ind: i32
+    x_strip: str = _lpython_str_strip(x)
+    if (x_strip == ""): 
+        return res
+    while True:
+        while (start < len(x_strip) and x_strip[start] == ' '):
+            start += 1
+        ind = _lpython_str_find(x_strip[start:len(x_strip)], sep)
+        if ind == -1:
+            res.append(x_strip[start:len(x_strip)])
+            break
+        else:
+            res.append(x_strip[start:start + ind])
+            start += ind + len(sep)
+    return res
+    
+@overload
+def _lpython_str_split(x: str, sep:str) -> list[str]:
+    if len(sep) == 0:
+        raise ValueError('empty separator')
+    res: list[str] = []
+    start:i32 = 0
+    ind: i32
+    while True:
+        ind = _lpython_str_find(x[start:len(x)], sep)
+        if ind == -1:
+            res.append(x[start:len(x)])
+            break
+        else:
+            res.append(x[start:start + ind])
+            start += ind + len(sep)
+    return res
+
 @overload
 def _lpython_str_swapcase(s: str) -> str:
     res :str = ""
@@ -870,7 +908,7 @@ def _lpython_str_partition(s:str, sep: str) -> tuple[str, str, str]:
     if len(s) == 0:
         raise ValueError('empty string cannot be partitioned')
     if len(sep) == 0:
-        raise ValueError('empty seperator')
+        raise ValueError('empty separator')
     res : tuple[str, str, str]
     ind : i32
     ind = _lpython_str_find(s, sep)

tests/reference/asr-array_01_decl-39cf894.json

@@ -6,7 +6,7 @@
     "outfile": null,
     "outfile_hash": null,
     "stdout": "asr-array_01_decl-39cf894.stdout",
-    "stdout_hash": "489d2e6a364cc6020f2942b94738849349928901f1269b975a6e2464",
+    "stdout_hash": "5d4751789e2ddcd882c4d6026f801ba32cfc227fafff7395a788bdd9",
     "stderr": null,
     "stderr_hash": null,
     "returncode": 0