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