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  return f
10
11# CHECK-LABEL: TEST: testAttributes
12@run
13def testAttributes():
14  with Context() as ctx, Location.unknown():
15    ctx.allow_unregistered_dialects = True
16
17    #
18    # Check op construction with attributes.
19    #
20
21    i32 = IntegerType.get_signless(32)
22    one = IntegerAttr.get(i32, 1)
23    two = IntegerAttr.get(i32, 2)
24    unit = UnitAttr.get()
25
26    # CHECK: "python_test.attributed_op"() {
27    # CHECK-DAG: mandatory_i32 = 1 : i32
28    # CHECK-DAG: optional_i32 = 2 : i32
29    # CHECK-DAG: unit
30    # CHECK: }
31    op = test.AttributedOp(one, two, unit)
32    print(f"{op}")
33
34    # CHECK: "python_test.attributed_op"() {
35    # CHECK: mandatory_i32 = 2 : i32
36    # CHECK: }
37    op2 = test.AttributedOp(two, None, None)
38    print(f"{op2}")
39
40    #
41    # Check generic "attributes" access and mutation.
42    #
43
44    assert "additional" not in op.attributes
45
46    # CHECK: "python_test.attributed_op"() {
47    # CHECK-DAG: additional = 1 : i32
48    # CHECK-DAG: mandatory_i32 = 2 : i32
49    # CHECK: }
50    op2.attributes["additional"] = one
51    print(f"{op2}")
52
53    # CHECK: "python_test.attributed_op"() {
54    # CHECK-DAG: additional = 2 : i32
55    # CHECK-DAG: mandatory_i32 = 2 : i32
56    # CHECK: }
57    op2.attributes["additional"] = two
58    print(f"{op2}")
59
60    # CHECK: "python_test.attributed_op"() {
61    # CHECK-NOT: additional = 2 : i32
62    # CHECK:     mandatory_i32 = 2 : i32
63    # CHECK: }
64    del op2.attributes["additional"]
65    print(f"{op2}")
66
67    try:
68      print(op.attributes["additional"])
69    except KeyError:
70      pass
71    else:
72      assert False, "expected KeyError on unknown attribute key"
73
74    #
75    # Check accessors to defined attributes.
76    #
77
78    # CHECK: Mandatory: 1
79    # CHECK: Optional: 2
80    # CHECK: Unit: True
81    print(f"Mandatory: {op.mandatory_i32.value}")
82    print(f"Optional: {op.optional_i32.value}")
83    print(f"Unit: {op.unit}")
84
85    # CHECK: Mandatory: 2
86    # CHECK: Optional: None
87    # CHECK: Unit: False
88    print(f"Mandatory: {op2.mandatory_i32.value}")
89    print(f"Optional: {op2.optional_i32}")
90    print(f"Unit: {op2.unit}")
91
92    # CHECK: Mandatory: 2
93    # CHECK: Optional: None
94    # CHECK: Unit: False
95    op.mandatory_i32 = two
96    op.optional_i32 = None
97    op.unit = False
98    print(f"Mandatory: {op.mandatory_i32.value}")
99    print(f"Optional: {op.optional_i32}")
100    print(f"Unit: {op.unit}")
101    assert "optional_i32" not in op.attributes
102    assert "unit" not in op.attributes
103
104    try:
105      op.mandatory_i32 = None
106    except ValueError:
107      pass
108    else:
109      assert False, "expected ValueError on setting a mandatory attribute to None"
110
111    # CHECK: Optional: 2
112    op.optional_i32 = two
113    print(f"Optional: {op.optional_i32.value}")
114
115    # CHECK: Optional: None
116    del op.optional_i32
117    print(f"Optional: {op.optional_i32}")
118
119    # CHECK: Unit: False
120    op.unit = None
121    print(f"Unit: {op.unit}")
122    assert "unit" not in op.attributes
123
124    # CHECK: Unit: True
125    op.unit = True
126    print(f"Unit: {op.unit}")
127
128    # CHECK: Unit: False
129    del op.unit
130    print(f"Unit: {op.unit}")
131
132
133# CHECK-LABEL: TEST: inferReturnTypes
134@run
135def inferReturnTypes():
136  with Context() as ctx, Location.unknown(ctx):
137    test.register_python_test_dialect(ctx)
138    module = Module.create()
139    with InsertionPoint(module.body):
140      op = test.InferResultsOp(
141          IntegerType.get_signless(32), IntegerType.get_signless(64))
142      dummy = test.DummyOp()
143
144    # CHECK: [Type(i32), Type(i64)]
145    iface = InferTypeOpInterface(op)
146    print(iface.inferReturnTypes())
147
148    # CHECK: [Type(i32), Type(i64)]
149    iface_static = InferTypeOpInterface(test.InferResultsOp)
150    print(iface.inferReturnTypes())
151
152    assert isinstance(iface.opview, test.InferResultsOp)
153    assert iface.opview == iface.operation.opview
154
155    try:
156      iface_static.opview
157    except TypeError:
158      pass
159    else:
160      assert False, ("not expected to be able to obtain an opview from a static"
161                     " interface")
162
163    try:
164      InferTypeOpInterface(dummy)
165    except ValueError:
166      pass
167    else:
168      assert False, "not expected dummy op to implement the interface"
169
170    try:
171      InferTypeOpInterface(test.DummyOp)
172    except ValueError:
173      pass
174    else:
175      assert False, "not expected dummy op class to implement the interface"
176