Merge pull request #5 from pavanagrawal123/pavan/fix_replace

add test + fix replace function implementation

Author: rayokota

src/jsonata/functions.py

@@ -755,22 +755,82 @@ def replace(string: Optional[str], pattern: Union[str, re.Pattern], replacement:
         if isinstance(pattern, str):
             if not pattern:
                 raise jexception.JException("Second argument of replace function cannot be an empty string", 0)
-        if limit is None:
-            if isinstance(pattern, str):
-                return re.sub(pattern, str(replacement), string)
-            else:
-                return Functions.safe_replace_all(string, pattern, replacement)
-        else:
+        
+        if limit is not None and limit < 0:
+            raise jexception.JException("Fourth argument of replace function must evaluate to a positive number", 0)
+
+        def string_replacer(match):
+            result = ''
+            position = 0
+            repl = str(replacement)
+            while position < len(repl):
+                index = repl.find('$', position)
+                if index == -1:
+                    result += repl[position:]
+                    break
+                result += repl[position:index]
+                position = index + 1
+                if position < len(repl):
+                    dollar_val = repl[position]
+                    if dollar_val == '$':
+                        result += '$'
+                        position += 1
+                    elif dollar_val == '0':
+                        result += match.group(0)
+                        position += 1
+                    else:
+                        max_digits = len(str(len(match.groups())))
+                        group_num = repl[position:position+max_digits]
+                        if group_num.isdigit():
+                            group_index = int(group_num)
+                            if 0 < group_index <= len(match.groups()):
+                                result += match.group(group_index) or ''
+                                position += len(group_num)
+                            else:
+                                result += '$'
+                        else:
+                            result += '$'
+                else:
+                    result += '$'
+            return result
 
-            if limit < 0:
-                raise jexception.JException("Fourth argument of replace function must evaluate to a positive number", 0)
+        if callable(replacement):
+            replacer = lambda m: replacement(m.groupdict())
+        elif isinstance(replacement, str):
+            replacer = string_replacer
+        else:
+            replacer = lambda m: str(replacement)
 
-            for i in range(0, limit):
-                if isinstance(pattern, str):
-                    string = re.sub(pattern, str(replacement), string, 1)
-                else:
-                    string = Functions.safe_replace_first(string, pattern, str(replacement))
-            return string
+        if isinstance(pattern, str):
+            # Use string methods for literal string patterns
+            result = ''
+            position = 0
+            count = 0
+            while True:
+                if limit is not None and count >= limit:
+                    result += string[position:]
+                    break
+                index = string.find(pattern, position)
+                if index == -1:
+                    result += string[position:]
+                    break
+                result += string[position:index]
+                match = re.match(re.escape(pattern), string[index:])
+                result += replacer(match)
+                position = index + len(pattern)
+                count += 1
+            return result
+        else:
+            # Use regex for pattern objects
+            if limit is None:
+                return Functions.safe_replace_all(string, pattern, replacement)
+            else:
+                count = 0
+                result = string
+                while count < limit:
+                    result = Functions.safe_replace_first(result, pattern, str(replacement))
+                    count += 1
+                return result
 
     #
     # Base64 encode a string

tests/string_test.py

@@ -32,6 +32,12 @@ def test_escape(self):
         assert jsonata.Jsonata("$string($)").evaluate({"a": str('\n')}) == "{\"a\":\"\\n\"}"
         assert jsonata.Jsonata("$string($)").evaluate({"a": "</"}) == "{\"a\":\"</\"}"
 
+    def test_replace(self):
+        assert jsonata.Jsonata("$replace('hello', '.', '')").evaluate(None) == "hello"
+        assert jsonata.Jsonata("$replace('hello', 'l', 'x')").evaluate(None) == "hexxo"
+        assert jsonata.Jsonata("$replace('h.ello', '.', '')").evaluate(None) == "hello"
+        assert jsonata.Jsonata("$replace('h.e.l.l.o', '.', '',2)").evaluate(None) == "hel.l.o"
+
     #
     # Additional $split tests
     #