xref: /llvm-project/mlir/test/python/dialects/python_test.py (revision 0f9e6451a836886f39137818c4f0cfd69ae31e62) 
1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4import mlir.dialects.python_test as test
5
6def run(f):
7  print("\nTEST:", f.__name__)
8  f()
9
10# CHECK-LABEL: TEST: testAttributes
11def testAttributes():
12  with Context() as ctx, Location.unknown():
13    ctx.allow_unregistered_dialects = True
14
15    #
16    # Check op construction with attributes.
17    #
18
19    i32 = IntegerType.get_signless(32)
20    one = IntegerAttr.get(i32, 1)
21    two = IntegerAttr.get(i32, 2)
22    unit = UnitAttr.get()
23
24    # CHECK: "python_test.attributed_op"() {
25    # CHECK-DAG: mandatory_i32 = 1 : i32
26    # CHECK-DAG: optional_i32 = 2 : i32
27    # CHECK-DAG: unit
28    # CHECK: }
29    op = test.AttributedOp(one, two, unit)
30    print(f"{op}")
31
32    # CHECK: "python_test.attributed_op"() {
33    # CHECK: mandatory_i32 = 2 : i32
34    # CHECK: }
35    op2 = test.AttributedOp(two, None, None)
36    print(f"{op2}")
37
38    #
39    # Check generic "attributes" access and mutation.
40    #
41
42    assert "additional" not in op.attributes
43
44    # CHECK: "python_test.attributed_op"() {
45    # CHECK-DAG: additional = 1 : i32
46    # CHECK-DAG: mandatory_i32 = 2 : i32
47    # CHECK: }
48    op2.attributes["additional"] = one
49    print(f"{op2}")
50
51    # CHECK: "python_test.attributed_op"() {
52    # CHECK-DAG: additional = 2 : i32
53    # CHECK-DAG: mandatory_i32 = 2 : i32
54    # CHECK: }
55    op2.attributes["additional"] = two
56    print(f"{op2}")
57
58    # CHECK: "python_test.attributed_op"() {
59    # CHECK-NOT: additional = 2 : i32
60    # CHECK:     mandatory_i32 = 2 : i32
61    # CHECK: }
62    del op2.attributes["additional"]
63    print(f"{op2}")
64
65    try:
66      print(op.attributes["additional"])
67    except KeyError:
68      pass
69    else:
70      assert False, "expected KeyError on unknown attribute key"
71
72    #
73    # Check accessors to defined attributes.
74    #
75
76    # CHECK: Mandatory: 1
77    # CHECK: Optional: 2
78    # CHECK: Unit: True
79    print(f"Mandatory: {op.mandatory_i32.value}")
80    print(f"Optional: {op.optional_i32.value}")
81    print(f"Unit: {op.unit}")
82
83    # CHECK: Mandatory: 2
84    # CHECK: Optional: None
85    # CHECK: Unit: False
86    print(f"Mandatory: {op2.mandatory_i32.value}")
87    print(f"Optional: {op2.optional_i32}")
88    print(f"Unit: {op2.unit}")
89
90    # CHECK: Mandatory: 2
91    # CHECK: Optional: None
92    # CHECK: Unit: False
93    op.mandatory_i32 = two
94    op.optional_i32 = None
95    op.unit = False
96    print(f"Mandatory: {op.mandatory_i32.value}")
97    print(f"Optional: {op.optional_i32}")
98    print(f"Unit: {op.unit}")
99    assert "optional_i32" not in op.attributes
100    assert "unit" not in op.attributes
101
102    try:
103      op.mandatory_i32 = None
104    except ValueError:
105      pass
106    else:
107      assert False, "expected ValueError on setting a mandatory attribute to None"
108
109    # CHECK: Optional: 2
110    op.optional_i32 = two
111    print(f"Optional: {op.optional_i32.value}")
112
113    # CHECK: Optional: None
114    del op.optional_i32
115    print(f"Optional: {op.optional_i32}")
116
117    # CHECK: Unit: False
118    op.unit = None
119    print(f"Unit: {op.unit}")
120    assert "unit" not in op.attributes
121
122    # CHECK: Unit: True
123    op.unit = True
124    print(f"Unit: {op.unit}")
125
126    # CHECK: Unit: False
127    del op.unit
128    print(f"Unit: {op.unit}")
129
130run(testAttributes)
131