Fix issue with slash at the end

AceLewis · Nov 11, 2018 · fc6c57e · fc6c57e
1 parent 459099c
commit fc6c57e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/spintax/spintax.py b/spintax/spintax.py
@@ -17,8 +17,9 @@ def _replace_string(match):
     """
     global spintax_seperator, random_string
     test_string = re.sub(spintax_seperator, lambda x: x.group(1)+random_string, match.group(2))
-    random_picked = random.choice(re.split(random_string, test_string))
-    return match.group(1) + random_picked + match.group(3)
+    split_strings = re.split(random_string, test_string)
+    random_picked = random.choice(split_strings)
+    return match.group(1) + random_picked + ['', match.group(3)][random_picked==split_strings[-1]]
 
 
 def spin(string, seed=None):

diff --git a/tests/spintax_tests.py b/tests/spintax_tests.py
@@ -46,6 +46,14 @@ def test_newline(self):
     def test_escaped_pipe_at_end(self):
         set_up = set([spintax.spin("{a\||b\|}") for a in range(50)])
         self.assertEqual(set_up, {'b|', 'a|'})
+
+    def test_escaped_at_end(self):
+        set_up = set([spintax.spin(r"{|a\\}") for a in range(50)])
+        self.assertEqual(set_up, {'', 'a\\'})
+
+    def test_escaped_at_end_only_one(self):
+        set_up = set([spintax.spin(r"{a\\}") for a in range(50)])
+        self.assertEqual(set_up, {'a\\'})
 
 
 if __name__ == '__main__':