From a661769f1f73783971a7eab3ea3bf2a87b737ff1 Mon Sep 17 00:00:00 2001 From: mrjsj Date: Mon, 26 Feb 2024 18:33:48 +0100 Subject: [PATCH] Fix surround substring to surround all substrings (#212) * surround_substring surround **all** substrings * added unit tests * ruff lint --------- Co-authored-by: Jimmy Jensen --- quinn/keyword_finder.py | 11 ++++------- tests/test_keyword_finder.py | 7 ++++++- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/quinn/keyword_finder.py b/quinn/keyword_finder.py index 848021d..79bcd3f 100644 --- 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 index e258d3d..f433afa 100644 --- 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", "**", "||") + +