Merge pull request #614 from RustPython/__dir__

Dir

tests/snippets/builtin_dir.py

@@ -5,7 +5,18 @@ class A:
 
 a = A()
 
-assert "test" in dir(a)
+assert "test" in dir(a), "test not in a"
+assert "test" in dir(A), "test not in A"
+
+class B(A):
+	def __dir__(self):
+		return ('q', 'h')
+
+# Gets sorted and turned into a list
+assert ['h', 'q'] == dir(B())
+
+# This calls type.__dir__ so isn't changed (but inheritance works)!
+assert 'test' in dir(A)
 
 import socket
 

vm/src/builtins.rs

@@ -40,18 +40,6 @@ fn dir_locals(vm: &mut VirtualMachine) -> PyObjectRef {
     get_locals(vm)
 }
 
-fn dir_object(vm: &mut VirtualMachine, obj: &PyObjectRef) -> PyObjectRef {
-    // Gather all members here:
-    let attributes = objtype::get_attributes(obj);
-    let mut members: Vec<String> = attributes.into_iter().map(|(n, _o)| n).collect();
-
-    // Sort members:
-    members.sort();
-
-    let members_pystr = members.into_iter().map(|m| vm.ctx.new_str(m)).collect();
-    vm.ctx.new_list(members_pystr)
-}
-
 fn builtin_abs(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(x, None)]);
     match vm.get_method(x.clone(), "__abs__") {
@@ -171,7 +159,9 @@ fn builtin_dir(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         Ok(dir_locals(vm))
     } else {
         let obj = args.args.into_iter().next().unwrap();
-        Ok(dir_object(vm, &obj))
+        let seq = vm.call_method(&obj, "__dir__", vec![])?;
+        let sorted = builtin_sorted(vm, PyFuncArgs::new(vec![seq], vec![]))?;
+        Ok(sorted)
     }
 }
 

vm/src/obj/mod.rs

@@ -18,6 +18,7 @@ pub mod objiter;
 pub mod objlist;
 pub mod objmap;
 pub mod objmemory;
+pub mod objmodule;
 pub mod objnone;
 pub mod objobject;
 pub mod objproperty;

vm/src/obj/objmodule.rs

@@ -0,0 +1,30 @@
+use crate::frame::ScopeRef;
+use crate::pyobject::{
+    DictProtocol, PyContext, PyFuncArgs, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
+};
+use crate::vm::VirtualMachine;
+
+fn module_dir(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(obj, Some(vm.ctx.module_type()))]);
+    let scope = get_scope(obj);
+    let keys = scope
+        .locals
+        .get_key_value_pairs()
+        .iter()
+        .map(|(k, _v)| k.clone())
+        .collect();
+    Ok(vm.ctx.new_list(keys))
+}
+
+pub fn init(context: &PyContext) {
+    let module_type = &context.module_type;
+    context.set_attr(&module_type, "__dir__", context.new_rustfunc(module_dir));
+}
+
+fn get_scope(obj: &PyObjectRef) -> &ScopeRef {
+    if let PyObjectPayload::Module { ref scope, .. } = &obj.payload {
+        &scope
+    } else {
+        panic!("Can't get scope from non-module.")
+    }
+}

vm/src/obj/objobject.rs

@@ -1,8 +1,8 @@
 use super::objstr;
 use super::objtype;
 use crate::pyobject::{
-    AttributeProtocol, DictProtocol, IdProtocol, PyContext, PyFuncArgs, PyObject, PyObjectPayload,
-    PyObjectRef, PyResult, TypeProtocol,
+    AttributeProtocol, DictProtocol, IdProtocol, PyAttributes, PyContext, PyFuncArgs, PyObject,
+    PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
 };
 use crate::vm::VirtualMachine;
 use std::cell::RefCell;
@@ -127,6 +127,18 @@ fn object_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.new_str(format!("<{} object at 0x{:x}>", type_name, address)))
 }
 
+pub fn object_dir(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(obj, None)]);
+
+    let attributes = get_attributes(&obj);
+    Ok(vm.ctx.new_list(
+        attributes
+            .keys()
+            .map(|k| vm.ctx.new_str(k.to_string()))
+            .collect(),
+    ))
+}
+
 fn object_format(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
@@ -161,6 +173,7 @@ pub fn init(context: &PyContext) {
         "__dict__",
         context.new_member_descriptor(object_dict),
     );
+    context.set_attr(&object, "__dir__", context.new_rustfunc(object_dir));
     context.set_attr(&object, "__hash__", context.new_rustfunc(object_hash));
     context.set_attr(&object, "__str__", context.new_rustfunc(object_str));
     context.set_attr(&object, "__repr__", context.new_rustfunc(object_repr));
@@ -238,3 +251,17 @@ fn object_getattribute(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         ))
     }
 }
+
+pub fn get_attributes(obj: &PyObjectRef) -> PyAttributes {
+    // Get class attributes:
+    let mut attributes = objtype::get_attributes(&obj.typ());
+
+    // Get instance attributes:
+    if let PyObjectPayload::Instance { dict } = &obj.payload {
+        for (name, value) in dict.borrow().iter() {
+            attributes.insert(name.to_string(), value.clone());
+        }
+    }
+
+    attributes
+}

vm/src/obj/objtype.rs

@@ -1,9 +1,8 @@
 use super::objdict;
 use super::objstr;
-use super::objtype; // Required for arg_check! to use isinstance
 use crate::pyobject::{
-    AttributeProtocol, DictProtocol, IdProtocol, PyAttributes, PyContext, PyFuncArgs, PyObject,
-    PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
+    AttributeProtocol, IdProtocol, PyAttributes, PyContext, PyFuncArgs, PyObject, PyObjectPayload,
+    PyObjectRef, PyResult, TypeProtocol,
 };
 use crate::vm::VirtualMachine;
 use std::cell::RefCell;
@@ -67,6 +66,7 @@ pub fn init(context: &PyContext) {
         context.new_rustfunc(type_subclass_check),
     );
     context.set_attr(&type_type, "__doc__", context.new_str(type_doc.to_string()));
+    context.set_attr(&type_type, "__dir__", context.new_rustfunc(type_dir));
 }
 
 fn type_mro(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -95,10 +95,6 @@ fn _mro(cls: PyObjectRef) -> Option<Vec<PyObjectRef>> {
     }
 }
 
-pub fn base_classes(obj: &PyObjectRef) -> Vec<PyObjectRef> {
-    _mro(obj.typ()).unwrap()
-}
-
 /// Determines if `obj` actually an instance of `cls`, this doesn't call __instancecheck__, so only
 /// use this if `cls` is known to have not overridden the base __instancecheck__ magic method.
 pub fn isinstance(obj: &PyObjectRef, cls: &PyObjectRef) -> bool {
@@ -238,12 +234,23 @@ pub fn type_getattribute(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult
     }
 }
 
+pub fn type_dir(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(obj, None)]);
+
+    let attributes = get_attributes(&obj);
+    Ok(vm.ctx.new_list(
+        attributes
+            .keys()
+            .map(|k| vm.ctx.new_str(k.to_string()))
+            .collect(),
+    ))
+}
+
 pub fn get_attributes(obj: &PyObjectRef) -> PyAttributes {
     // Gather all members here:
     let mut attributes = PyAttributes::new();
 
-    // Get class attributes:
-    let mut base_classes = objtype::base_classes(obj);
+    let mut base_classes = _mro(obj.clone()).expect("Type get_attributes on non-type");
     base_classes.reverse();
     for bc in base_classes {
         if let PyObjectPayload::Class { dict, .. } = &bc.payload {
@@ -253,19 +260,6 @@ pub fn get_attributes(obj: &PyObjectRef) -> PyAttributes {
         }
     }
 
-    // Get instance attributes:
-    if let PyObjectPayload::Instance { dict } = &obj.payload {
-        for (name, value) in dict.borrow().iter() {
-            attributes.insert(name.to_string(), value.clone());
-        }
-    }
-
-    // Get module attributes:
-    if let PyObjectPayload::Module { ref scope, .. } = &obj.payload {
-        for (name, value) in scope.locals.get_key_value_pairs().iter() {
-            attributes.insert(objstr::get_value(name).to_string(), value.clone());
-        }
-    }
     attributes
 }
 

vm/src/pyobject.rs

@@ -31,6 +31,7 @@ use crate::obj::objiter;
 use crate::obj::objlist;
 use crate::obj::objmap;
 use crate::obj::objmemory;
+use crate::obj::objmodule;
 use crate::obj::objnone;
 use crate::obj::objobject;
 use crate::obj::objproperty;
@@ -317,6 +318,7 @@ impl PyContext {
         objcode::init(&context);
         objframe::init(&context);
         objnone::init(&context);
+        objmodule::init(&context);
         exceptions::init(&context);
         context
     }
@@ -357,6 +359,10 @@ impl PyContext {
         self.list_type.clone()
     }
 
+    pub fn module_type(&self) -> PyObjectRef {
+        self.module_type.clone()
+    }
+
     pub fn set_type(&self) -> PyObjectRef {
         self.set_type.clone()
     }