mlir/test/mlir-tblgen/enums-python-bindings.td

   1 // RUN: mlir-tblgen -gen-python-enum-bindings %s -I %S/../../include | FileCheck %s
   2
   3 include "mlir/IR/EnumAttr.td"
   4
   5 def Test_Dialect : Dialect {
   6   let name = "TestDialect";
   7   let cppNamespace = "::test";
   8 }
   9
  10 // CHECK: Autogenerated by mlir-tblgen; don't manually edit.
  11
  12 // CHECK: from enum import IntEnum, auto, IntFlag
  13 // CHECK: from ._ods_common import _cext as _ods_cext
  14 // CHECK: from ..ir import register_attribute_builder
  15 // CHECK: _ods_ir = _ods_cext.ir
  16
  17 def One : I32EnumAttrCase<"CaseOne", 1, "one">;
  18 def Two : I32EnumAttrCase<"CaseTwo", 2, "two">;
  19 def NegOne : I32EnumAttrCase<"CaseNegOne", -1, "negone">;
  20
  21 def MyEnum : I32EnumAttr<"MyEnum", "An example 32-bit enum", [One, Two, NegOne]>;
  22 // CHECK-LABEL: class MyEnum(IntEnum):
  23 // CHECK:     """An example 32-bit enum"""
  24
  25 // CHECK:     CaseOne = 1
  26 // CHECK:     CaseTwo = 2
  27 // CHECK:     CaseNegOne = auto()
  28
  29 // CHECK:     def __str__(self):
  30 // CHECK:         if self is MyEnum.CaseOne:
  31 // CHECK:             return "one"
  32 // CHECK:         if self is MyEnum.CaseTwo:
  33 // CHECK:             return "two"
  34 // CHECK:         if self is MyEnum.CaseNegOne:
  35 // CHECK:             return "negone"
  36 // CHECK:         raise ValueError("Unknown MyEnum enum entry.")
  37
  38 // CHECK: @register_attribute_builder("MyEnum")
  39 // CHECK: def _myenum(x, context):
  40 // CHECK:     return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))
  41
  42 def TestMyEnum_Attr : EnumAttr<Test_Dialect, MyEnum, "enum">;
  43
  44 def One64 : I64EnumAttrCase<"CaseOne64", 1, "one">;
  45 def Two64 : I64EnumAttrCase<"CaseTwo64", 2, "two">;
  46
  47 def MyEnum64 : I64EnumAttr<"MyEnum64", "An example 64-bit enum", [One64, Two64]>;
  48 // CHECK-LABEL: class MyEnum64(IntEnum):
  49 // CHECK:     """An example 64-bit enum"""
  50
  51 // CHECK:     CaseOne64 = 1
  52 // CHECK:     CaseTwo64 = 2
  53
  54 // CHECK:     def __str__(self):
  55 // CHECK:         if self is MyEnum64.CaseOne64:
  56 // CHECK:             return "one"
  57 // CHECK:         if self is MyEnum64.CaseTwo64:
  58 // CHECK:             return "two"
  59 // CHECK:         raise ValueError("Unknown MyEnum64 enum entry.")
  60
  61 // CHECK: @register_attribute_builder("MyEnum64")
  62 // CHECK: def _myenum64(x, context):
  63 // CHECK:     return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(64, context=context), int(x))
  64
  65 def TestBitEnum
  66     : I32BitEnumAttr<"TestBitEnum", "", [
  67         I32BitEnumAttrCaseBit<"User", 0, "user">,
  68         I32BitEnumAttrCaseBit<"Group", 1, "group">,
  69         I32BitEnumAttrCaseBit<"Other", 2, "other">,
  70       ]> {
  71   let genSpecializedAttr = 0;
  72   let separator = " | ";
  73 }
  74
  75 def TestBitEnum_Attr : EnumAttr<Test_Dialect, TestBitEnum, "testbitenum">;
  76
  77 // CHECK-LABEL: class TestBitEnum(IntFlag):
  78
  79 // CHECK:     User = 1
  80 // CHECK:     Group = 2
  81 // CHECK:     Other = 4
  82
  83 // CHECK:     def __iter__(self):
  84 // CHECK:         return iter([case for case in type(self) if (self & case) is case])
  85 // CHECK:     def __len__(self):
  86 // CHECK:         return bin(self).count("1")
  87
  88 // CHECK:     def __str__(self):
  89 // CHECK:         if len(self) > 1:
  90 // CHECK:             return " | ".join(map(str, self))
  91 // CHECK:         if self is TestBitEnum.User:
  92 // CHECK:             return "user"
  93 // CHECK:         if self is TestBitEnum.Group:
  94 // CHECK:             return "group"
  95 // CHECK:         if self is TestBitEnum.Other:
  96 // CHECK:             return "other"
  97 // CHECK:         raise ValueError("Unknown TestBitEnum enum entry.")
  98
  99 // CHECK: @register_attribute_builder("TestBitEnum")
 100 // CHECK: def _testbitenum(x, context):
 101 // CHECK:     return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))
 102
 103 // CHECK: @register_attribute_builder("TestBitEnum_Attr")
 104 // CHECK: def _testbitenum_attr(x, context):
 105 // CHECK:     return _ods_ir.Attribute.parse(f'#TestDialect<testbitenum {str(x)}>', context=context)
 106
 107 // CHECK: @register_attribute_builder("TestMyEnum_Attr")
 108 // CHECK: def _testmyenum_attr(x, context):
 109 // CHECK:     return _ods_ir.Attribute.parse(f'#TestDialect<enum {str(x)}>', context=context)