Preserve str subclasses in ascii() for ASCII-only __repr__ results (#7455)

Lib/test/test_str.py

@@ -112,7 +112,6 @@ class StrTest(string_tests.StringLikeTest,
         # raw strings should not have unicode escapes
         self.assertNotEqual(r"\u0020", " ")
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: <class 'str'> is not <class 'test.test_str.StrSubclass'>
     def test_ascii(self):
         self.assertEqual(ascii('abc'), "'abc'")
         self.assertEqual(ascii('ab\\c'), "'ab\\\\c'")

crates/vm/src/cformat.rs

@@ -28,7 +28,7 @@ fn spec_format_bytes(
             // Unlike strings, %r and %a are identical for bytes: the behaviour corresponds to
             // %a for strings (not %r)
             CFormatConversion::Repr | CFormatConversion::Ascii => {
-                let b = builtins::ascii(obj, vm)?.into();
+                let b = builtins::ascii(obj, vm)?.as_bytes().to_vec();
                 Ok(b)
             }
             CFormatConversion::Str | CFormatConversion::Bytes => {
@@ -132,7 +132,7 @@ fn spec_format_string(
     match &spec.format_type {
         CFormatType::String(conversion) => {
             let result = match conversion {
-                CFormatConversion::Ascii => builtins::ascii(obj, vm)?.into(),
+                CFormatConversion::Ascii => builtins::ascii(obj, vm)?.as_wtf8().to_owned(),
                 CFormatConversion::Str => obj.str(vm)?.as_wtf8().to_owned(),
                 CFormatConversion::Repr => obj.repr(vm)?.as_wtf8().to_owned(),
                 CFormatConversion::Bytes => {

crates/vm/src/format.rs

@@ -159,9 +159,7 @@ fn format_internal(
                 let argument = match conversion_spec.and_then(FormatConversion::from_char) {
                     Some(FormatConversion::Str) => argument.str(vm)?.into(),
                     Some(FormatConversion::Repr) => argument.repr(vm)?.into(),
-                    Some(FormatConversion::Ascii) => {
-                        vm.ctx.new_str(builtins::ascii(argument, vm)?).into()
-                    }
+                    Some(FormatConversion::Ascii) => builtins::ascii(argument, vm)?.into(),
                     Some(FormatConversion::Bytes) => {
                         vm.call_method(&argument, identifier!(vm, decode).as_str(), ())?
                     }

crates/vm/src/frame.rs

@@ -7104,7 +7104,7 @@ impl ExecutingFrame<'_> {
         let value = match conversion {
             ConvertValueOparg::Str => value.str(vm)?.into(),
             ConvertValueOparg::Repr => value.repr(vm)?.into(),
-            ConvertValueOparg::Ascii => vm.ctx.new_str(builtins::ascii(value, vm)?).into(),
+            ConvertValueOparg::Ascii => builtins::ascii(value, vm)?.into(),
             ConvertValueOparg::None => value,
         };
 

crates/vm/src/protocol/object.rs

@@ -370,10 +370,13 @@ impl PyObject {
         })
     }
 
-    pub fn ascii(&self, vm: &VirtualMachine) -> PyResult<ascii::AsciiString> {
+    pub fn ascii(&self, vm: &VirtualMachine) -> PyResult<PyRef<PyStr>> {
         let repr = self.repr(vm)?;
-        let ascii = to_ascii(repr.as_wtf8());
-        Ok(ascii)
+        if repr.as_wtf8().is_ascii() {
+            Ok(repr)
+        } else {
+            Ok(vm.ctx.new_str(to_ascii(repr.as_wtf8())))
+        }
     }
 
     pub fn str_utf8(&self, vm: &VirtualMachine) -> PyResult<PyRef<PyUtf8Str>> {

crates/vm/src/stdlib/builtins.rs

@@ -18,7 +18,7 @@ mod builtins {
             iter::PyCallableIterator,
             list::{PyList, SortOptions},
         },
-        common::{hash::PyHash, str::to_ascii},
+        common::hash::PyHash,
         function::{
             ArgBytesLike, ArgCallable, ArgIndex, ArgIntoBool, ArgIterable, ArgMapping,
             ArgStrOrBytesLike, Either, FsPath, FuncArgs, KwArgs, OptionalArg, OptionalOption,
@@ -64,10 +64,8 @@ mod builtins {
     }
 
     #[pyfunction]
-    pub fn ascii(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<ascii::AsciiString> {
-        let repr = obj.repr(vm)?;
-        let ascii = to_ascii(repr.as_wtf8());
-        Ok(ascii)
+    pub fn ascii(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyStrRef> {
+        obj.ascii(vm)
     }
 
     #[pyfunction]

extra_tests/snippets/builtin_ascii.py

@@ -5,3 +5,23 @@ assert ascii(5) == "5"
 assert ascii(chr(0x10001)) == "'\\U00010001'"
 assert ascii(chr(0x9999)) == "'\\u9999'"
 assert ascii(chr(0x0A)) == "'\\n'"
+
+
+class MyStr(str):
+    pass
+
+
+class Foo:
+    def __repr__(self):
+        return MyStr("hello")
+
+
+class Bar:
+    def __repr__(self):
+        return MyStr("héllo")
+
+
+assert type(ascii(Foo())) is MyStr
+assert ascii(Foo()) == MyStr("hello")
+assert type(ascii(Bar())) is str
+assert ascii(Bar()) == "h\\xe9llo"