diff --git a/sdc/datatypes/hpat_pandas_stringmethods_functions.py b/sdc/datatypes/hpat_pandas_stringmethods_functions.py index e5cefd569..e322a12eb 100644 --- a/sdc/datatypes/hpat_pandas_stringmethods_functions.py +++ b/sdc/datatypes/hpat_pandas_stringmethods_functions.py @@ -1002,15 +1002,19 @@ def strip_usecase(s, to_strip): def gen_sdc_pandas_series_str_strip_impl(usecase): """Generate series.str.lstrip/rstrip/strip implementations based on usecase func""" def impl(self, to_strip=None): + mask = get_nan_mask(self._data._data) item_count = len(self._data) - result = [''] * item_count + res_list = [''] * item_count for it in range(item_count): item = self._data._data[it] if len(item) > 0: - result[it] = usecase(item, to_strip) + res_list[it] = usecase(item, to_strip) else: - result[it] = item + res_list[it] = item + + str_arr = create_str_arr_from_list(res_list) + result = str_arr_set_na_by_mask(str_arr, mask) return pandas.Series(result, self._data._index, name=self._data._name) @@ -1192,19 +1196,19 @@ def hpat_pandas_stringmethods_strip(self, to_strip=None): 'method': hpat_pandas_stringmethods_strip, 'caption': 'Remove leading and trailing characters.', 'seealso': seealso_strip_methods, - 'limitations': limitation_nans_unsupported + 'limitations': '' }, 'lstrip': { 'method': hpat_pandas_stringmethods_lstrip, 'caption': 'Remove leading and trailing characters.', 'seealso': seealso_strip_methods, - 'limitations': limitation_nans_unsupported + 'limitations': '' }, 'rstrip': { 'method': hpat_pandas_stringmethods_rstrip, 'caption': 'Remove leading and trailing characters.', 'seealso': seealso_strip_methods, - 'limitations': limitation_nans_unsupported + 'limitations': '' } } diff --git a/sdc/tests/test_series.py b/sdc/tests/test_series.py index 58c5c25d8..aebbaabaa 100644 --- a/sdc/tests/test_series.py +++ b/sdc/tests/test_series.py @@ -5977,19 +5977,19 @@ def test_series_islower_str(self): pd.testing.assert_series_equal(cfunc(S), islower_usecase(S)) def test_series_strip_str(self): - s = pd.Series(['1. Ant. ', '2. Bee!\n', '3. Cat?\t']) + s = pd.Series(['1. Ant. ', None, '2. Bee!\n', np.nan, '3. Cat?\t']) cfunc = self.jit(strip_usecase) for to_strip in [None, '123.', '.!? \n\t', '123.!? \n\t']: pd.testing.assert_series_equal(cfunc(s, to_strip), strip_usecase(s, to_strip)) def test_series_lstrip_str(self): - s = pd.Series(['1. Ant. ', '2. Bee!\n', '3. Cat?\t']) + s = pd.Series(['1. Ant. ', None, '2. Bee!\n', np.nan, '3. Cat?\t']) cfunc = self.jit(lstrip_usecase) for to_strip in [None, '123.', '.!? \n\t', '123.!? \n\t']: pd.testing.assert_series_equal(cfunc(s, to_strip), lstrip_usecase(s, to_strip)) def test_series_rstrip_str(self): - s = pd.Series(['1. Ant. ', '2. Bee!\n', '3. Cat?\t']) + s = pd.Series(['1. Ant. ', None, '2. Bee!\n', np.nan, '3. Cat?\t']) cfunc = self.jit(rstrip_usecase) for to_strip in [None, '123.', '.!? \n\t', '123.!? \n\t']: pd.testing.assert_series_equal(cfunc(s, to_strip), rstrip_usecase(s, to_strip))