Merge pull request RustPython#4612 from youknowone/PyObject_Format

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/builtins/bool.rs

@@ -116,7 +116,7 @@ impl PyBool {
 
     #[pymethod(magic)]
     fn format(obj: PyObjectRef, format_spec: PyStrRef, vm: &VirtualMachine) -> PyResult<PyStrRef> {
-        if format_spec.as_str().is_empty() {
+        if format_spec.is_empty() {
             obj.str(vm)
         } else {
             Err(vm.new_type_error("unsupported format string passed to bool.__format__".to_owned()))

vm/src/builtins/int.rs

@@ -1,10 +1,10 @@
-use super::{float, PyByteArray, PyBytes, PyStr, PyStrRef, PyType, PyTypeRef};
+use super::{float, PyByteArray, PyBytes, PyStr, PyType, PyTypeRef};
 use crate::{
     atomic_func,
+    builtins::PyStrRef,
     bytesinner::PyBytesInner,
     class::PyClassImpl,
-    common::format::FormatSpec,
-    common::hash,
+    common::{format::FormatSpec, hash},
     convert::{IntoPyException, ToPyObject, ToPyResult},
     function::{
         ArgByteOrder, ArgIntoBool, OptionalArg, OptionalOption, PyArithmeticValue,

vm/src/builtins/object.rs

@@ -183,6 +183,7 @@ impl PyBaseObject {
     /// Return str(self).
     #[pymethod(magic)]
     fn str(zelf: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyStrRef> {
+        // FIXME: try tp_repr first and fallback to object.__repr__
         zelf.repr(vm)
     }
 
@@ -243,14 +244,13 @@ impl PyBaseObject {
 
     #[pymethod(magic)]
     fn format(obj: PyObjectRef, format_spec: PyStrRef, vm: &VirtualMachine) -> PyResult<PyStrRef> {
-        if format_spec.as_str().is_empty() {
-            obj.str(vm)
-        } else {
-            Err(vm.new_type_error(format!(
+        if !format_spec.is_empty() {
+            return Err(vm.new_type_error(format!(
                 "unsupported format string passed to {}.__format__",
                 obj.class().name()
-            )))
+            )));
         }
+        obj.str(vm)
     }
 
     #[pyslot]

vm/src/builtins/str.rs

@@ -753,10 +753,20 @@ impl PyStr {
     }
 
     #[pymethod(name = "__format__")]
-    fn format_str(&self, spec: PyStrRef, vm: &VirtualMachine) -> PyResult<String> {
-        FormatSpec::parse(spec.as_str())
-            .and_then(|format_spec| format_spec.format_string(self.borrow()))
-            .map_err(|err| err.into_pyexception(vm))
+    fn __format__(zelf: PyRef<Self>, spec: PyStrRef, vm: &VirtualMachine) -> PyResult<PyStrRef> {
+        let spec = spec.as_str();
+        if spec.is_empty() {
+            return if zelf.class().is(vm.ctx.types.str_type) {
+                Ok(zelf)
+            } else {
+                zelf.as_object().str(vm)
+            };
+        }
+
+        let s = FormatSpec::parse(spec)
+            .and_then(|format_spec| format_spec.format_string(zelf.borrow()))
+            .map_err(|err| err.into_pyexception(vm))?;
+        Ok(vm.ctx.new_str(s))
     }
 
     /// Return a titlecased version of the string where words start with an

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, vm.ctx.new_str(format_spec))?;
                 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, spec.downcast::<PyStr>().unwrap())?;
         self.push_value(formatted.into());
         Ok(None)
     }

vm/src/protocol/object.rs

@@ -3,7 +3,7 @@
 
 use crate::{
     builtins::{
-        pystr::IntoPyStrRef, PyBytes, PyDict, PyDictRef, PyGenericAlias, PyInt, PyStrRef,
+        pystr::IntoPyStrRef, PyBytes, PyDict, PyDictRef, PyGenericAlias, PyInt, PyStr, PyStrRef,
         PyTupleRef, PyTypeRef,
     },
     bytesinner::ByteInnerNewOptions,
@@ -330,12 +330,22 @@ impl PyObject {
 
     // Container of the virtual machine state:
     pub fn str(&self, vm: &VirtualMachine) -> PyResult<PyStrRef> {
-        if self.class().is(vm.ctx.types.str_type) {
-            Ok(self.to_owned().downcast().unwrap())
-        } else {
-            let s = vm.call_special_method(self.to_owned(), identifier!(vm, __str__), ())?;
-            s.try_into_value(vm)
-        }
+        let obj = match self.to_owned().downcast_exact::<PyStr>(vm) {
+            Ok(s) => return Ok(s.into_pyref()),
+            Err(obj) => obj,
+        };
+        // TODO: replace to obj.class().slots.str
+        let str_method = match vm.get_special_method(obj, identifier!(vm, __str__))? {
+            Ok(str_method) => str_method,
+            Err(obj) => return obj.repr(vm),
+        };
+        let s = str_method.invoke((), vm)?;
+        s.downcast::<PyStr>().map_err(|obj| {
+            vm.new_type_error(format!(
+                "__str__ returned non-string (type {})",
+                obj.class().name()
+            ))
+        })
     }
 
     // Equivalent to check_class. Masks Attribute errors (into TypeErrors) and lets everything

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.unwrap_or(vm.ctx.new_str("")))
     }
 
     #[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,34 @@ impl VirtualMachine {
             .invoke((), self)
     }
 
+    // PyObject_Format
+    pub fn format(&self, obj: &PyObject, format_spec: PyStrRef) -> PyResult<PyStrRef> {
+        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,