ffi (Object): make class dict visible in instances (#5843)

apache · Jun 18, 2020 · 9ba98be · 9ba98be
1 parent 082874c
commit 9ba98be
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/python/tvm/runtime/object.py b/python/tvm/runtime/object.py
@@ -46,9 +46,10 @@ def __repr__(self):
         return _ffi_node_api.AsRepr(self)
 
     def __dir__(self):
+        class_names = dir(self.__class__)
         fnames = _ffi_node_api.NodeListAttrNames(self)
         size = fnames(-1)
-        return [fnames(i) for i in range(size)]
+        return sorted([fnames(i) for i in range(size)] + class_names)
 
     def __getattr__(self, name):
         try:

diff --git a/tests/python/unittest/test_node_reflection.py b/tests/python/unittest/test_node_reflection.py
@@ -130,6 +130,12 @@ def test_pass_config():
             "tir.UnrollLoop": 1
         })
 
+def test_dict():
+    x = tvm.tir.const(1) # a class that has Python-defined methods
+    # instances should see the full class dict
+    assert set(dir(x.__class__)) <= set(dir(x))
+
+
 if __name__ == "__main__":
     test_string()
     test_env_func()
@@ -138,3 +144,4 @@ def test_pass_config():
     test_const_saveload_json()
     test_make_sum()
     test_pass_config()
+    test_dict()