diff --git a/Lib/test/string_tests.py b/Lib/test/string_tests.py index 3f82b515b..c5831c47f 100644 --- a/Lib/test/string_tests.py +++ b/Lib/test/string_tests.py @@ -154,8 +154,6 @@ class BaseTest: self.assertEqual(rem, 0, '%s != 0 for %s' % (rem, i)) self.assertEqual(r1, r2, '%s != %s for %s' % (r1, r2, i)) - # TODO: RUSTPYTHON; TypeError: Unexpected keyword argument count - @unittest.expectedFailure def test_count_keyword(self): self.assertEqual('aa'.replace('a', 'b', 0), 'aa'.replace('a', 'b', count=0)) self.assertEqual('aa'.replace('a', 'b', 1), 'aa'.replace('a', 'b', count=1)) diff --git a/vm/src/builtins/str.rs b/vm/src/builtins/str.rs index f822a124e..9f86da3da 100644 --- a/vm/src/builtins/str.rs +++ b/vm/src/builtins/str.rs @@ -1018,20 +1018,27 @@ impl PyStr { } #[pymethod] - fn replace(&self, old: PyStrRef, new: PyStrRef, count: OptionalArg) -> Wtf8Buf { + fn replace(&self, args: ReplaceArgs) -> Wtf8Buf { + use std::cmp::Ordering; + let s = self.as_wtf8(); - match count { - OptionalArg::Present(max_count) if max_count >= 0 => { - if max_count == 0 || (s.is_empty() && !old.is_empty()) { - // nothing to do; return the original bytes + let ReplaceArgs { old, new, count } = args; + + match count.cmp(&0) { + Ordering::Less => s.replace(old.as_wtf8(), new.as_wtf8()), + Ordering::Equal => s.to_owned(), + Ordering::Greater => { + let s_is_empty = s.is_empty(); + let old_is_empty = old.is_empty(); + + if s_is_empty && !old_is_empty { s.to_owned() - } else if s.is_empty() && old.is_empty() { + } else if s_is_empty && old_is_empty { new.as_wtf8().to_owned() } else { - s.replacen(old.as_wtf8(), new.as_wtf8(), max_count as usize) + s.replacen(old.as_wtf8(), new.as_wtf8(), count as usize) } } - _ => s.replace(old.as_wtf8(), new.as_wtf8()), } } @@ -1685,6 +1692,18 @@ impl FindArgs { } } +#[derive(FromArgs)] +struct ReplaceArgs { + #[pyarg(positional)] + old: PyStrRef, + + #[pyarg(positional)] + new: PyStrRef, + + #[pyarg(any, default = -1)] + count: isize, +} + pub fn init(ctx: &Context) { PyStr::extend_class(ctx, ctx.types.str_type);