Merge pull request RustPython#4732 from youknowone/into-exact-or

Fix str.__str__ and PyRef::into_exact_or for int() or str() pattern

Author: fanninpm

vm/src/builtins/int.rs

@@ -15,7 +15,7 @@ use crate::{
     },
     protocol::PyNumberMethods,
     types::{AsNumber, Comparable, Constructor, Hashable, PyComparisonOp, Representable},
-    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
+    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
     TryFromBorrowedObject, VirtualMachine,
 };
 use num_bigint::{BigInt, Sign};
@@ -321,7 +321,7 @@ impl PyInt {
 
 #[pyclass(
     flags(BASETYPE),
-    with(Comparable, Hashable, Constructor, AsNumber, Representable)
+    with(PyRef, Comparable, Hashable, Constructor, AsNumber, Representable)
 )]
 impl PyInt {
     #[pymethod(name = "__radd__")]
@@ -521,11 +521,6 @@ impl PyInt {
         Ok(zelf)
     }
 
-    #[pymethod(magic)]
-    fn int(zelf: PyRef<Self>) -> PyRef<Self> {
-        zelf
-    }
-
     #[pymethod(magic)]
     fn pos(&self) -> BigInt {
         self.value.clone()
@@ -537,23 +532,23 @@ impl PyInt {
     }
 
     #[pymethod(magic)]
-    fn trunc(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn trunc(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pymethod(magic)]
-    fn floor(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn floor(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pymethod(magic)]
-    fn ceil(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn ceil(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pymethod(magic)]
-    fn index(zelf: PyRef<Self>) -> PyRef<Self> {
-        zelf
+    fn index(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pymethod(magic)]
@@ -589,8 +584,8 @@ impl PyInt {
     }
 
     #[pymethod]
-    fn conjugate(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn conjugate(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pyclassmethod]
@@ -659,18 +654,9 @@ impl PyInt {
         Ok(bytes.into())
     }
 
-    #[inline]
-    fn clone_if_subclass(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        if zelf.class().is(vm.ctx.types.int_type) {
-            return zelf;
-        }
-
-        vm.ctx.new_bigint(&zelf.value)
-    }
-
     #[pygetset]
-    fn real(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn real(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pygetset]
@@ -679,8 +665,8 @@ impl PyInt {
     }
 
     #[pygetset]
-    fn numerator(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn numerator(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pygetset]
@@ -701,6 +687,17 @@ impl PyInt {
     }
 }
 
+#[pyclass]
+impl PyRef<PyInt> {
+    #[pymethod(magic)]
+    fn int(self, vm: &VirtualMachine) -> PyRefExact<PyInt> {
+        self.into_exact_or(&vm.ctx, |zelf| unsafe {
+            // TODO: this is actually safe. we need better interface
+            PyRefExact::new_unchecked(vm.ctx.new_bigint(&zelf.value))
+        })
+    }
+}
+
 impl Comparable for PyInt {
     fn cmp(
         zelf: &Py<Self>,

vm/src/builtins/str.rs

@@ -371,6 +371,7 @@ impl PyStr {
 #[pyclass(
     flags(BASETYPE),
     with(
+        PyRef,
         AsMapping,
         AsNumber,
         AsSequence,
@@ -488,11 +489,6 @@ impl PyStr {
         Self::repeat(zelf, value.into(), vm)
     }
 
-    #[pymethod(magic)]
-    fn str(zelf: PyRef<Self>) -> PyStrRef {
-        zelf
-    }
-
     #[inline]
     pub(crate) fn repr(&self, vm: &VirtualMachine) -> PyResult<String> {
         rustpython_common::str::repr(self.as_str())
@@ -1279,6 +1275,17 @@ impl PyStr {
     }
 }
 
+#[pyclass]
+impl PyRef<PyStr> {
+    #[pymethod(magic)]
+    fn str(self, vm: &VirtualMachine) -> PyRefExact<PyStr> {
+        self.into_exact_or(&vm.ctx, |zelf| unsafe {
+            // Creating a copy with same kind is safe
+            PyStr::new_str_unchecked(zelf.bytes.to_vec(), zelf.kind.kind()).into_exact_ref(&vm.ctx)
+        })
+    }
+}
+
 impl PyStrRef {
     pub fn concat_in_place(&mut self, other: &str, vm: &VirtualMachine) {
         // TODO: call [A]Rc::get_mut on the str to try to mutate the data in place

vm/src/object/ext.rs

@@ -9,6 +9,7 @@ use crate::common::{
 use crate::{
     builtins::{PyBaseExceptionRef, PyStrInterned, PyType},
     convert::{IntoPyException, ToPyObject, ToPyResult, TryFromObject},
+    vm::Context,
     VirtualMachine,
 };
 use std::{borrow::Borrow, fmt, marker::PhantomData, ops::Deref, ptr::null_mut};
@@ -107,6 +108,20 @@ impl<T: PyPayload> std::borrow::ToOwned for PyExact<T> {
     }
 }
 
+impl<T: PyPayload> PyRef<T> {
+    pub fn into_exact_or(
+        self,
+        ctx: &Context,
+        f: impl FnOnce(Self) -> PyRefExact<T>,
+    ) -> PyRefExact<T> {
+        if self.class().is(T::class(ctx)) {
+            unsafe { PyRefExact::new_unchecked(self) }
+        } else {
+            f(self)
+        }
+    }
+}
+
 /// PyRef but guaranteed not to be a subtype instance
 #[derive(Debug)]
 #[repr(transparent)]