From dd288904dbaaa6f085252b7457dd10e5abfdb1f2 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Fri, 28 Sep 2018 23:16:56 +1000
Subject: [PATCH] py/objtype: Support full object model for get/set/delitem
 special meths.

This makes these special methods have the same calling behaviour as other
methods in a class instance (mp_convert_member_lookup() is already called
by mp_obj_class_lookup()).
---
 py/objtype.c                            | 15 ++++-----------
 tests/basics/class_staticclassmethod.py | 15 +++++++++++++++
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/py/objtype.c b/py/objtype.c
index f84e2d7465..5499196923 100644
--- a/py/objtype.c
+++ b/py/objtype.c
@@ -797,36 +797,29 @@ STATIC void mp_obj_instance_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
 
 STATIC mp_obj_t instance_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
     mp_obj_instance_t *self = MP_OBJ_TO_PTR(self_in);
-    mp_obj_t member[2] = {MP_OBJ_NULL};
+    mp_obj_t member[4] = {MP_OBJ_NULL, MP_OBJ_NULL, index, value};
     struct class_lookup_data lookup = {
         .obj = self,
         .meth_offset = offsetof(mp_obj_type_t, subscr),
         .dest = member,
         .is_type = false,
     };
-    size_t meth_args;
     if (value == MP_OBJ_NULL) {
         // delete item
         lookup.attr = MP_QSTR___delitem__;
-        mp_obj_class_lookup(&lookup, self->base.type);
-        meth_args = 2;
     } else if (value == MP_OBJ_SENTINEL) {
         // load item
         lookup.attr = MP_QSTR___getitem__;
-        mp_obj_class_lookup(&lookup, self->base.type);
-        meth_args = 2;
     } else {
         // store item
         lookup.attr = MP_QSTR___setitem__;
-        mp_obj_class_lookup(&lookup, self->base.type);
-        meth_args = 3;
     }
+    mp_obj_class_lookup(&lookup, self->base.type);
     if (member[0] == MP_OBJ_SENTINEL) {
         return mp_obj_subscr(self->subobj[0], index, value);
     } else if (member[0] != MP_OBJ_NULL) {
-        mp_obj_t args[3] = {self_in, index, value};
-        // TODO probably need to call mp_convert_member_lookup, and use mp_call_method_n_kw
-        mp_obj_t ret = mp_call_function_n_kw(member[0], meth_args, 0, args);
+        size_t n_args = value == MP_OBJ_NULL || value == MP_OBJ_SENTINEL ? 1 : 2;
+        mp_obj_t ret = mp_call_method_n_kw(n_args, 0, member);
         if (value == MP_OBJ_SENTINEL) {
             return ret;
         } else {
diff --git a/tests/basics/class_staticclassmethod.py b/tests/basics/class_staticclassmethod.py
index 1cb59d5c7b..edde419271 100644
--- a/tests/basics/class_staticclassmethod.py
+++ b/tests/basics/class_staticclassmethod.py
@@ -17,9 +17,24 @@ class C:
     def __add__(self, rhs):
         print('add', rhs)
 
+    # subscript special methods wrapped in staticmethod
+    @staticmethod
+    def __getitem__(item):
+        print('static get', item)
+        return 'item'
+    @staticmethod
+    def __setitem__(item, value):
+        print('static set', item, value)
+    @staticmethod
+    def __delitem__(item):
+        print('static del', item)
+
 c = C()
 
 c.f(0)
 c.g(0)
 c - 1
 c + 2
+print(c[1])
+c[1] = 2
+del c[3]