Add vm.format to call __format__ like PyObject_Format

Author: youknowone

Lib/test/test_unicode.py

@@ -2399,8 +2399,6 @@ def test_raiseMemError(self):
             self.assertRaises(MemoryError, alloc)
             self.assertRaises(MemoryError, alloc)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_format_subclass(self):
         class S(str):
             def __str__(self):

vm/src/format.rs

@@ -1,10 +1,10 @@
 use crate::{
-    builtins::{PyBaseExceptionRef, PyStrRef},
+    builtins::PyBaseExceptionRef,
     common::format::*,
     convert::{IntoPyException, ToPyException},
     function::FuncArgs,
     stdlib::builtins,
-    AsObject, PyObject, PyObjectRef, PyResult, VirtualMachine,
+    PyObject, PyResult, VirtualMachine,
 };
 
 impl IntoPyException for FormatSpecError {
@@ -94,7 +94,20 @@ fn format_internal(
                     FormatString::from_str(format_spec).map_err(|e| e.to_pyexception(vm))?;
                 let format_spec = format_internal(vm, &nested_format, field_func)?;
 
-                pystr = call_object_format(vm, argument, *conversion_spec, &format_spec)?;
+                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::Bytes) => {
+                        vm.call_method(&argument, identifier!(vm, decode).as_str(), ())?
+                    }
+                    None => argument,
+                };
+
+                // FIXME: compiler can intern specs using parser tree. Then this call can be interned_str
+                pystr = vm.format(&argument, Some(vm.ctx.intern_str(format_spec).to_owned()))?;
                 pystr.as_ref()
             }
             FormatPart::Literal(literal) => literal,
@@ -158,27 +171,3 @@ pub(crate) fn format_map(
         FieldType::Keyword(keyword) => dict.get_item(&keyword, vm),
     })
 }
-
-pub fn call_object_format(
-    vm: &VirtualMachine,
-    argument: PyObjectRef,
-    conversion_spec: Option<char>,
-    format_spec: &str,
-) -> PyResult<PyStrRef> {
-    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::Bytes) => {
-            vm.call_method(&argument, identifier!(vm, decode).as_str(), ())?
-        }
-        None => argument,
-    };
-    let result = vm.call_special_method(argument, identifier!(vm, __format__), (format_spec,))?;
-    result.downcast().map_err(|result| {
-        vm.new_type_error(format!(
-            "__format__ must return a str, not {}",
-            &result.class().name()
-        ))
-    })
-}

vm/src/frame.rs

@@ -11,7 +11,6 @@ use crate::{
     convert::{IntoObject, ToPyResult},
     coroutine::Coro,
     exceptions::ExceptionCtor,
-    format::call_object_format,
     function::{ArgMapping, Either, FuncArgs},
     protocol::{PyIter, PyIterReturn},
     scope::Scope,
@@ -1766,12 +1765,7 @@ impl ExecutingFrame<'_> {
         };
 
         let spec = self.pop_value();
-        let formatted = call_object_format(
-            vm,
-            value,
-            None,
-            spec.downcast_ref::<PyStr>().unwrap().as_str(),
-        )?;
+        let formatted = vm.format(&value, Some(spec.downcast::<PyStr>().unwrap()))?;
         self.push_value(formatted.into());
         Ok(None)
     }

vm/src/stdlib/builtins.rs

@@ -17,7 +17,6 @@ mod builtins {
         },
         common::{hash::PyHash, str::to_ascii},
         convert::ToPyException,
-        format::call_object_format,
         function::{
             ArgBytesLike, ArgCallable, ArgIntoBool, ArgIterable, ArgMapping, ArgStrOrBytesLike,
             Either, FuncArgs, KwArgs, OptionalArg, OptionalOption, PosArgs, PyArithmeticValue,
@@ -317,11 +316,7 @@ mod builtins {
         format_spec: OptionalArg<PyStrRef>,
         vm: &VirtualMachine,
     ) -> PyResult<PyStrRef> {
-        let format_spec = format_spec
-            .into_option()
-            .unwrap_or_else(|| vm.ctx.empty_str.clone());
-
-        call_object_format(vm, value, None, format_spec.as_str())
+        vm.format(&value, format_spec.into())
     }
 
     #[pyfunction]

vm/src/vm/vm_ops.rs

@@ -1,6 +1,6 @@
 use super::{PyMethod, VirtualMachine};
 use crate::{
-    builtins::{PyInt, PyIntRef, PyStrInterned},
+    builtins::{PyInt, PyIntRef, PyStr, PyStrInterned, PyStrRef},
     function::PyArithmeticValue,
     object::{AsObject, PyObject, PyObjectRef, PyResult},
     protocol::PyIterReturn,
@@ -494,6 +494,37 @@ impl VirtualMachine {
             .invoke((), self)
     }
 
+    // PyObject_Format
+    pub fn format(&self, obj: &PyObject, format_spec: Option<PyStrRef>) -> PyResult<PyStrRef> {
+        let format_spec_is_empty = format_spec
+            .as_ref()
+            .map_or(true, |spec| spec.as_str().is_empty());
+        if format_spec_is_empty {
+            let obj = match obj.to_owned().downcast_exact::<PyStr>(self) {
+                Ok(s) => return Ok(s.into_pyref()),
+                Err(obj) => obj,
+            };
+            if obj.class().is(self.ctx.types.int_type) {
+                return obj.str(self);
+            }
+        }
+        let bound_format = self
+            .get_special_method(obj.to_owned(), identifier!(self, __format__))?
+            .map_err(|_| {
+                self.new_type_error(format!(
+                    "Type {} doesn't define __format__",
+                    obj.class().name()
+                ))
+            })?;
+        let formatted = bound_format.invoke((format_spec,), self)?;
+        formatted.downcast().map_err(|result| {
+            self.new_type_error(format!(
+                "__format__ must return a str, not {}",
+                &result.class().name()
+            ))
+        })
+    }
+
     // https://docs.python.org/3/reference/expressions.html#membership-test-operations
     fn _membership_iter_search(
         &self,