From 931f2a2cc2f86294081b9f7309b9e19c3bdf6bd9 Mon Sep 17 00:00:00 2001
From: bechtt <bechtelt@fusion.gat.com>
Date: Thu, 28 Mar 2024 17:49:35 -0700
Subject: [PATCH] Fixes to get machine mappings working for other experiments
 again

Source all compatible machines, protect against any that aren't,
add missing dependencies, and include missing checks
---
 omas/machine_mappings/d3d.py                |  1 +
 omas/machine_mappings/east.py               |  2 ++
 omas/machine_mappings/mast.py               |  2 ++
 omas/machine_mappings/nstxu.py              |  4 ++++
 omas/omas_machine.py                        |  5 ++++-
 omas/utilities/machine_mapping_decorator.py | 11 +++++++----
 6 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/omas/machine_mappings/d3d.py b/omas/machine_mappings/d3d.py
index 6c07a0173..bf3ce439a 100644
--- a/omas/machine_mappings/d3d.py
+++ b/omas/machine_mappings/d3d.py
@@ -11,6 +11,7 @@
 from omas.utilities.omas_mds import mdsvalue
 from omas.omas_core import ODS
 from omas.omas_structure import add_extra_structures
+from omas.omas_physics import omas_environment
 
 __all__ = []
 __regression_arguments__ = {'__all__': __all__}
diff --git a/omas/machine_mappings/east.py b/omas/machine_mappings/east.py
index aecf676eb..f8eddc42a 100644
--- a/omas/machine_mappings/east.py
+++ b/omas/machine_mappings/east.py
@@ -3,6 +3,8 @@
 from omas import *
 from omas.omas_utils import printd
 from omas.machine_mappings._common import *
+from omas.utilities.machine_mapping_decorator import machine_mapping_function
+from omas.utilities.omas_mds import mdsvalue
 
 __all__ = []
 __regression_arguments__ = {'__all__': __all__}
diff --git a/omas/machine_mappings/mast.py b/omas/machine_mappings/mast.py
index cb97b6335..778a71d17 100644
--- a/omas/machine_mappings/mast.py
+++ b/omas/machine_mappings/mast.py
@@ -4,6 +4,8 @@
 from omas import *
 from omas.omas_utils import printd, printe, unumpy
 from omas.machine_mappings._common import *
+from omas.utilities.machine_mapping_decorator import machine_mapping_function
+from omas.omas_core import ODS
 import glob
 import pyuda
 
diff --git a/omas/machine_mappings/nstxu.py b/omas/machine_mappings/nstxu.py
index 178c6d492..d02353135 100644
--- a/omas/machine_mappings/nstxu.py
+++ b/omas/machine_mappings/nstxu.py
@@ -4,6 +4,10 @@
 from omas import *
 from omas.omas_utils import printd, printe, unumpy
 from omas.machine_mappings._common import *
+from omas.utilities.machine_mapping_decorator import machine_mapping_function
+from omas.utilities.omas_mds import mdsvalue
+from omas.omas_core import ODS
+from omas.omas_physics import omas_environment
 import glob
 
 # NOTES:
diff --git a/omas/omas_machine.py b/omas/omas_machine.py
index 3c662df2f..4502acda0 100644
--- a/omas/omas_machine.py
+++ b/omas/omas_machine.py
@@ -6,7 +6,7 @@
 from .omas_utils import *
 from .omas_core import ODS, dynamic_ODS, omas_environment, omas_info_node, imas_json_dir, omas_rcparams
 from .omas_physics import cocos_signals
-from omas.machine_mappings import d3d
+from omas.machine_mappings import d3d, nstx, nstxu, east
 from omas.machine_mappings.d3d import __regression_arguments__
 from omas.utilities.machine_mapping_decorator import machine_mapping_function
 from omas.utilities.omas_mds import mdsvalue
@@ -215,6 +215,9 @@ def resolve_mapped(ods, machine, pulse,  mappings, location, idm, options_with_d
 
     # PYTHON
     elif 'PYTHON' in mapped:
+        if 'mast' in machine:
+            printe(f"MAST is currently not supported because of UDA", topic='machine')
+            return ods
         call = mapped['PYTHON'].format(**options_with_defaults)
         # python functions tend to set multiple locations at once
         # it is thus very beneficial to cache that
diff --git a/omas/utilities/machine_mapping_decorator.py b/omas/utilities/machine_mapping_decorator.py
index 9ee4ae395..52286bb9b 100644
--- a/omas/utilities/machine_mapping_decorator.py
+++ b/omas/utilities/machine_mapping_decorator.py
@@ -52,12 +52,15 @@ def machine_mapping_caller(*args, **kwargs):
                     default_options = {item: value for item, value in default_options.items() if not item.startswith('_')}
 
             # call
-            update_mapping = kwargs.pop("update_callback")
+            update_mapping = None
+            if "update_callback" in kwargs:
+                update_mapping = kwargs.pop("update_callback")
             out = f(*args, **kwargs)
             #update mappings definitions
-            if clean_ods and omas_git_repo:
-                for ulocation in numpy.unique(list(map(o2u, args[0].flat().keys()))):
-                    update_mapping(machine, ulocation, {'PYTHON': call}, 11, default_options, update_path=True)
+            if not update_mapping is None:
+                if clean_ods and omas_git_repo:
+                    for ulocation in numpy.unique(list(map(o2u, args[0].flat().keys()))):
+                        update_mapping(machine, ulocation, {'PYTHON': call}, 11, default_options, update_path=True)
 
             return out