use re for str.replace with regex

Author: Ehsan Totoni

hpat/hiframes/hiframes_typed.py

@@ -1,5 +1,6 @@
 import operator
 from collections import defaultdict, namedtuple
+import re
 import numpy as np
 import pandas as pd
 import warnings
@@ -1746,7 +1747,8 @@ def _run_series_str_replace(self, assign, lhs, arr, rhs, nodes):
             extra_globals={'unicode_to_std_str': unicode_to_std_str,
                             'std_str_to_unicode': std_str_to_unicode,
                             'pre_alloc_string_array': pre_alloc_string_array,
-                            'get_utf8_size': get_utf8_size}
+                            'get_utf8_size': get_utf8_size,
+                            're': re}
         )
 
 

hpat/hiframes/series_kernels.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
 import numpy as np
+import re
 
 import numba
 from numba import types
@@ -363,20 +364,38 @@ def _series_astype_str_impl(arr):
     return hpat.hiframes.api.init_series(A)
 
 
+# def _str_replace_regex_impl(str_arr, pat, val):
+#     numba.parfor.init_prange()
+#     e = hpat.str_ext.compile_regex(unicode_to_std_str(pat))
+#     val = unicode_to_std_str(val)
+#     n = len(str_arr)
+#     n_total_chars = 0
+#     str_list = hpat.str_ext.alloc_str_list(n)
+#     for i in numba.parfor.internal_prange(n):
+#         # TODO: support unicode
+#         in_str = unicode_to_std_str(str_arr[i])
+#         out_str = std_str_to_unicode(
+#             hpat.str_ext.str_replace_regex(in_str, e, val))
+#         str_list[i] = out_str
+#         n_total_chars += len(out_str)
+#     numba.parfor.init_prange()
+#     out_arr = pre_alloc_string_array(n, n_total_chars)
+#     for i in numba.parfor.internal_prange(n):
+#         _str = str_list[i]
+#         out_arr[i] = _str
+#     return hpat.hiframes.api.init_series(out_arr)
+
+
 def _str_replace_regex_impl(str_arr, pat, val):
     numba.parfor.init_prange()
-    e = hpat.str_ext.compile_regex(unicode_to_std_str(pat))
-    val = unicode_to_std_str(val)
+    e = re.compile(pat)
     n = len(str_arr)
     n_total_chars = 0
     str_list = hpat.str_ext.alloc_str_list(n)
     for i in numba.parfor.internal_prange(n):
-        # TODO: support unicode
-        in_str = unicode_to_std_str(str_arr[i])
-        out_str = std_str_to_unicode(
-            hpat.str_ext.str_replace_regex(in_str, e, val))
+        out_str = e.sub(val, str_arr[i])
         str_list[i] = out_str
-        n_total_chars += len(out_str)
+        n_total_chars += get_utf8_size(out_str)
     numba.parfor.init_prange()
     out_arr = pre_alloc_string_array(n, n_total_chars)
     for i in numba.parfor.internal_prange(n):

hpat/str_ext.py

@@ -105,7 +105,9 @@ def unbox_re_pattern(typ, obj, c):
     c.pyapi.incref(obj)
     return NativeValue(obj)
 
-@overload(re.compile)
+
+# jitoptions until numba #4020 is resolved
+@overload(re.compile, jit_options={'no_cpython_wrapper': False})
 def re_compile_overload(pattern, flags=0):
     def _re_compile_impl(pattern, flags=0):
         with numba.objmode(pat='re_pattern_type'):

hpat/tests/test_strings.py

@@ -106,7 +106,6 @@ def test_impl(a):
             # XXX: use startswith since hpat output can have extra characters
             self.assertTrue(h_res.startswith(py_res))
 
-    @unittest.skip("pending numba #4020")
     def test_re_sub(self):
         def test_impl(_str):
             p = re.compile('ab*')