diff --git a/cfbs/utils.py b/cfbs/utils.py index 81cfe335..17d70656 100644 --- a/cfbs/utils.py +++ b/cfbs/utils.py @@ -128,6 +128,24 @@ def strip_left(string, beginning): return string[len(beginning) :] +def strip_right_any(string, suffixes: tuple): + if not string.endswith(suffixes): + return string + + for suffix in suffixes: + if string.endswith(suffix): + return string[0 : -len(suffix)] + + +def strip_left_any(string, prefixes: tuple): + if not string.startswith(prefixes): + return string + + for prefix in prefixes: + if string.startswith(prefix): + return string[len(prefix) :] + + def read_file(path): try: with open(path, "r") as f: diff --git a/tests/test_utils.py b/tests/test_utils.py index 4331788d..042c436b 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -23,7 +23,9 @@ read_json, string_sha256, strip_left, + strip_left_any, strip_right, + strip_right_any, ) @@ -55,6 +57,28 @@ def test_strip_left(): assert strip_left(s, "b") == "abab" +def test_strip_right_any(): + s = "a.b.c" + + assert strip_right_any(s, (".b", ".c")) == "a.b" + assert strip_right_any(s, (".c", ".b")) == "a.b" + + s = "a.b.b" + + assert strip_right_any(s, (".b", ".b")) == "a.b" + + +def test_strip_left_any(): + s = "a.b.c" + + assert strip_left_any(s, ("a.", "b.")) == "b.c" + assert strip_left_any(s, ("b.", "a.")) == "b.c" + + s = "a.a.b" + + assert strip_left_any(s, ("a.", "a.")) == "a.b" + + def test_read_file(): file_path = "tests/sample/sample_dir/sample_file_1.txt" expected_str = "sample_string\n123"