Skip to content

Commit ea0757f

Browse files
sinhrkssinhrks
committed
ENH: Add StringMethods.partition and rpartition
1 parent cb8c130 commit ea0757f

File tree

5 files changed

+128
-11
lines changed

5 files changed

+128
-11
lines changed

‎doc/source/api.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -539,10 +539,12 @@ strings and apply several methods to it. These can be acccessed like
539539
Series.str.lstrip
540540
Series.str.match
541541
Series.str.pad
542+
Series.str.partition
542543
Series.str.repeat
543544
Series.str.replace
544545
Series.str.rfind
545546
Series.str.rjust
547+
Series.str.rpartition
546548
Series.str.rstrip
547549
Series.str.slice
548550
Series.str.slice_replace

‎doc/source/text.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,8 @@ Method Summary
229229
:meth:`~Series.str.strip`,Equivalent to ``str.strip``
230230
:meth:`~Series.str.rstrip`,Equivalent to ``str.rstrip``
231231
:meth:`~Series.str.lstrip`,Equivalent to ``str.lstrip``
232+
:meth:`~Series.str.partition`,Equivalent to ``str.partition``
233+
:meth:`~Series.str.rpartition`,Equivalent to ``str.rpartition``
232234
:meth:`~Series.str.lower`,Equivalent to ``str.lower``
233235
:meth:`~Series.str.upper`,Equivalent to ``str.upper``
234236
:meth:`~Series.str.find`,Equivalent to ``str.find``

‎doc/source/whatsnew/v0.16.1.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ Enhancements
1919

2020

2121

22-
22+
- Added ``StringMethods.partition()`` and ``rpartition()`` which behave as the same as standard ``str`` (:issue:`xxxx`)
2323

2424

2525
.. _whatsnew_0161.api:

‎pandas/core/strings.py

Lines changed: 45 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -622,6 +622,19 @@ def str_pad(arr, width, side='left', fillchar=' '):
622622
return _na_map(f, arr)
623623

624624

625+
def _return_type_wrapper(f, arr, return_type):
626+
if return_type not in ('series', 'frame'):
627+
raise ValueError("return_type must be {'series', 'frame'}")
628+
629+
if return_type == 'frame':
630+
from pandas.core.frame import DataFrame
631+
from pandas.core.series import Series
632+
return DataFrame((Series(x) for x in _na_map(f, arr)), index=arr.index)
633+
634+
else:
635+
return _na_map(f, arr)
636+
637+
625638
def str_split(arr, pat=None, n=None, return_type='series'):
626639
"""
627640
Split each string (a la re.split) in array by given pattern, propagating NA
@@ -644,11 +657,6 @@ def str_split(arr, pat=None, n=None, return_type='series'):
644657
-------
645658
split : array
646659
"""
647-
from pandas.core.series import Series
648-
from pandas.core.frame import DataFrame
649-
650-
if return_type not in ('series', 'frame'):
651-
raise ValueError("return_type must be {'series', 'frame'}")
652660
if pat is None:
653661
if n is None or n == 0:
654662
n = -1
@@ -663,11 +671,7 @@ def str_split(arr, pat=None, n=None, return_type='series'):
663671
n = 0
664672
regex = re.compile(pat)
665673
f = lambda x: regex.split(x, maxsplit=n)
666-
if return_type == 'frame':
667-
res = DataFrame((Series(x) for x in _na_map(f, arr)), index=arr.index)
668-
else:
669-
res = _na_map(f, arr)
670-
return res
674+
return _return_type_wrapper(f, arr, return_type)
671675

672676

673677
def str_slice(arr, start=None, stop=None, step=None):
@@ -978,6 +982,37 @@ def split(self, pat=None, n=-1, return_type='series'):
978982
result = str_split(self.series, pat, n=n, return_type=return_type)
979983
return self._wrap_result(result)
980984

985+
_shared_docs['str_partition'] = ("""
986+
Split the string at the %(side)s occurrence of sep, and return a 3-tuple containing the part
987+
before the separator, the separator itself, and the part after the separator.
988+
If the separator is not found, return %(return)s.
989+
990+
Parameters
991+
----------
992+
pat : string, default whitespace
993+
String to split on.
994+
return_type : {'series', 'frame'}, default 'series
995+
If frame, returns a DataFrame (elements are strings)
996+
If series, returns an Series (elements are lists of strings).
997+
998+
Returns
999+
-------
1000+
split : array
1001+
""")
1002+
@Appender(_shared_docs['str_partition'] % {'side': 'first',
1003+
'return': 'a 3-tuple containing the string itself, followed by two empty strings'})
1004+
def partition(self, pat=' ', return_type='series'):
1005+
f = lambda x: x.partition(pat)
1006+
result = _return_type_wrapper(f, self.series, return_type)
1007+
return self._wrap_result(result)
1008+
1009+
@Appender(_shared_docs['str_partition'] % {'side': 'last',
1010+
'return': 'a 3-tuple containing two empty strings, followed by the string itself'})
1011+
def rpartition(self, pat=' ', return_type='series'):
1012+
f = lambda x: x.rpartition(pat)
1013+
result = _return_type_wrapper(f, self.series, return_type)
1014+
return self._wrap_result(result)
1015+
9811016
@copy(str_get)
9821017
def get(self, i):
9831018
result = str_get(self.series, i)

‎pandas/tests/test_strings.py

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -617,6 +617,8 @@ def test_empty_str_methods(self):
617617
tm.assert_series_equal(empty_str, empty.str.pad(42))
618618
tm.assert_series_equal(empty_str, empty.str.center(42))
619619
tm.assert_series_equal(empty_list, empty.str.split('a'))
620+
tm.assert_series_equal(empty_list, empty.str.partition('a'))
621+
tm.assert_series_equal(empty_list, empty.str.rpartition('a'))
620622
tm.assert_series_equal(empty_str, empty.str.slice(stop=1))
621623
tm.assert_series_equal(empty_str, empty.str.slice(step=1))
622624
tm.assert_series_equal(empty_str, empty.str.strip())
@@ -1125,6 +1127,82 @@ def test_split_to_dataframe(self):
11251127
with tm.assertRaisesRegexp(ValueError, "return_type must be"):
11261128
s.str.split('_', return_type="some_invalid_type")
11271129

1130+
def test_partition(self):
1131+
values = Series(['a_b_c', 'c_d_e', NA, 'f_g_h'])
1132+
1133+
result = values.str.partition('_')
1134+
exp = Series([['a', '_', 'b_c'], ['c', '_', 'd_e'], NA, ['f', '_', 'g_h']])
1135+
tm.assert_series_equal(result, exp)
1136+
1137+
result = values.str.rpartition('_')
1138+
exp = Series([['a_b', '_', 'c'], ['c_d', '_', 'e'], NA, ['f_g', '_', 'h']])
1139+
tm.assert_series_equal(result, exp)
1140+
1141+
# more than one char
1142+
values = Series(['a__b__c', 'c__d__e', NA, 'f__g__h'])
1143+
result = values.str.partition('__')
1144+
exp = Series([['a', '__', 'b__c'], ['c', '__', 'd__e'], NA, ['f', '__', 'g__h']])
1145+
tm.assert_series_equal(result, exp)
1146+
1147+
result = values.str.rpartition('__')
1148+
exp = Series([['a__b', '__', 'c'], ['c__d', '__', 'e'], NA, ['f__g', '__', 'h']])
1149+
tm.assert_series_equal(result, exp)
1150+
1151+
# None
1152+
values = Series(['a b c', 'c d e', NA, 'f g h'])
1153+
result = values.str.partition()
1154+
exp = Series([['a', ' ', 'b c'], ['c', ' ', 'd e'], NA, ['f', ' ', 'g h']])
1155+
tm.assert_series_equal(result, exp)
1156+
1157+
result = values.str.rpartition()
1158+
exp = Series([['a b', ' ', 'c'], ['c d', ' ', 'e'], NA, ['f g', ' ', 'h']])
1159+
tm.assert_series_equal(result, exp)
1160+
1161+
# Not splited
1162+
values = Series(['abc', 'cde', NA, 'fgh'])
1163+
result = values.str.partition('_')
1164+
exp = Series([['abc', '', ''], ['cde', '', ''], NA, ['fgh', '', '']])
1165+
tm.assert_series_equal(result, exp)
1166+
1167+
result = values.str.rpartition('_')
1168+
exp = Series([['', '', 'abc'], ['', '', 'cde'], NA, ['', '', 'fgh']])
1169+
tm.assert_series_equal(result, exp)
1170+
1171+
# unicode
1172+
values = Series([u('a_b_c'), u('c_d_e'), NA, u('f_g_h')])
1173+
1174+
result = values.str.partition('_')
1175+
exp = Series([[u('a'), u('_'), u('b_c')], [u('c'), u('_'), u('d_e')],
1176+
NA, [u('f'), u('_'), u('g_h')]])
1177+
tm.assert_series_equal(result, exp)
1178+
1179+
result = values.str.rpartition('_')
1180+
exp = Series([[u('a_b'), u('_'), u('c')], [u('c_d'), u('_'), u('e')],
1181+
NA, [u('f_g'), u('_'), u('h')]])
1182+
tm.assert_series_equal(result, exp)
1183+
1184+
# compare to standard lib
1185+
values = Series(['A_B_C', 'B_C_D', 'E_F_G', 'EFGHEF'])
1186+
self.assertEqual(values.str.partition('_').tolist(), [v.partition('_') for v in values])
1187+
self.assertEqual(values.str.rpartition('_').tolist(), [v.rpartition('_') for v in values])
1188+
1189+
def test_partition_to_dataframe(self):
1190+
values = Series(['a_b_c', 'c_d_e', NA, 'f_g_h'])
1191+
result = values.str.partition('_', return_type='frame')
1192+
exp = DataFrame({0: ['a', 'c', np.nan, 'f'],
1193+
1: ['_', '_', np.nan, '_'],
1194+
2: ['b_c', 'd_e', np.nan, 'g_h']})
1195+
tm.assert_frame_equal(result, exp)
1196+
1197+
result = values.str.rpartition('_', return_type='frame')
1198+
exp = DataFrame({0: ['a_b', 'c_d', np.nan, 'f_g'],
1199+
1: ['_', '_', np.nan, '_'],
1200+
2: ['c', 'e', np.nan, 'h']})
1201+
tm.assert_frame_equal(result, exp)
1202+
1203+
with tm.assertRaisesRegexp(ValueError, "return_type must be"):
1204+
values.str.partition('_', return_type="some_invalid_type")
1205+
11281206
def test_pipe_failures(self):
11291207
# #2119
11301208
s = Series(['A|B|C'])

0 commit comments

Comments
 (0)