[Mlir-commits] [mlir] [MLIR] [Python] The generated op definitions now use typed parameters (PR #188635)

[Maksim Levental]

================
@@ -1108,54 +1143,115 @@ static SmallVector<std::string> emitDefaultOpBuilder(const Operator &op,
   populateBuilderLinesSuccessors(op, successorArgNames, builderLines);
   populateBuilderRegions(op, builderArgs, builderLines);
 
-  // Layout of builderArgs vector elements:
-  // [ result_args  operand_attr_args successor_args regions ]
+  // Compute type annotations for each builder arg.
+  SmallVector<std::string> argTypes(builderArgs.size());
+
+  // Result args: user passes Type objects.
+  for (size_t i = 0; i < numResultArgs; ++i) {
+    const NamedTypeConstraint &result = op.getResult(i);
+    if (result.isVariadic())
+      argTypes[i] = "_Sequence[_ods_ir.Type]";
+    else if (result.isOptional())
+      argTypes[i] = "_Optional[_ods_ir.Type]";
+    else
+      argTypes[i] = "_ods_ir.Type";
+  }
+
+  // Operand and attribute args.
+  for (size_t i = 0; i < numOperandAttrArgs; ++i) {
+    size_t idx = numResultArgs + i;
+    Argument arg = op.getArg(i);
+    if (auto *nattr = llvm::dyn_cast_if_present<NamedAttribute *>(arg)) {
+      if (nattr->attr.getStorageType().trim() == "::mlir::UnitAttr") {
+        argTypes[idx] = "bool";
+      } else {
+        std::string attrType = "_ods_ir." + getPythonAttrName(nattr->attr);
+        StringRef rawType = getPythonAttrRawType(nattr->attr);
+        argTypes[idx] =
+            llvm::formatv("_Union[{0}, {1}]",
+                          rawType.empty() ? "_Any" : rawType, attrType)
+                .str();
+      }
+    } else if (auto *ntype =
+                   llvm::dyn_cast_if_present<NamedTypeConstraint *>(arg)) {
+      if (ntype->isVariadic()) {
+        std::string type = "_ods_ir.Value";
+        if (StringRef pythonType =
+                getPythonType(ntype->constraint.getCppType());
+            !pythonType.empty())
+          type = llvm::formatv("{0}[{1}]", type, pythonType);
+        argTypes[idx] = llvm::formatv("_Sequence[{0}]", type);
+      } else {
+        std::string type = "_ods_ir.Value";
+        if (StringRef pythonType =
+                getPythonType(ntype->constraint.getCppType());
+            !pythonType.empty())
+          type = llvm::formatv("{0}[{1}]", type, pythonType);
+        argTypes[idx] = type;
+      }
+    }
+    // NamedProperty args are skipped (no type hint).
+  }
+
+  // Successor args.
+  for (size_t i = 0; i < numSuccessorArgs; ++i) {
+    size_t idx = numResultArgs + numOperandAttrArgs + i;
+    const NamedSuccessor &successor = op.getSuccessor(i);
+    argTypes[idx] =
+        successor.isVariadic() ? "_Sequence[_ods_ir.Block]" : "_ods_ir.Block";
+  }
+
+  // Region args (variadic region count).
+  for (size_t i = numResultArgs + numOperandAttrArgs + numSuccessorArgs;
+       i < builderArgs.size(); ++i) {
+    argTypes[i] = "int";
+  }
 
-  // Determine whether the argument corresponding to a given index into the
-  // builderArgs vector is a python keyword argument or not.
-  auto isKeywordArgFn = [&](size_t builderArgIndex) -> bool {
-    // All result, successor, and region arguments are positional arguments.
-    if ((builderArgIndex < numResultArgs) ||
-        (builderArgIndex >= (numResultArgs + numOperandAttrArgs)))
+  // Determine whether a builder arg is a keyword argument.
+  auto isKeywordArg = [&](size_t i) -> bool {
+    // Only operand/attr args can be keyword; results, successors, and regions
+    // are always positional.
+    if (i < numResultArgs || i >= numResultArgs + numOperandAttrArgs)
       return false;
-    // Keyword arguments:
-    // - optional named attributes (including unit attributes)
-    // - default-valued named attributes
-    // - optional operands
-    Argument a = op.getArg(builderArgIndex - numResultArgs);
+    Argument a = op.getArg(i - numResultArgs);
     if (auto *nattr = llvm::dyn_cast_if_present<NamedAttribute *>(a))
-      return (nattr->attr.isOptional() || nattr->attr.hasDefaultValue());
+      return nattr->attr.isOptional() || nattr->attr.hasDefaultValue();
     if (auto *ntype = llvm::dyn_cast_if_present<NamedTypeConstraint *>(a))
       return ntype->isOptional();
     return false;
   };
 
-  // StringRefs in functionArgs refer to strings allocated by builderArgs.
-  SmallVector<StringRef> functionArgs;
+  // Format a single function argument with optional type hint and default.
+  auto formatArg = [](StringRef name, StringRef typeHint,
+                      bool isKeyword) -> std::string {
+    std::string result = name.str();
+    if (isKeyword && !typeHint.empty()) {
+      result += ": _Optional[" + typeHint.str() + "] = None";
+    } else if (isKeyword) {
+      result += "=None";
+    } else if (!typeHint.empty()) {
+      result += ": " + typeHint.str();
+    }
+    return result;
+  };
 
-  // Add positional arguments.
-  for (size_t i = 0, cnt = builderArgs.size(); i < cnt; ++i) {
-    if (!isKeywordArgFn(i))
-      functionArgs.push_back(builderArgs[i]);
-  }
+  // Build the function argument list: positional args, *, keyword args.
+  SmallVector<std::string> functionArgs;
+  for (size_t i = 0, cnt = builderArgs.size(); i < cnt; ++i)
+    if (!isKeywordArg(i))
+      functionArgs.push_back(formatArg(builderArgs[i], argTypes[i], false));
 
-  // Add a bare '*' to indicate that all following arguments must be keyword
-  // arguments.
   functionArgs.push_back("*");
 
-  // Add a default 'None' value to each keyword arg string, and then add to the
-  // function args list.
----------------
makslevental wrote:

same nit as above

https://github.com/llvm/llvm-project/pull/188635