From 07605df9b0ef9a49e2dca0f09f7e66426d9a7bd5 Mon Sep 17 00:00:00 2001
From: tangkong <tangkongrobert@gmail.com>
Date: Tue, 4 Jun 2024 13:49:52 -0700
Subject: [PATCH] MNT: add allowed types to FillTemplatePage to prevent active
 template in passive checkouts

---
 atef/widgets/config/find_replace.py | 21 ++++++++++++++++++++-
 atef/widgets/config/page.py         | 19 +++++++++++++------
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/atef/widgets/config/find_replace.py b/atef/widgets/config/find_replace.py
index 689a238d..3f73e2fc 100644
--- a/atef/widgets/config/find_replace.py
+++ b/atef/widgets/config/find_replace.py
@@ -9,7 +9,7 @@
 from functools import partial
 from pathlib import Path
 from typing import (TYPE_CHECKING, Any, ClassVar, Iterable, List, Optional,
-                    Union)
+                    Tuple, Union)
 
 import happi
 import qtawesome as qta
@@ -403,11 +403,14 @@ def __init__(
         *args,
         filepath: Optional[str] = None,
         window: Optional[Window] = None,
+        allowed_types: Optional[Tuple[Any]] = None,
         **kwargs
     ) -> None:
         super().__init__(*args, **kwargs)
         self._window = window
         self.fp = filepath
+        self.orig_file = None
+        self.allowed_types = allowed_types
         self.staged_actions: List[FindReplaceAction] = []
         self._signals: List[str] = []
         self._devices: List[str] = []
@@ -501,6 +504,18 @@ def load_file(self, filepath: str) -> None:
                 logger.error('failed to open file as either active '
                              'or passive checkout')
 
+        if self.allowed_types and not isinstance(data, self.allowed_types):
+            logger.error("loaded checkout is of a disallowed type: "
+                         f"({type(data)})")
+            QtWidgets.QMessageBox.warning(
+                self,
+                'Template Checkout type error',
+                f'Loaded checkout is one of the allowed types: {self.allowed_types}'
+            )
+            self.fp = None
+            self.orig_file = None
+            return
+
         self.fp = filepath
         self.orig_file = data
 
@@ -540,6 +555,10 @@ def _fill_devices_list(self) -> None:
 
     def setup_tree_view(self) -> None:
         """Populate tree view with preview of loaded file"""
+        if self.orig_file is None:
+            # clear tree
+            self.tree_view.setModel(None)
+            return
         root_item = create_tree_from_file(data=self.orig_file)
 
         model = ConfigTreeModel(data=root_item)
diff --git a/atef/widgets/config/page.py b/atef/widgets/config/page.py
index df28b9c9..c0369aec 100644
--- a/atef/widgets/config/page.py
+++ b/atef/widgets/config/page.py
@@ -33,7 +33,7 @@
 from atef.check import (ALL_COMPARISONS, AnyComparison, AnyValue, Comparison,
                         Equals, Greater, GreaterOrEqual, Less, LessOrEqual,
                         NotEquals, Range, ValueSet)
-from atef.config import (Configuration, ConfigurationGroup,
+from atef.config import (Configuration, ConfigurationFile, ConfigurationGroup,
                          DeviceConfiguration, PreparedConfiguration,
                          PreparedGroup, PreparedTemplateConfiguration,
                          PVConfiguration, TemplateConfiguration,
@@ -41,8 +41,9 @@
 from atef.procedure import (ComparisonToTarget, DescriptionStep, PassiveStep,
                             PreparedDescriptionStep, PreparedPassiveStep,
                             PreparedProcedureStep, PreparedSetValueStep,
-                            PreparedTemplateStep, ProcedureGroup,
-                            ProcedureStep, SetValueStep, TemplateStep)
+                            PreparedTemplateStep, ProcedureFile,
+                            ProcedureGroup, ProcedureStep, SetValueStep,
+                            TemplateStep)
 from atef.tools import Ping, PingResult, Tool, ToolResult
 from atef.type_hints import AnyDataclass
 from atef.widgets.config.data_active import (CheckRowWidget,
@@ -1414,16 +1415,22 @@ class TemplateConfigurationPage(DesignerDisplay, PageWidget):
     template_page_widget: FillTemplatePage
     template_page_placeholder: QWidget
 
-    data: TemplateConfiguration
+    data: Union[TemplateConfiguration, TemplateStep]
+    ALLOWED_TYPE_MAP: ClassVar[Dict[Any, Tuple[Any]]] = {
+        TemplateConfiguration: (ConfigurationFile),
+        TemplateStep: (ConfigurationFile, ProcedureFile)
+    }
 
-    def __init__(self, data: TemplateConfiguration, **kwargs):
+    def __init__(self, data: Union[TemplateConfiguration, TemplateStep], **kwargs):
         super().__init__(data=data, **kwargs)
         self.setup_name_desc_tags_init()
         self.setup_template_widget_init()
         self.post_tree_setup()
 
     def setup_template_widget_init(self) -> None:
-        self.template_page_widget = FillTemplatePage()
+        self.template_page_widget = FillTemplatePage(
+            allowed_types=self.ALLOWED_TYPE_MAP[type(self.data)]
+        )
 
         def finish_widget_setup(*args, **kwargs):
             # only run this once, when we're loading an existing template checkout