[Mlir-commits] [mlir] 94d96c2 - [MLIR] [Python] Pybind adaptors: coerce None to default MlirLocation

[John Demme]

Author: John Demme
Date: 2022-04-05T17:10:20-07:00
New Revision: 94d96c2a6201e0927e835b5c0e5b4cb8f95eeadf

URL: https://github.com/llvm/llvm-project/commit/94d96c2a6201e0927e835b5c0e5b4cb8f95eeadf
DIFF: https://github.com/llvm/llvm-project/commit/94d96c2a6201e0927e835b5c0e5b4cb8f95eeadf.diff

LOG: [MLIR] [Python] Pybind adaptors: coerce None to default MlirLocation

Add default source location coercion to enable location elision in
Python code.

Added: 
    

Modified: 
    mlir/include/mlir/Bindings/Python/PybindAdaptors.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Bindings/Python/PybindAdaptors.h b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
index 661ed48f9c1dd..582855a33afb4 100644
--- a/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
@@ -125,11 +125,16 @@ struct type_caster<MlirContext> {
 };
 
 /// Casts object <-> MlirLocation.
-// TODO: Coerce None to default MlirLocation.
 template <>
 struct type_caster<MlirLocation> {
   PYBIND11_TYPE_CASTER(MlirLocation, _("MlirLocation"));
   bool load(handle src, bool) {
+    if (src.is_none()) {
+      // Gets the current thread-bound context.
+      src = py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+                .attr("Location")
+                .attr("current");
+    }
     py::object capsule = mlirApiObjectToCapsule(src);
     value = mlirPythonCapsuleToLocation(capsule.ptr());
     return !mlirLocationIsNull(value);