From ea959dfc0489875804c8b7126195c28af7a1c764 Mon Sep 17 00:00:00 2001
From: miquelcampos <miquel.campos@gmail.com>
Date: Thu, 21 Mar 2024 10:59:09 +0900
Subject: [PATCH] Rig builder:keyerror catch for all config and black format
 #395

---
 .../scripts/mgear/shifter/rig_builder/ui.py   | 61 +++++++++++++------
 1 file changed, 43 insertions(+), 18 deletions(-)

diff --git a/release/scripts/mgear/shifter/rig_builder/ui.py b/release/scripts/mgear/shifter/rig_builder/ui.py
index 612a77c9..fac2ef75 100644
--- a/release/scripts/mgear/shifter/rig_builder/ui.py
+++ b/release/scripts/mgear/shifter/rig_builder/ui.py
@@ -8,6 +8,7 @@
 from mgear.shifter.rig_builder import builder
 
 from functools import partial
+import pymel.core as pm
 
 
 class RigBuilderUI(
@@ -55,7 +56,6 @@ def create_layout(self):
             icon="mgear_folder", width=25
         )
 
-
         output_folder_layout.addWidget(QtWidgets.QLabel("Output Folder"))
         output_folder_layout.addWidget(self.output_folder_line_edit)
         output_folder_layout.addWidget(self.output_folder_button)
@@ -114,8 +114,6 @@ def create_layout(self):
             3, QtWidgets.QHeaderView.ResizeToContents
         )
 
-        
-
         self.layout.addWidget(self.table_widget)
 
         # Pre-scrtip
@@ -256,16 +254,24 @@ def collect_table_data(self):
         for i in range(row_count):
             file_path_item = self.table_widget.item(i, 0)
             output_name_item = self.table_widget.item(i, 1)
-            custom_output_path_item = self.table_widget.item(i, 2) 
+            custom_output_path_item = self.table_widget.item(i, 2)
 
             file_path = file_path_item.text().strip() if file_path_item else ""
             output_name = (
                 output_name_item.text().strip() if output_name_item else ""
             )
-            custom_output_path = custom_output_path_item.text().strip() if custom_output_path_item else ""
+            custom_output_path = (
+                custom_output_path_item.text().strip()
+                if custom_output_path_item
+                else ""
+            )
 
             data["rows"].append(
-                {"file_path": file_path, "output_name": output_name, "custom_output_path":custom_output_path}
+                {
+                    "file_path": file_path,
+                    "output_name": output_name,
+                    "custom_output_path": custom_output_path,
+                }
             )
 
         return json.dumps(data, indent=4)
@@ -291,10 +297,16 @@ def add_file(self, file_path):
         custom_output_path = QtWidgets.QTableWidgetItem("")
         self.table_widget.setItem(row_position, 2, custom_output_path)
 
-        set_custom_output_bttn = widgets.create_button(icon="mgear_folder", setMax=False, size=37.5)
-        
-        self.table_widget.setCellWidget(row_position, 3, set_custom_output_bttn)
-        set_custom_output_bttn.clicked.connect(partial(self.on_custom_path_clicked, custom_output_path))
+        set_custom_output_bttn = widgets.create_button(
+            icon="mgear_folder", setMax=False, size=37.5
+        )
+
+        self.table_widget.setCellWidget(
+            row_position, 3, set_custom_output_bttn
+        )
+        set_custom_output_bttn.clicked.connect(
+            partial(self.on_custom_path_clicked, custom_output_path)
+        )
 
     def create_results_popup(self, results_dict):
         """Launches a pop-up containing validator results.
@@ -326,7 +338,10 @@ def import_config(self, file_path=""):
 
         self.output_folder_line_edit.setText(data["output_folder"])
         self.pre_script_line_edit.setText(data["pre_script"])
-        self.post_script_line_edit.setText(data["post_script"])
+        try:
+            self.post_script_line_edit.setText(data["post_script"])
+        except KeyError:
+            pm.displayInfo("Post Script not available in Config")
 
         self.table_widget.clearContents()
         data_rows = data["rows"]
@@ -343,14 +358,25 @@ def import_config(self, file_path=""):
             output_item = QtWidgets.QTableWidgetItem(output_name)
             self.table_widget.setItem(row_position, 1, output_item)
 
-            custom_output_path = row["custom_output_path"]
-            custom_output_path_item = QtWidgets.QTableWidgetItem(custom_output_path)
+            try:
+                custom_output_path = row["custom_output_path"]
+            except KeyError:
+                custom_output_path = ""
+            custom_output_path_item = QtWidgets.QTableWidgetItem(
+                custom_output_path
+            )
             self.table_widget.setItem(row_position, 2, custom_output_path_item)
 
-            set_custom_output_bttn = widgets.create_button(icon="mgear_folder", setMax=False, size=37.5)
-            
-            self.table_widget.setCellWidget(row_position, 3, set_custom_output_bttn)
-            set_custom_output_bttn.clicked.connect(partial(self.on_custom_path_clicked, custom_output_path_item))
+            set_custom_output_bttn = widgets.create_button(
+                icon="mgear_folder", setMax=False, size=37.5
+            )
+
+            self.table_widget.setCellWidget(
+                row_position, 3, set_custom_output_bttn
+            )
+            set_custom_output_bttn.clicked.connect(
+                partial(self.on_custom_path_clicked, custom_output_path_item)
+            )
 
     def export_config(self):
         data_string = self.collect_table_data()
@@ -365,7 +391,6 @@ def on_custom_path_clicked(self, custom_path_widget):
             custom_path_widget.setText(folder_path)
 
 
-
 class ResultsPopupDialog(QtWidgets.QDialog):
     """
     A custom pop-up to display Pyblish validator results.