NanobindAdaptors.h

Go to the documentation of this file.

//===- NanobindAdaptors.h - Interop with MLIR APIs via nanobind -----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
// This file contains adaptors for clients of the core MLIR Python APIs to
// interop via MLIR CAPI types, using nanobind. The facilities here do not
// depend on implementation details of the MLIR Python API and do not introduce
// C++-level dependencies with it (requiring only Python and CAPI-level
// dependencies).
//
// It is encouraged to be used both in-tree and out-of-tree. For in-tree use
// cases, it should be used for dialect implementations (versus relying on
// Pybind-based internals of the core libraries).
//===----------------------------------------------------------------------===//
 
#ifndef MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H
#define MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H
 
#include <cstdint>
#include <memory>
#include <optional>
 
#include "mlir-c/Diagnostics.h"
#include "mlir-c/IR.h"
// clang-format off
#include "mlir/Bindings/Python/Nanobind.h"
#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
// clang-format on
#include "mlir/Bindings/Python/NanobindUtils.h"
 
namespace mlir {
namespace python {
namespace {
 
nanobind::module_ &irModule() {
  static SafeInit<nanobind::module_> init([]() {
    return std::make_unique<nanobind::module_>(
        nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir")));
  });
  return init.get();
}
 
} // namespace
} // namespace python
} // namespace mlir
 
// Raw CAPI type casters need to be declared before use, so always include them
// first.
namespace nanobind {
namespace detail {
 
/// Helper to convert a presumed MLIR API object to a capsule, accepting either
/// an explicit Capsule (which can happen when two C APIs are communicating
/// directly via Python) or indirectly by querying the MLIR_PYTHON_CAPI_PTR_ATTR
/// attribute (through which supported MLIR Python API objects export their
/// contained API pointer as a capsule). Throws a type error if the object is
/// neither. This is intended to be used from type casters, which are invoked
/// with a raw handle (unowned). The returned object's lifetime may not extend
/// beyond the apiObject handle without explicitly having its refcount increased
/// (i.e. on return).
static std::optional<nanobind::object>
mlirApiObjectToCapsule(nanobind::handle apiObject) {
  if (PyCapsule_CheckExact(apiObject.ptr()))
    return nanobind::borrow<nanobind::object>(apiObject);
  nanobind::object api =
      nanobind::getattr(apiObject, MLIR_PYTHON_CAPI_PTR_ATTR, nanobind::none());
  if (api.is_none())
    return {};
  return api;
}
 
/// Clears the Python error indicator if the given condition `val` is false and
/// returns the condition. This is needed in `from_python` of the type casters
/// below, where a failed conversion from a Python capsule sets the Python error
/// indicator but the caller of the caster expects or may even need a clean
/// error indicator.
inline bool pyErrClearIfFalse(bool val) {
  if (!val)
    PyErr_Clear();
  return val;
}
 
// Note: Currently all of the following support cast from nanobind::object to
// the Mlir* C-API type, but only a few light-weight, context-bound ones
// implicitly cast the other way because the use case has not yet emerged and
// ownership is unclear.
 
/// Casts object <-> MlirAffineMap.
template <>
struct type_caster<MlirAffineMap> {
  NB_TYPE_CASTER(MlirAffineMap,
                 const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.AffineMap")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToAffineMap(capsule->ptr());
      return pyErrClearIfFalse(!mlirAffineMapIsNull(value));
    }
    return false;
  }
  static handle from_cpp(MlirAffineMap v, rv_policy,
                         cleanup_list *cleanup) noexcept {
    nanobind::object capsule =
        nanobind::steal<nanobind::object>(mlirPythonAffineMapToCapsule(v));
    return mlir::python::irModule()
        .attr("AffineMap")
        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
        .release();
  }
};
 
/// Casts object <-> MlirAttribute.
template <>
struct type_caster<MlirAttribute> {
  NB_TYPE_CASTER(MlirAttribute,
                 const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Attribute")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToAttribute(capsule->ptr());
      return pyErrClearIfFalse(!mlirAttributeIsNull(value));
    }
    return false;
  }
  static handle from_cpp(MlirAttribute v, rv_policy,
                         cleanup_list *cleanup) noexcept {
    nanobind::object capsule =
        nanobind::steal<nanobind::object>(mlirPythonAttributeToCapsule(v));
    return mlir::python::irModule()
        .attr("Attribute")
        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
        .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
        .release();
  }
};
 
/// Casts object -> MlirBlock.
template <>
struct type_caster<MlirBlock> {
  NB_TYPE_CASTER(MlirBlock, const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Block")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToBlock(capsule->ptr());
      return pyErrClearIfFalse(!mlirBlockIsNull(value));
    }
    return false;
  }
};
 
/// Casts object -> MlirContext.
template <>
struct type_caster<MlirContext> {
  NB_TYPE_CASTER(MlirContext,
                 const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Context")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (src.is_none()) {
      // Gets the current thread-bound context.
      src = mlir::python::irModule().attr("Context").attr("current");
    }
    // If there is no context, including thread-bound, emit a warning (since
    // this function is not allowed to throw) and fail to cast.
    if (src.is_none()) {
      PyErr_WarnEx(
          PyExc_RuntimeWarning,
          "Passing None as MLIR Context is only allowed inside "
          "the " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " context manager.",
          /*stacklevel=*/1);
      return false;
    }
    if (std::optional<nanobind::object> capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToContext(capsule->ptr());
      return pyErrClearIfFalse(!mlirContextIsNull(value));
    }
    return false;
  }
};
 
/// Casts object <-> MlirDialectRegistry.
template <>
struct type_caster<MlirDialectRegistry> {
  NB_TYPE_CASTER(MlirDialectRegistry,
                 const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.DialectRegistry")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToDialectRegistry(capsule->ptr());
      return pyErrClearIfFalse(!mlirDialectRegistryIsNull(value));
    }
    return false;
  }
  static handle from_cpp(MlirDialectRegistry v, rv_policy,
                         cleanup_list *cleanup) noexcept {
    nanobind::object capsule = nanobind::steal<nanobind::object>(
        mlirPythonDialectRegistryToCapsule(v));
    return mlir::python::irModule()
        .attr("DialectRegistry")
        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
        .release();
  }
};
 
/// Casts object <-> MlirLocation.
template <>
struct type_caster<MlirLocation> {
  NB_TYPE_CASTER(MlirLocation,
                 const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Location")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (src.is_none()) {
      // Gets the current thread-bound context.
      src = mlir::python::irModule().attr("Location").attr("current");
    }
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToLocation(capsule->ptr());
      return pyErrClearIfFalse(!mlirLocationIsNull(value));
    }
    return false;
  }
  static handle from_cpp(MlirLocation v, rv_policy,
                         cleanup_list *cleanup) noexcept {
    nanobind::object capsule =
        nanobind::steal<nanobind::object>(mlirPythonLocationToCapsule(v));
    return mlir::python::irModule()
        .attr("Location")
        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
        .release();
  }
};
 
/// Casts object <-> MlirModule.
template <>
struct type_caster<MlirModule> {
  NB_TYPE_CASTER(MlirModule, const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Module")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToModule(capsule->ptr());
      return pyErrClearIfFalse(!mlirModuleIsNull(value));
    }
    return false;
  }
  static handle from_cpp(MlirModule v, rv_policy,
                         cleanup_list *cleanup) noexcept {
    nanobind::object capsule =
        nanobind::steal<nanobind::object>(mlirPythonModuleToCapsule(v));
    return mlir::python::irModule()
        .attr("Module")
        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
        .release();
  };
};
 
/// Casts object <-> MlirFrozenRewritePatternSet.
template <>
struct type_caster<MlirFrozenRewritePatternSet> {
  NB_TYPE_CASTER(
      MlirFrozenRewritePatternSet,
      const_name(MAKE_MLIR_PYTHON_QUALNAME("rewrite.FrozenRewritePatternSet")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToFrozenRewritePatternSet(capsule->ptr());
      return pyErrClearIfFalse(value.ptr != nullptr);
    }
    return false;
  }
  static handle from_cpp(MlirFrozenRewritePatternSet v, rv_policy,
                         handle) noexcept {
    nanobind::object capsule = nanobind::steal<nanobind::object>(
        mlirPythonFrozenRewritePatternSetToCapsule(v));
    return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("rewrite"))
        .attr("FrozenRewritePatternSet")
        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
        .release();
  };
};
 
/// Casts object <-> MlirOperation.
template <>
struct type_caster<MlirOperation> {
  NB_TYPE_CASTER(MlirOperation,
                 const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Operation")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToOperation(capsule->ptr());
      return pyErrClearIfFalse(!mlirOperationIsNull(value));
    }
    return false;
  }
  static handle from_cpp(MlirOperation v, rv_policy,
                         cleanup_list *cleanup) noexcept {
    if (v.ptr == nullptr)
      return nanobind::none();
    nanobind::object capsule =
        nanobind::steal<nanobind::object>(mlirPythonOperationToCapsule(v));
    return mlir::python::irModule()
        .attr("Operation")
        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
        .release();
  };
};
 
/// Casts object <-> MlirValue.
template <>
struct type_caster<MlirValue> {
  NB_TYPE_CASTER(MlirValue, const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Value")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToValue(capsule->ptr());
      return pyErrClearIfFalse(!mlirValueIsNull(value));
    }
    return false;
  }
  static handle from_cpp(MlirValue v, rv_policy,
                         cleanup_list *cleanup) noexcept {
    if (v.ptr == nullptr)
      return nanobind::none();
    nanobind::object capsule =
        nanobind::steal<nanobind::object>(mlirPythonValueToCapsule(v));
    return mlir::python::irModule()
        .attr("Value")
        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
        .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
        .release();
  };
};
 
/// Casts object -> MlirPassManager.
template <>
struct type_caster<MlirPassManager> {
  NB_TYPE_CASTER(MlirPassManager, const_name(MAKE_MLIR_PYTHON_QUALNAME(
                                      "passmanager.PassManager")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToPassManager(capsule->ptr());
      return pyErrClearIfFalse(!mlirPassManagerIsNull(value));
    }
    return false;
  }
};
 
/// Casts object <-> MlirTypeID.
template <>
struct type_caster<MlirTypeID> {
  NB_TYPE_CASTER(MlirTypeID, const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToTypeID(capsule->ptr());
      return pyErrClearIfFalse(!mlirTypeIDIsNull(value));
    }
    return false;
  }
  static handle from_cpp(MlirTypeID v, rv_policy,
                         cleanup_list *cleanup) noexcept {
    if (v.ptr == nullptr)
      return nanobind::none();
    nanobind::object capsule =
        nanobind::steal<nanobind::object>(mlirPythonTypeIDToCapsule(v));
    return mlir::python::irModule()
        .attr("TypeID")
        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
        .release();
  };
};
 
/// Casts object <-> MlirType.
template <>
struct type_caster<MlirType> {
  NB_TYPE_CASTER(MlirType, const_name(MAKE_MLIR_PYTHON_QUALNAME("ir.Type")))
  bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
    if (auto capsule = mlirApiObjectToCapsule(src)) {
      value = mlirPythonCapsuleToType(capsule->ptr());
      return pyErrClearIfFalse(!mlirTypeIsNull(value));
    }
    return false;
  }
  static handle from_cpp(MlirType t, rv_policy,
                         cleanup_list *cleanup) noexcept {
    nanobind::object capsule =
        nanobind::steal<nanobind::object>(mlirPythonTypeToCapsule(t));
    return mlir::python::irModule()
        .attr("Type")
        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
        .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
        .release();
  }
};
 
/// Casts MlirStringRef -> object.
template <>
struct type_caster<MlirStringRef> {
  NB_TYPE_CASTER(MlirStringRef, const_name("str"))
  static handle from_cpp(MlirStringRef s, rv_policy,
                         cleanup_list *cleanup) noexcept {
    return nanobind::str(s.data, s.length).release();
  }
};
 
} // namespace detail
} // namespace nanobind
 
namespace mlir {
namespace python {
namespace nanobind_adaptors {
 
/// Provides a facility like nanobind::class_ for defining a new class in a
/// scope, but this allows extension of an arbitrary Python class, defining
/// methods on it is a similar way. Classes defined in this way are very similar
/// to if defined in Python in the usual way but use nanobind machinery to
/// do it. These are not "real" nanobind classes but pure Python classes
/// with no relation to a concrete C++ class.
///
/// Derived from a discussion upstream:
///   https://github.com/pybind/pybind11/issues/1193
///   (plus a fair amount of extra curricular poking)
///   TODO: If this proves useful, see about including it in nanobind.
class pure_subclass {
public:
  pure_subclass(nanobind::handle scope, const char *derivedClassName,
                const nanobind::object &superClass) {
    nanobind::object pyType =
        nanobind::borrow<nanobind::object>((PyObject *)&PyType_Type);
    nanobind::object metaclass = pyType(superClass);
    nanobind::dict attributes;
 
    thisClass = metaclass(derivedClassName, nanobind::make_tuple(superClass),
                          attributes);
    scope.attr(derivedClassName) = thisClass;
    thisClass.attr("__module__") = scope.attr("__name__");
  }
 
  template <typename Func, typename... Extra>
  pure_subclass &def(const char *name, Func &&f, const Extra &...extra) {
    nanobind::object cf = nanobind::cpp_function(
        std::forward<Func>(f), nanobind::name(name), nanobind::is_method(),
        nanobind::scope(thisClass), extra...);
    thisClass.attr(name) = cf;
    return *this;
  }
 
  template <typename Func, typename... Extra>
  pure_subclass &def_property_readonly(const char *name, Func &&f,
                                       const Extra &...extra) {
    nanobind::object cf = nanobind::cpp_function(
        std::forward<Func>(f), nanobind::name(name), nanobind::is_method(),
        nanobind::scope(thisClass), extra...);
    auto builtinProperty =
        nanobind::borrow<nanobind::object>((PyObject *)&PyProperty_Type);
    thisClass.attr(name) = builtinProperty(cf);
    return *this;
  }
 
  template <typename Func, typename... Extra>
  pure_subclass &def_staticmethod(const char *name, Func &&f,
                                  const Extra &...extra) {
    static_assert(!std::is_member_function_pointer<Func>::value,
                  "def_staticmethod(...) called with a non-static member "
                  "function pointer");
    nanobind::object cf = nanobind::cpp_function(
        std::forward<Func>(f),
        nanobind::name(name), // nanobind::scope(thisClass),
        extra...);
    thisClass.attr(name) = cf;
    return *this;
  }
 
  template <typename Func, typename... Extra>
  pure_subclass &def_classmethod(const char *name, Func &&f,
                                 const Extra &...extra) {
    static_assert(!std::is_member_function_pointer<Func>::value,
                  "def_classmethod(...) called with a non-static member "
                  "function pointer");
    nanobind::object cf = nanobind::cpp_function(
        std::forward<Func>(f),
        nanobind::name(name), // nanobind::scope(thisClass),
        extra...);
    static SafeInit<nanobind::object> classmethodFn([]() {
      return std::make_unique<nanobind::object>(
          nanobind::module_::import_("builtins").attr("classmethod"));
    });
    thisClass.attr(name) = classmethodFn.get()(cf);
    return *this;
  }
 
  nanobind::object get_class() const { return thisClass; }
 
protected:
  nanobind::object superClass;
  nanobind::object thisClass;
};
 
/// Creates a custom subclass of mlir.ir.Attribute, implementing a casting
/// constructor and type checking methods.
class mlir_attribute_subclass : public pure_subclass {
public:
  using IsAFunctionTy = bool (*)(MlirAttribute);
  using GetTypeIDFunctionTy = MlirTypeID (*)();
 
  /// Subclasses by looking up the super-class dynamically.
  mlir_attribute_subclass(nanobind::handle scope, const char *attrClassName,
                          IsAFunctionTy isaFunction,
                          GetTypeIDFunctionTy getTypeIDFunction = nullptr)
      : mlir_attribute_subclass(scope, attrClassName, isaFunction,
                                irModule().attr("Attribute"),
                                getTypeIDFunction) {}
 
  /// Subclasses with a provided mlir.ir.Attribute super-class. This must
  /// be used if the subclass is being defined in the same extension module
  /// as the mlir.ir class (otherwise, it will trigger a recursive
  /// initialization).
  mlir_attribute_subclass(nanobind::handle scope, const char *typeClassName,
                          IsAFunctionTy isaFunction,
                          const nanobind::object &superCls,
                          GetTypeIDFunctionTy getTypeIDFunction = nullptr)
      : pure_subclass(scope, typeClassName, superCls) {
    // Casting constructor. Note that it is hard, if not impossible, to properly
    // call chain to parent `__init__` in nanobind due to its special handling
    // for init functions that don't have a fully constructed self-reference,
    // which makes it impossible to forward it to `__init__` of a superclass.
    // Instead, provide a custom `__new__` and call that of a superclass, which
    // eventually calls `__init__` of the superclass. Since attribute subclasses
    // have no additional members, we can just return the instance thus created
    // without amending it.
    std::string captureTypeName(
        typeClassName); // As string in case if typeClassName is not static.
    nanobind::object newCf = nanobind::cpp_function(
        [superCls, isaFunction, captureTypeName](
            nanobind::object cls, nanobind::object otherAttribute) {
          MlirAttribute rawAttribute;
          if (!nanobind::try_cast<MlirAttribute>(otherAttribute,
                                                 rawAttribute) ||
              !isaFunction(rawAttribute)) {
            auto origRepr =
                nanobind::cast<std::string>(nanobind::repr(otherAttribute));
            throw std::invalid_argument(nanobind::detail::join(
                "Cannot cast attribute to ", captureTypeName, " (from ",
                origRepr, ")"));
          }
          nanobind::object self = superCls.attr("__new__")(cls, otherAttribute);
          return self;
        },
        nanobind::name("__new__"), nanobind::arg("cls"),
        nanobind::arg("cast_from_attr"));
    thisClass.attr("__new__") = newCf;
 
    // 'isinstance' method.
    static const char kIsinstanceSig[] =
        "def isinstance(other_attribute: " MAKE_MLIR_PYTHON_QUALNAME(
            "ir") ".Attribute) -> bool";
    def_staticmethod(
        "isinstance",
        [isaFunction](MlirAttribute other) { return isaFunction(other); },
        nanobind::arg("other_attribute"), nanobind::sig(kIsinstanceSig));
    def("__repr__", [superCls, captureTypeName](nanobind::object self) {
      return nanobind::cast<std::string>(
          nanobind::repr(superCls(self))
              .attr("replace")(superCls.attr("__name__"), captureTypeName));
    });
    if (getTypeIDFunction) {
      def_staticmethod(
          "get_static_typeid",
          [getTypeIDFunction]() { return getTypeIDFunction(); },
          // clang-format off
          nanobind::sig("def get_static_typeid() -> " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID"))
          // clang-format on
      );
      nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
          .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
              getTypeIDFunction())(nanobind::cpp_function(
              [thisClass = thisClass](const nanobind::object &mlirAttribute) {
                return thisClass(mlirAttribute);
              }));
    }
  }
};
 
/// Creates a custom subclass of mlir.ir.Type, implementing a casting
/// constructor and type checking methods.
class mlir_type_subclass : public pure_subclass {
public:
  using IsAFunctionTy = bool (*)(MlirType);
  using GetTypeIDFunctionTy = MlirTypeID (*)();
 
  /// Subclasses by looking up the super-class dynamically.
  mlir_type_subclass(nanobind::handle scope, const char *typeClassName,
                     IsAFunctionTy isaFunction,
                     GetTypeIDFunctionTy getTypeIDFunction = nullptr)
      : mlir_type_subclass(scope, typeClassName, isaFunction,
                           irModule().attr("Type"), getTypeIDFunction) {}
 
  /// Subclasses with a provided mlir.ir.Type super-class. This must
  /// be used if the subclass is being defined in the same extension module
  /// as the mlir.ir class (otherwise, it will trigger a recursive
  /// initialization).
  mlir_type_subclass(nanobind::handle scope, const char *typeClassName,
                     IsAFunctionTy isaFunction,
                     const nanobind::object &superCls,
                     GetTypeIDFunctionTy getTypeIDFunction = nullptr)
      : pure_subclass(scope, typeClassName, superCls) {
    // Casting constructor. Note that it is hard, if not impossible, to properly
    // call chain to parent `__init__` in nanobind due to its special handling
    // for init functions that don't have a fully constructed self-reference,
    // which makes it impossible to forward it to `__init__` of a superclass.
    // Instead, provide a custom `__new__` and call that of a superclass, which
    // eventually calls `__init__` of the superclass. Since attribute subclasses
    // have no additional members, we can just return the instance thus created
    // without amending it.
    std::string captureTypeName(
        typeClassName); // As string in case if typeClassName is not static.
    nanobind::object newCf = nanobind::cpp_function(
        [superCls, isaFunction, captureTypeName](nanobind::object cls,
                                                 nanobind::object otherType) {
          MlirType rawType;
          if (!nanobind::try_cast<MlirType>(otherType, rawType) ||
              !isaFunction(rawType)) {
            auto origRepr =
                nanobind::cast<std::string>(nanobind::repr(otherType));
            throw std::invalid_argument(
                nanobind::detail::join("Cannot cast type to ", captureTypeName,
                                       " (from ", origRepr, ")"));
          }
          nanobind::object self = superCls.attr("__new__")(cls, otherType);
          return self;
        },
        nanobind::name("__new__"), nanobind::arg("cls"),
        nanobind::arg("cast_from_type"));
    thisClass.attr("__new__") = newCf;
 
    // 'isinstance' method.
    static const char kIsinstanceSig[] =
        // clang-format off
        "def isinstance(other_type: " MAKE_MLIR_PYTHON_QUALNAME("ir.Type") ") -> bool";
    // clang-format on
    def_staticmethod(
        "isinstance",
        [isaFunction](MlirType other) { return isaFunction(other); },
        nanobind::arg("other_type"), nanobind::sig(kIsinstanceSig));
    def("__repr__", [superCls, captureTypeName](nanobind::object self) {
      return nanobind::cast<std::string>(
          nanobind::repr(superCls(self))
              .attr("replace")(superCls.attr("__name__"), captureTypeName));
    });
    if (getTypeIDFunction) {
      // 'get_static_typeid' method.
      // This is modeled as a static method instead of a static property because
      // `def_property_readonly_static` is not available in `pure_subclass` and
      // we do not want to introduce the complexity that pybind uses to
      // implement it.
      def_staticmethod(
          "get_static_typeid",
          [getTypeIDFunction]() { return getTypeIDFunction(); },
          // clang-format off
          nanobind::sig("def get_static_typeid() -> " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID"))
          // clang-format on
      );
      nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
          .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
              getTypeIDFunction())(nanobind::cpp_function(
              [thisClass = thisClass](const nanobind::object &mlirType) {
                return thisClass(mlirType);
              }));
    }
  }
};
 
/// Creates a custom subclass of mlir.ir.Value, implementing a casting
/// constructor and type checking methods.
class mlir_value_subclass : public pure_subclass {
public:
  using IsAFunctionTy = bool (*)(MlirValue);
 
  /// Subclasses by looking up the super-class dynamically.
  mlir_value_subclass(nanobind::handle scope, const char *valueClassName,
                      IsAFunctionTy isaFunction)
      : mlir_value_subclass(scope, valueClassName, isaFunction,
                            irModule().attr("Value")) {}
 
  /// Subclasses with a provided mlir.ir.Value super-class. This must
  /// be used if the subclass is being defined in the same extension module
  /// as the mlir.ir class (otherwise, it will trigger a recursive
  /// initialization).
  mlir_value_subclass(nanobind::handle scope, const char *valueClassName,
                      IsAFunctionTy isaFunction,
                      const nanobind::object &superCls)
      : pure_subclass(scope, valueClassName, superCls) {
    // Casting constructor. Note that it is hard, if not impossible, to properly
    // call chain to parent `__init__` in nanobind due to its special handling
    // for init functions that don't have a fully constructed self-reference,
    // which makes it impossible to forward it to `__init__` of a superclass.
    // Instead, provide a custom `__new__` and call that of a superclass, which
    // eventually calls `__init__` of the superclass. Since attribute subclasses
    // have no additional members, we can just return the instance thus created
    // without amending it.
    std::string captureValueName(
        valueClassName); // As string in case if valueClassName is not static.
    nanobind::object newCf = nanobind::cpp_function(
        [superCls, isaFunction, captureValueName](nanobind::object cls,
                                                  nanobind::object otherValue) {
          MlirValue rawValue;
          if (!nanobind::try_cast<MlirValue>(otherValue, rawValue) ||
              !isaFunction(rawValue)) {
            auto origRepr =
                nanobind::cast<std::string>(nanobind::repr(otherValue));
            throw std::invalid_argument(nanobind::detail::join(
                "Cannot cast value to ", captureValueName, " (from ", origRepr,
                ")"));
          }
          nanobind::object self = superCls.attr("__new__")(cls, otherValue);
          return self;
        },
        nanobind::name("__new__"), nanobind::arg("cls"),
        nanobind::arg("cast_from_value"));
    thisClass.attr("__new__") = newCf;
 
    // 'isinstance' method.
    static const char kIsinstanceSig[] =
        // clang-format off
        "def isinstance(other_value: " MAKE_MLIR_PYTHON_QUALNAME("ir.Value") ") -> bool";
    // clang-format on
    def_staticmethod(
        "isinstance",
        [isaFunction](MlirValue other) { return isaFunction(other); },
        nanobind::arg("other_value"), nanobind::sig(kIsinstanceSig));
  }
};
 
} // namespace nanobind_adaptors
 
} // namespace python
} // namespace mlir
 
#endif // MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H