Fix surround substring to surround all substrings (#212)

* surround_substring surround **all** substrings * added unit tests * ruff lint --------- Co-authored-by: Jimmy Jensen <[email protected]>
mrpowers-io · Feb 26, 2024 · a661769 · a661769
1 parent 6bb7fcd
commit a661769
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/quinn/keyword_finder.py b/quinn/keyword_finder.py
@@ -111,10 +111,7 @@ def surround_substring(input: str, substring: str, surround_start: str, surround
     :rtype: str
 
     """
-    index = input.find(substring)
-    res = ""
-    if index == -1:
-        res = input
-    else:
-        res = input[:index] + surround_start + substring + surround_end + input[(index + len(substring)) :]
-    return res
+    return input.replace(
+        substring,
+        surround_start + substring + surround_end,
+    )
diff --git a/tests/test_keyword_finder.py b/tests/test_keyword_finder.py
@@ -17,4 +17,9 @@ def test_keyword_format():
 
 
 def test_surround_substring():
-    print(surround_substring("spark rdd stuff", "rdd", "**", "||"))
+
+    assert "spark **rdd|| stuff" == surround_substring("spark rdd stuff", "rdd", "**", "||")
+    assert "spark **rdd|| stuff with **rdd||" == surround_substring("spark rdd stuff with rdd", "rdd", "**", "||")
+    assert "spark **rdd||dd stuff" == surround_substring("spark rdddd stuff", "rdd", "**", "||")
+
+