String#partition (#3435)

crystal-lang · Oct 20, 2016 · 9cdc89e · 9cdc89e
1 parent 2754433
commit 9cdc89e
Showing 2 changed files with 69 additions and 0 deletions.
diff --git a/spec/std/string_spec.cr b/spec/std/string_spec.cr
@@ -627,6 +627,36 @@ describe "String" do
     end
   end
 
+  describe "partition" do
+    describe "by char" do
+      "hello".partition('h').should eq ({"", "h", "ello"})
+      "hello".partition('o').should eq ({"hell", "o", ""})
+      "hello".partition('l').should eq ({"he", "l", "lo"})
+      "hello".partition('x').should eq ({"hello", "", ""})
+    end
+
+    describe "by string" do
+      "hello".partition("h").should eq ({"", "h", "ello"})
+      "hello".partition("o").should eq ({"hell", "o", ""})
+      "hello".partition("l").should eq ({"he", "l", "lo"})
+      "hello".partition("ll").should eq ({"he", "ll", "o"})
+      "hello".partition("x").should eq ({"hello", "", ""})
+    end
+
+    describe "by regex" do
+      "hello".partition(/h/).should eq ({"", "h", "ello"})
+      "hello".partition(/o/).should eq ({"hell", "o", ""})
+      "hello".partition(/l/).should eq ({"he", "l", "lo"})
+      "hello".partition(/ll/).should eq ({"he", "ll", "o"})
+      "hello".partition(/.l/).should eq ({"h", "el", "lo"})
+      "hello".partition(/.h/).should eq ({"hello", "", ""})
+      "hello".partition(/h./).should eq ({"", "he", "llo"})
+      "hello".partition(/o./).should eq ({"hello", "", ""})
+      "hello".partition(/.o/).should eq ({"hel", "lo", ""})
+      "hello".partition(/x/).should eq ({"hello", "", ""})
+    end
+  end
+
   describe "byte_index" do
     assert { "foo".byte_index('o'.ord).should eq(1) }
     assert { "foo bar booz".byte_index('o'.ord, 3).should eq(9) }

diff --git a/src/string.cr b/src/string.cr
@@ -2200,6 +2200,45 @@ class String
     match_result.try &.begin(0)
   end
 
+  # Searches sep or pattern (regexp) in the string,
+  # and returns the part before it, the match, and the part after it.
+  # If it is not found, returns str followed by two empty strings.
+  #
+  # ```
+  # "hello".partition("l") # => {"he", "l", "lo"}
+  # "hello".partition("x") # => {"hello", "", ""}
+  # ```
+  def partition(search : (Char | String)) : Tuple(String, String, String)
+    pre = mid = post = ""
+    search_size = search.is_a?(Char) ? 1 : search.size
+    case pos = self.index(search)
+    when .nil?
+      pre = self
+    when 0
+      mid = search.to_s
+      post = self[(pos + search_size)..-1]
+    else
+      pre = self[0..(pos - 1)]
+      mid = search.to_s
+      post = self[(pos + search_size)..-1]
+    end
+    {pre, mid, post}
+  end
+
+  # ditto
+  def partition(search : Regex) : Tuple(String, String, String)
+    pre = mid = post = ""
+    case m = self.match(search)
+    when .nil?
+      pre = self
+    else
+      pre = m.pre_match
+      mid = m[0]
+      post = m.post_match
+    end
+    {pre, mid, post}
+  end
+
   def byte_index(byte : Int, offset = 0)
     offset.upto(bytesize - 1) do |i|
       if to_unsafe[i] == byte