Buffer arguments as PyObject

youknowone · youknowone · commit e18dda3e087c · 2023-04-23T10:57:24.000+09:00
diff --git a/vm/src/function/buffer.rs b/vm/src/function/buffer.rs
@@ -2,7 +2,8 @@ use crate::{
     builtins::{PyStr, PyStrRef},
     common::borrow::{BorrowedValue, BorrowedValueMut},
     protocol::PyBuffer,
-    PyObject, PyObjectRef, PyResult, TryFromBorrowedObject, TryFromObject, VirtualMachine,
+    AsObject, PyObject, PyObjectRef, PyResult, TryFromBorrowedObject, TryFromObject,
+    VirtualMachine,
 };
 
 // Python/getargs.c
@@ -57,6 +58,10 @@ impl ArgBytesLike {
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
+
+    pub fn as_object(&self) -> &PyObject {
+        &self.0.obj
+    }
 }
 
 impl From<ArgBytesLike> for PyBuffer {
@@ -126,6 +131,15 @@ pub enum ArgStrOrBytesLike {
     Str(PyStrRef),
 }
 
+impl ArgStrOrBytesLike {
+    pub fn as_object(&self) -> &PyObject {
+        match self {
+            Self::Buf(b) => b.as_object(),
+            Self::Str(s) => s.as_object(),
+        }
+    }
+}
+
 impl TryFromObject for ArgStrOrBytesLike {
     fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
         obj.downcast()
