Address review comments

Author: jpienaar

mlir/python/mlir/_mlir_libs/__init__.py

@@ -156,16 +156,23 @@ def __init__(self, load_on_create_dialects=None, *args, **kwargs):
             if not disable_multithreading:
                 self.enable_multithreading(True)
             if load_on_create_dialects is not None:
-                logger.debug("Loading all dialects from load_on_create_dialects arg %r", _load_on_create_dialects)
+                logger.debug(
+                    "Loading all dialects from load_on_create_dialects arg %r",
+                    load_on_create_dialects,
+                )
                 for dialect in load_on_create_dialects:
-                    # Load dialect.
+                    # This triggers loading the dialect into the context.
                     _ = self.dialects[dialect]
             else:
                 if disable_load_all_available_dialects:
-                    if _load_on_create_dialects:
-                        logger.debug("Loading all dialects from global load_on_create_dialects %r", _load_on_create_dialects)
-                        for dialect in _load_on_create_dialects:
-                            # Load dialect.
+                    dialects = get_load_on_create_dialects()
+                    if dialects:
+                        logger.debug(
+                            "Loading all dialects from global load_on_create_dialects %r",
+                            dialects,
+                        )
+                        for dialect in dialects:
+                            # This triggers loading the dialect into the context.
                             _ = self.dialects[dialect]
                 else:
                     logger.debug("Loading all available dialects")

mlir/python/mlir/ir.py

@@ -5,7 +5,11 @@
 from ._mlir_libs._mlir.ir import *
 from ._mlir_libs._mlir.ir import _GlobalDebug
 from ._mlir_libs._mlir import register_type_caster, register_value_caster
-from ._mlir_libs import get_dialect_registry, append_load_on_create_dialect, get_load_on_create_dialects
+from ._mlir_libs import (
+    get_dialect_registry,
+    append_load_on_create_dialect,
+    get_load_on_create_dialects,
+)
 
 
 # Convenience decorator for registering user-friendly Attribute builders.

mlir/test/python/ir/dialects.py

@@ -129,7 +129,7 @@ def testDialectLoadOnCreate():
     with Context(load_on_create_dialects=[]) as ctx:
         ctx.emit_error_diagnostics = True
         ctx.allow_unregistered_dialects = True
-        
+
         def callback(d):
             # CHECK: DIAGNOSTIC
             # CHECK-SAME: op created with unregistered dialect
@@ -139,13 +139,21 @@ def callback(d):
         handler = ctx.attach_diagnostic_handler(callback)
         loc = Location.unknown(ctx)
         try:
-          op = Operation.create("arith.addi", loc=loc)
-          ctx.allow_unregistered_dialects = False
-          op.verify()
+            op = Operation.create("arith.addi", loc=loc)
+            ctx.allow_unregistered_dialects = False
+            op.verify()
         except MLIRError as e:
-          pass
-  
-    with Context(load_on_create_dialects=["func"]) as ctx:
-      loc = Location.unknown(ctx)
-      fn = Operation.create("func.func", loc=loc)
+            pass
 
+    with Context(load_on_create_dialects=["func"]) as ctx:
+        loc = Location.unknown(ctx)
+        fn = Operation.create("func.func", loc=loc)
+
+    # TODO: This may require an update if a site wide policy is set.
+    # CHECK: Load on create: []
+    print(f"Load on create: {get_load_on_create_dialects()}")
+    append_load_on_create_dialect("func")
+    # CHECK: Load on create:
+    # CHECK-SAME: func
+    print(f"Load on create: {get_load_on_create_dialects()}")
+    print(get_load_on_create_dialects())