xref: /llvm-project/mlir/test/mlir-tblgen/enums-python-bindings.td (revision 92233062c17590d3157bdc6db430fcdfc54312fe)

// RUN: mlir-tblgen -gen-python-enum-bindings %s -I %S/../../include | FileCheck %s

include "mlir/IR/EnumAttr.td"

def Test_Dialect : Dialect {
  let name = "TestDialect";
  let cppNamespace = "::test";
}

// CHECK: Autogenerated by mlir-tblgen; don't manually edit.

// CHECK: from enum import IntEnum, auto, IntFlag
// CHECK: from ._ods_common import _cext as _ods_cext
// CHECK: from ..ir import register_attribute_builder
// CHECK: _ods_ir = _ods_cext.ir

def One : I32EnumAttrCase<"CaseOne", 1, "one">;
def Two : I32EnumAttrCase<"CaseTwo", 2, "two">;
def NegOne : I32EnumAttrCase<"CaseNegOne", -1, "negone">;

def MyEnum : I32EnumAttr<"MyEnum", "An example 32-bit enum", [One, Two, NegOne]>;
// CHECK-LABEL: class MyEnum(IntEnum):
// CHECK:     """An example 32-bit enum"""

// CHECK:     CaseOne = 1
// CHECK:     CaseTwo = 2
// CHECK:     CaseNegOne = auto()

// CHECK:     def __str__(self):
// CHECK:         if self is MyEnum.CaseOne:
// CHECK:             return "one"
// CHECK:         if self is MyEnum.CaseTwo:
// CHECK:             return "two"
// CHECK:         if self is MyEnum.CaseNegOne:
// CHECK:             return "negone"
// CHECK:         raise ValueError("Unknown MyEnum enum entry.")

// CHECK: @register_attribute_builder("MyEnum")
// CHECK: def _myenum(x, context):
// CHECK:     return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))

def TestMyEnum_Attr : EnumAttr<Test_Dialect, MyEnum, "enum">;

def One64 : I64EnumAttrCase<"CaseOne64", 1, "one">;
def Two64 : I64EnumAttrCase<"CaseTwo64", 2, "two">;

def MyEnum64 : I64EnumAttr<"MyEnum64", "An example 64-bit enum", [One64, Two64]>;
// CHECK-LABEL: class MyEnum64(IntEnum):
// CHECK:     """An example 64-bit enum"""

// CHECK:     CaseOne64 = 1
// CHECK:     CaseTwo64 = 2

// CHECK:     def __str__(self):
// CHECK:         if self is MyEnum64.CaseOne64:
// CHECK:             return "one"
// CHECK:         if self is MyEnum64.CaseTwo64:
// CHECK:             return "two"
// CHECK:         raise ValueError("Unknown MyEnum64 enum entry.")

// CHECK: @register_attribute_builder("MyEnum64")
// CHECK: def _myenum64(x, context):
// CHECK:     return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(64, context=context), int(x))

def TestBitEnum
    : I32BitEnumAttr<"TestBitEnum", "", [
        I32BitEnumAttrCaseBit<"User", 0, "user">,
        I32BitEnumAttrCaseBit<"Group", 1, "group">,
        I32BitEnumAttrCaseBit<"Other", 2, "other">,
      ]> {
  let genSpecializedAttr = 0;
  let separator = " | ";
}

def TestBitEnum_Attr : EnumAttr<Test_Dialect, TestBitEnum, "testbitenum">;

// CHECK-LABEL: class TestBitEnum(IntFlag):

// CHECK:     User = 1
// CHECK:     Group = 2
// CHECK:     Other = 4

// CHECK:     def __iter__(self):
// CHECK:         return iter([case for case in type(self) if (self & case) is case])
// CHECK:     def __len__(self):
// CHECK:         return bin(self).count("1")

// CHECK:     def __str__(self):
// CHECK:         if len(self) > 1:
// CHECK:             return " | ".join(map(str, self))
// CHECK:         if self is TestBitEnum.User:
// CHECK:             return "user"
// CHECK:         if self is TestBitEnum.Group:
// CHECK:             return "group"
// CHECK:         if self is TestBitEnum.Other:
// CHECK:             return "other"
// CHECK:         raise ValueError("Unknown TestBitEnum enum entry.")

// CHECK: @register_attribute_builder("TestBitEnum")
// CHECK: def _testbitenum(x, context):
// CHECK:     return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))

// CHECK: @register_attribute_builder("TestBitEnum_Attr")
// CHECK: def _testbitenum_attr(x, context):
// CHECK:     return _ods_ir.Attribute.parse(f'#TestDialect<testbitenum {str(x)}>', context=context)

// CHECK: @register_attribute_builder("TestMyEnum_Attr")
// CHECK: def _testmyenum_attr(x, context):
// CHECK:     return _ods_ir.Attribute.parse(f'#TestDialect<enum {str(x)}>', context=context)