Skip to content

Commit 3d14ae1

Browse files
cthulahoopscthulahoops
committed
Introduce IntoPyStringRef to make vm.get_attribute more usable.
1 parent 6c660c9 commit 3d14ae1

File tree

6 files changed

+48
-18
lines changed

6 files changed

+48
-18
lines changed

‎vm/src/builtins.rs‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -310,7 +310,7 @@ fn builtin_getattr(
310310
default: OptionalArg<PyObjectRef>,
311311
vm: &mut VirtualMachine,
312312
) -> PyResult {
313-
let ret = vm.get_attribute(obj.clone(), attr.into_object());
313+
let ret = vm.get_attribute(obj.clone(), attr);
314314
if let OptionalArg::Present(default) = default {
315315
ret.or_else(|ex| catch_attr_exception(ex, default, vm))
316316
} else {
@@ -323,7 +323,7 @@ fn builtin_globals(vm: &mut VirtualMachine, _args: PyFuncArgs) -> PyResult {
323323
}
324324

325325
fn builtin_hasattr(obj: PyObjectRef, attr: PyStringRef, vm: &mut VirtualMachine) -> PyResult<bool> {
326-
if let Err(ex) = vm.get_attribute(obj.clone(), attr.into_object()) {
326+
if let Err(ex) = vm.get_attribute(obj.clone(), attr) {
327327
catch_attr_exception(ex, false, vm)
328328
} else {
329329
Ok(true)
@@ -832,8 +832,7 @@ pub fn builtin_build_class_(vm: &mut VirtualMachine, mut args: PyFuncArgs) -> Py
832832
let bases = vm.context().new_tuple(bases);
833833

834834
// Prepare uses full __getattribute__ resolution chain.
835-
let prepare_name = vm.new_str("__prepare__".to_string());
836-
let prepare = vm.get_attribute(metaclass.clone(), prepare_name)?;
835+
let prepare = vm.get_attribute(metaclass.clone(), "__prepare__")?;
837836
let namespace = vm.invoke(prepare, vec![name_arg.clone(), bases.clone()])?;
838837

839838
let cells = vm.new_dict();

‎vm/src/exceptions.rs‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
use crate::function::PyFuncArgs;
22
use crate::obj::objsequence;
33
use crate::obj::objtype;
4-
use crate::pyobject::{
5-
create_type, AttributeProtocol, PyContext, PyObjectRef, PyResult, TypeProtocol,
6-
};
4+
use crate::pyobject::{create_type, PyContext, PyObjectRef, PyResult, TypeProtocol};
75
use crate::vm::VirtualMachine;
86

97
fn exception_init(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -21,7 +19,7 @@ fn exception_init(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
2119

2220
// Print exception including traceback:
2321
pub fn print_exception(vm: &mut VirtualMachine, exc: &PyObjectRef) {
24-
if let Some(tb) = exc.get_attr("__traceback__") {
22+
if let Ok(tb) = vm.get_attribute(exc.clone(), "__traceback__") {
2523
println!("Traceback (most recent call last):");
2624
if objtype::isinstance(&tb, &vm.ctx.list_type()) {
2725
let mut elements = objsequence::get_elements(&tb).to_vec();
@@ -70,7 +68,7 @@ fn exception_str(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
7068
required = [(exc, Some(vm.ctx.exceptions.exception_type.clone()))]
7169
);
7270
let type_name = objtype::get_type_name(&exc.typ());
73-
let msg = if let Some(m) = exc.get_attr("msg") {
71+
let msg = if let Ok(m) = vm.get_attribute(exc.clone(), "msg") {
7472
match vm.to_pystr(&m) {
7573
Ok(msg) => msg,
7674
_ => "<exception str() failed>".to_string(),

‎vm/src/frame.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -245,8 +245,9 @@ impl Frame {
245245
&exception,
246246
&vm.ctx.exceptions.base_exception_type
247247
));
248-
let traceback_name = vm.new_str("__traceback__".to_string());
249-
let traceback = vm.get_attribute(exception.clone(), traceback_name).unwrap();
248+
let traceback = vm
249+
.get_attribute(exception.clone(), "__traceback__")
250+
.unwrap();
250251
trace!("Adding to traceback: {:?} {:?}", traceback, lineno);
251252
let pos = vm.ctx.new_tuple(vec![
252253
vm.ctx.new_str(filename.clone()),
@@ -1137,7 +1138,6 @@ impl Frame {
11371138

11381139
fn load_attr(&self, vm: &mut VirtualMachine, attr_name: &str) -> FrameResult {
11391140
let parent = self.pop_value();
1140-
let attr_name = vm.new_str(attr_name.to_string());
11411141
let obj = vm.get_attribute(parent, attr_name)?;
11421142
self.push_value(obj);
11431143
Ok(None)

‎vm/src/obj/objstr.rs‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ pub struct PyString {
2424
// TODO: shouldn't be public
2525
pub value: String,
2626
}
27+
pub type PyStringRef = PyRef<PyString>;
2728

2829
impl<T: ToString> From<T> for PyString {
2930
fn from(t: T) -> PyString {
@@ -33,7 +34,36 @@ impl<T: ToString> From<T> for PyString {
3334
}
3435
}
3536

36-
pub type PyStringRef = PyRef<PyString>;
37+
pub trait IntoPyStringRef {
38+
fn into_pystringref(self, vm: &mut VirtualMachine) -> PyResult<PyStringRef>;
39+
}
40+
41+
impl IntoPyStringRef for PyStringRef {
42+
fn into_pystringref(self, _vm: &mut VirtualMachine) -> PyResult<PyStringRef> {
43+
Ok(self)
44+
}
45+
}
46+
47+
impl IntoPyStringRef for String {
48+
fn into_pystringref(self, vm: &mut VirtualMachine) -> PyResult<PyStringRef> {
49+
Ok(PyString { value: self }.into_ref(vm))
50+
}
51+
}
52+
53+
impl IntoPyStringRef for PyObjectRef {
54+
fn into_pystringref(self, vm: &mut VirtualMachine) -> PyResult<PyStringRef> {
55+
TryFromObject::try_from_object(vm, self)
56+
}
57+
}
58+
59+
impl IntoPyStringRef for &str {
60+
fn into_pystringref(self, vm: &mut VirtualMachine) -> PyResult<PyStringRef> {
61+
Ok(PyString {
62+
value: self.to_string(),
63+
}
64+
.into_ref(vm))
65+
}
66+
}
3767

3868
impl PyStringRef {
3969
fn add(self, rhs: PyObjectRef, vm: &mut VirtualMachine) -> PyResult<String> {

‎vm/src/stdlib/dis.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,7 @@ fn dis_dis(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
77
arg_check!(vm, args, required = [(obj, None)]);
88

99
// Method or function:
10-
let code_name = vm.new_str("__code__".to_string());
11-
if let Ok(co) = vm.get_attribute(obj.clone(), code_name) {
10+
if let Ok(co) = vm.get_attribute(obj.clone(), "__code__") {
1211
return dis_disassemble(vm, PyFuncArgs::new(vec![co], vec![]));
1312
}
1413

‎vm/src/vm.rs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ use crate::obj::objgenerator;
2424
use crate::obj::objiter;
2525
use crate::obj::objlist::PyList;
2626
use crate::obj::objsequence;
27-
use crate::obj::objstr::PyStringRef;
27+
use crate::obj::objstr::{IntoPyStringRef, PyStringRef};
2828
use crate::obj::objtuple::PyTuple;
2929
use crate::obj::objtype;
3030
use crate::pyobject::{
@@ -549,9 +549,13 @@ impl VirtualMachine {
549549
}
550550

551551
// get_attribute should be used for full attribute access (usually from user code).
552-
pub fn get_attribute(&mut self, obj: PyObjectRef, attr_name: PyObjectRef) -> PyResult {
552+
pub fn get_attribute<T>(&mut self, obj: PyObjectRef, attr_name: T) -> PyResult
553+
where
554+
T: IntoPyStringRef,
555+
{
556+
let attr_name = attr_name.into_pystringref(self)?;
553557
trace!("vm.__getattribute__: {:?} {:?}", obj, attr_name);
554-
self.call_method(&obj, "__getattribute__", vec![attr_name])
558+
self.call_method(&obj, "__getattribute__", vec![attr_name.into_object()])
555559
}
556560

557561
pub fn set_attr(

0 commit comments

Comments
 (0)