Define __eq__ for basic types.

Author: cthulahoops

tests/snippets/math.py

@@ -11,6 +11,7 @@
 assert a ** 3 == 64
 assert a * 3 == 12
 assert a / 2 == 2
+assert 2 == a / 2
 # assert a % 3 == 1
 assert a - 3 == 1
 assert -a == -4

vm/src/obj/objfloat.rs

@@ -42,6 +42,25 @@ fn set_value(obj: &PyObjectRef, value: f64) {
     obj.borrow_mut().kind = PyObjectKind::Float { value };
 }
 
+fn float_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(zelf, Some(vm.ctx.float_type())), (other, None)]
+    );
+    let zelf = get_value(zelf);
+    let result = if objtype::isinstance(other.clone(), vm.ctx.float_type()) {
+        let other = get_value(other);
+        zelf == other
+    } else if objtype::isinstance(other.clone(), vm.ctx.int_type()) {
+        let other = objint::get_value(other) as f64;
+        zelf == other
+    } else {
+        false
+    };
+    Ok(vm.ctx.new_bool(result))
+}
+
 fn float_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
@@ -97,6 +116,7 @@ fn float_pow(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 
 pub fn init(context: &PyContext) {
     let ref float_type = context.float_type;
+    float_type.set_attr("__eq__", context.new_rustfunc(float_eq));
     float_type.set_attr("__add__", context.new_rustfunc(float_add));
     float_type.set_attr("__init__", context.new_rustfunc(float_init));
     float_type.set_attr("__pow__", context.new_rustfunc(float_pow));

vm/src/obj/objint.rs

@@ -49,6 +49,26 @@ impl FromPyObjectRef for i32 {
     }
 }
 
+fn int_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(zelf, Some(vm.ctx.int_type())), (other, None)]
+    );
+    let result = if objtype::isinstance(other.clone(), vm.ctx.int_type()) {
+        let zelf = i32::from_pyobj(zelf);
+        let other = i32::from_pyobj(other);
+        zelf == other
+    } else if objtype::isinstance(other.clone(), vm.ctx.float_type()) {
+        let zelf = i32::from_pyobj(zelf) as f64;
+        let other = objfloat::get_value(other);
+        zelf == other
+    } else {
+        false
+    };
+    Ok(vm.ctx.new_bool(result))
+}
+
 fn int_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
@@ -144,6 +164,7 @@ fn int_pow(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 
 pub fn init(context: &PyContext) {
     let ref int_type = context.int_type;
+    int_type.set_attr("__eq__", context.new_rustfunc(int_eq));
     int_type.set_attr("__add__", context.new_rustfunc(int_add));
     int_type.set_attr("__init__", context.new_rustfunc(int_init));
     int_type.set_attr("__mod__", context.new_rustfunc(int_mod));

vm/src/obj/objstr.rs

@@ -8,6 +8,7 @@ use super::objtype;
 
 pub fn init(context: &PyContext) {
     let ref str_type = context.str_type;
+    str_type.set_attr("__eq__", context.new_rustfunc(str_eq));
     str_type.set_attr("__add__", context.new_rustfunc(str_add));
     str_type.set_attr("__len__", context.new_rustfunc(str_len));
     str_type.set_attr("__mul__", context.new_rustfunc(str_mul));
@@ -23,6 +24,21 @@ pub fn get_value(obj: &PyObjectRef) -> String {
     }
 }
 
+fn str_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(a, Some(vm.ctx.str_type())), (b, None)]
+    );
+
+    let result = if objtype::isinstance(b.clone(), vm.ctx.str_type()) {
+        get_value(a) == get_value(b)
+    } else {
+        false
+    };
+    Ok(vm.ctx.new_bool(result))
+}
+
 fn str_str(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(s, Some(vm.ctx.str_type()))]);
     Ok(s.clone())

vm/src/obj/objtype.rs

@@ -29,7 +29,6 @@ pub fn init(context: &PyContext) {
 }
 
 fn type_mro(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    println!("{:?}", args);
     arg_check!(
         vm,
         args,

vm/src/objbool.rs

@@ -35,6 +35,22 @@ pub fn init(context: &PyContext) {
     let ref bool_type = context.bool_type;
     bool_type.set_attr("__new__", context.new_rustfunc(bool_new));
     bool_type.set_attr("__str__", context.new_rustfunc(bool_str));
+    bool_type.set_attr("__eq__", context.new_rustfunc(bool_eq));
+}
+
+fn bool_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(zelf, Some(vm.ctx.bool_type())), (other, None)]
+    );
+
+    let result = if objtype::isinstance(zelf.clone(), vm.ctx.bool_type()) {
+        get_value(zelf) == get_value(other)
+    } else {
+        false
+    };
+    Ok(vm.ctx.new_bool(result))
 }
 
 // Retrieve inner int value:

vm/src/objobject.rs

@@ -29,7 +29,16 @@ pub fn create_object(type_type: PyObjectRef, object_type: PyObjectRef, dict_type
     (*object_type.borrow_mut()).typ = Some(type_type.clone());
 }
 
-fn obj_str(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+fn object_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(zelf, Some(vm.ctx.object())), (other, None)]
+    );
+    Ok(vm.ctx.new_bool(zelf.is(other)))
+}
+
+fn object_str(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(obj, Some(vm.ctx.object()))]);
     let type_name = objtype::get_type_name(&obj.typ());
     let address = obj.get_id();
@@ -39,8 +48,9 @@ fn obj_str(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 pub fn init(context: &PyContext) {
     let ref object = context.object;
     object.set_attr("__new__", context.new_rustfunc(new_instance));
+    object.set_attr("__eq__", context.new_rustfunc(object_eq));
     object.set_attr("__dict__", context.new_member_descriptor(object_dict));
-    object.set_attr("__str__", context.new_rustfunc(obj_str));
+    object.set_attr("__str__", context.new_rustfunc(object_str));
 }
 
 fn object_dict(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {