py/runtime: Allow types to use both .attr and .locals_dict.

Make it possible to proceed to a regular lookup in locals_dict if the
custom type->attr fails.  This allows type->attr to extend rather than
completely replace the lookup in locals_dict.

This is useful for custom builtin classes that have mostly regular methods
but just a few special attributes/properties.  This way, type->attr needs
to deal with the special cases only and the default lookup will be used for
generic methods.

Signed-off-by: Laurens Valk <laurens@pybricks.com>

2 files changed, 14 insertions, 4 deletions

diff --git a/py/obj.h b/py/obj.h
index 11918ba17..7730059e6 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -551,6 +551,7 @@ struct _mp_obj_type_t {
     //
     // dest[0] = MP_OBJ_NULL means load
     //  return: for fail, do nothing
+    //          for fail but continue lookup in locals_dict, dest[1] = MP_OBJ_SENTINEL
     //          for attr, dest[0] = value
     //          for method, dest[0] = method, dest[1] = self
     //
diff --git a/py/runtime.c b/py/runtime.c
index 27e77fc29..0120b70d7 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -1100,12 +1100,20 @@ void mp_load_method_maybe(mp_obj_t obj, qstr attr, mp_obj_t *dest) {
     if (attr == MP_QSTR___next__ && type->iternext != NULL) {
         dest[0] = MP_OBJ_FROM_PTR(&mp_builtin_next_obj);
         dest[1] = obj;
-
-    } else if (type->attr != NULL) {
+        return;
+    }
+    if (type->attr != NULL) {
         // this type can do its own load, so call it
         type->attr(obj, attr, dest);
-
-    } else if (type->locals_dict != NULL) {
+        // If type->attr has set dest[1] = MP_OBJ_SENTINEL, we should proceed
+        // with lookups below (i.e. in locals_dict). If not, return right away.
+        if (dest[1] != MP_OBJ_SENTINEL) {
+            return;
+        }
+        // Clear the fail flag set by type->attr so it's like it never ran.
+        dest[1] = MP_OBJ_NULL;
+    }
+    if (type->locals_dict != NULL) {
         // generic method lookup
         // this is a lookup in the object (ie not class or type)
         assert(type->locals_dict->base.type == &mp_type_dict); // MicroPython restriction, for now
@@ -1114,6 +1122,7 @@ void mp_load_method_maybe(mp_obj_t obj, qstr attr, mp_obj_t *dest) {
         if (elem != NULL) {
             mp_convert_member_lookup(obj, type, elem->value, dest);
         }
+        return;
     }
 }