[Mlir-commits] [mlir] 9fb1086 - [mlir][python] Add a __contains__ method to the python bindings for DictionaryAttr.

[Adrian Kuegel]

Author: Adrian Kuegel
Date: 2021-10-29T15:19:16+02:00
New Revision: 9fb1086b94f06982e80dcfd985db911bf94599be

URL: https://github.com/llvm/llvm-project/commit/9fb1086b94f06982e80dcfd985db911bf94599be
DIFF: https://github.com/llvm/llvm-project/commit/9fb1086b94f06982e80dcfd985db911bf94599be.diff

LOG: [mlir][python] Add a __contains__ method to the python bindings for DictionaryAttr.

This makes it easier to check in python whether a certain attribute is there.

Differential Revision: https://reviews.llvm.org/D112814

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRAttributes.cpp
    mlir/test/python/ir/attributes.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index 066350a0a8728..7db90ec8dff5d 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -698,7 +698,13 @@ class PyDictAttribute : public PyConcreteAttribute<PyDictAttribute> {
 
   intptr_t dunderLen() { return mlirDictionaryAttrGetNumElements(*this); }
 
+  bool dunderContains(const std::string &name) {
+    return !mlirAttributeIsNull(
+        mlirDictionaryAttrGetElementByName(*this, toMlirStringRef(name)));
+  }
+
   static void bindDerived(ClassTy &c) {
+    c.def("__contains__", &PyDictAttribute::dunderContains);
     c.def("__len__", &PyDictAttribute::dunderLen);
     c.def_static(
         "get",

diff  --git a/mlir/test/python/ir/attributes.py b/mlir/test/python/ir/attributes.py
index 661b3ce1ccaba..9bbf23cf20855 100644
--- a/mlir/test/python/ir/attributes.py
+++ b/mlir/test/python/ir/attributes.py
@@ -340,6 +340,12 @@ def testDictAttr():
     # CHECK: "string"
     print(a['stringattr'])
 
+    # CHECK: True
+    print('stringattr' in a)
+
+    # CHECK: False
+    print('not_in_dict' in a)
+
     # Check that exceptions are raised as expected.
     try:
       _ = a['does_not_exist']