From 9ba98beea56c156d841c501a3bc2b91c78996375 Mon Sep 17 00:00:00 2001 From: Thomas Viehmann Date: Thu, 18 Jun 2020 20:14:21 +0200 Subject: [PATCH] ffi (Object): make class dict visible in instances (#5843) --- python/tvm/runtime/object.py | 3 ++- tests/python/unittest/test_node_reflection.py | 7 +++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/python/tvm/runtime/object.py b/python/tvm/runtime/object.py index ec9f22b82c25..2a34b3455492 100644 --- 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 index b10951691715..3a7318c75761 100644 --- 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()