strip, lstrip, rstrip work with NaN (#739)

Rubtsowa · web-flow · commit f5b7dba5fd37 · 2020-03-20T12:27:56.000+03:00
diff --git 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
@@ -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))
