Implement Series.str.center() (#385)

Author: densmirn

sdc/datatypes/hpat_pandas_stringmethods_functions.py

@@ -189,6 +189,53 @@ def hpat_pandas_stringmethods_{methodname}_impl(self{methodparams}):
 """
 
 
+@overload_method(StringMethodsType, 'center')
+def hpat_pandas_stringmethods_center(self, width, fillchar=' '):
+    """
+    Pandas Series method :meth:`pandas.core.strings.StringMethods.center()` implementation.
+
+    Note: Unicode type of list elements are supported only. Numpy.NaN is not supported as elements.
+
+    .. only:: developer
+
+    Test: python -m sdc.runtests -k sdc.tests.test_series.TestSeries.test_series_center
+
+    Parameters
+    ----------
+    self: :class:`pandas.core.strings.StringMethods`
+        input arg
+    width: :obj:`int`
+        Minimum width of resulting string
+    fillchar: :obj:`str`
+        Additional character for filling, default is whitespace
+
+    Returns
+    -------
+    :obj:`pandas.Series`
+         returns :obj:`pandas.Series` object
+    """
+
+    ty_checker = TypeChecker('Method center().')
+    ty_checker.check(self, StringMethodsType)
+
+    if not isinstance(width, Integer):
+        ty_checker.raise_exc(width, 'int', 'width')
+
+    accepted_types = (Omitted, StringLiteral, UnicodeType)
+    if not isinstance(fillchar, accepted_types) and fillchar != ' ':
+        ty_checker.raise_exc(fillchar, 'str', 'fillchar')
+
+    def hpat_pandas_stringmethods_center_impl(self, width, fillchar=' '):
+        item_count = len(self._data)
+        result = [''] * item_count
+        for idx, item in enumerate(self._data._data):
+            result[idx] = item.center(width, fillchar)
+
+        return pandas.Series(result, self._data._index, name=self._data._name)
+
+    return hpat_pandas_stringmethods_center_impl
+
+
 @overload_method(StringMethodsType, 'endswith')
 def hpat_pandas_stringmethods_endswith(self, pat, na=None):
     """

sdc/hiframes/pd_series_ext.py

@@ -758,8 +758,10 @@ def resolve_head(self, ary, args, kws):
     Functions which are still overloaded by HPAT compiler pipeline
 """
 
-str2str_methods_excluded = ['upper', 'endswith', 'find', 'isupper', 'len',
-                            'lower', 'lstrip', 'rstrip', 'startswith', 'strip']
+str2str_methods_excluded = [
+    'upper', 'center', 'endswith', 'find', 'isupper', 'len',
+    'lower', 'lstrip', 'rstrip', 'startswith', 'strip'
+]
 """
     Functions which are used from Numba directly by calling from StringMethodsType
 

sdc/tests/test_series.py

@@ -2424,6 +2424,67 @@ def test_impl(S1, S2):
                 hpat_func(S1, S2), test_impl(S1, S2),
                 err_msg='S1={}\nS2={}'.format(S1, S2))
 
+    def test_series_str_center_default_fillchar(self):
+        def test_impl(series, width):
+            return series.str.center(width)
+
+        hpat_func = self.jit(test_impl)
+
+        data = test_global_input_data_unicode_kind1
+        series = pd.Series(data)
+        width = max(len(s) for s in data) + 10
+
+        pd.testing.assert_series_equal(hpat_func(series, width),
+                                       test_impl(series, width))
+
+    def test_series_str_center(self):
+        def test_impl(series, width, fillchar):
+            return series.str.center(width, fillchar)
+
+        hpat_func = self.jit(test_impl)
+
+        data = test_global_input_data_unicode_kind1
+        data_lengths = [len(s) for s in data]
+        widths = [max(data_lengths) + 10, min(data_lengths)]
+
+        for index in [None, list(range(len(data)))[::-1], data[::-1]]:
+            series = pd.Series(data, index, name='A')
+            for width, fillchar in product(widths, ['\t']):
+                jit_result = hpat_func(series, width, fillchar)
+                ref_result = test_impl(series, width, fillchar)
+                pd.testing.assert_series_equal(jit_result, ref_result)
+
+    def test_series_str_center_exception_unsupported_fillchar(self):
+        def test_impl(series, width, fillchar):
+            return series.str.center(width, fillchar)
+
+        hpat_func = self.jit(test_impl)
+
+        data = test_global_input_data_unicode_kind1
+        series = pd.Series(data)
+        width = max(len(s) for s in data) + 10
+
+        with self.assertRaises(TypingError) as raises:
+            hpat_func(series, width, 10)
+        msg_tmpl = 'Method center(). The object fillchar\n {}'
+        msg = msg_tmpl.format('given: int64\n expected: str')
+        self.assertIn(msg, str(raises.exception))
+
+    def test_series_str_center_exception_unsupported_kind4(self):
+        def test_impl(series, width):
+            return series.str.center(width)
+
+        hpat_func = self.jit(test_impl)
+
+        data = test_global_input_data_unicode_kind4
+        series = pd.Series(data)
+        width = max(len(s) for s in data) + 10
+
+        with self.assertRaises(SystemError) as raises:
+            hpat_func(series, width)
+        msg = 'NULL object passed to Py_BuildValue'
+        self.assertIn(msg, str(raises.exception))
+
     def test_series_str_endswith(self):
         def test_impl(series, pat):
             return series.str.endswith(pat)