From 77ae23538577f907665fceb444c2a3860047c71a Mon Sep 17 00:00:00 2001
From: DarkPhoenix <phoenix@mail.ru>
Date: Tue, 30 Jul 2019 22:01:33 +0300
Subject: [PATCH] Offload actual set adding logic to windows calling the
 context menu

---
 gui/builtinAdditionPanes/implantView.py  |  5 +++++
 gui/builtinContextMenus/implantSetAdd.py | 23 +++--------------------
 gui/characterEditor.py                   | 10 ++++++++++
 3 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/gui/builtinAdditionPanes/implantView.py b/gui/builtinAdditionPanes/implantView.py
index a39b41e7f..136664401 100644
--- a/gui/builtinAdditionPanes/implantView.py
+++ b/gui/builtinAdditionPanes/implantView.py
@@ -291,3 +291,8 @@ class ImplantDisplay(d.Display):
                 continue
             implants.append(implant)
         return implants
+
+    def addImplantSet(self, impSet):
+        self.mainFrame.command.Submit(cmd.GuiAddImplantSetCommand(
+            fitID=self.mainFrame.getActiveFit(),
+            itemIDs=[i.itemID for i in impSet.implants]))
diff --git a/gui/builtinContextMenus/implantSetAdd.py b/gui/builtinContextMenus/implantSetAdd.py
index e85f149b4..f44ebc2c1 100644
--- a/gui/builtinContextMenus/implantSetAdd.py
+++ b/gui/builtinContextMenus/implantSetAdd.py
@@ -11,9 +11,6 @@ from service.implantSet import ImplantSets as s_ImplantSets
 
 class AddImplantSet(ContextMenuUnconditional):
 
-    def __init__(self):
-        self.mainFrame = gui.mainFrame.MainFrame.getInstance()
-
     def display(self, callingWindow, srcContext):
 
         sIS = s_ImplantSets.getInstance()
@@ -48,26 +45,12 @@ class AddImplantSet(ContextMenuUnconditional):
         return m
 
     def handleSelection(self, event):
-        set = self.idmap.get(event.Id, None)
-
-        if set is None:
+        impSet = self.idmap.get(event.Id, None)
+        if impSet is None:
             event.Skip()
             return
 
-        if self.context == "implantEditor":
-            charEditor = self.callingWindow.Parent.Parent
-            # we are calling from character editor, the implant source is different
-            sChar = Character.getInstance()
-            char = charEditor.entityEditor.getActiveEntity()
-
-            for implant in set.implants:
-                sChar.addImplant(char.ID, implant.item.ID)
-
-            wx.PostEvent(charEditor, GE.CharChanged())
-        else:
-            self.mainFrame.command.Submit(cmd.GuiAddImplantSetCommand(
-                fitID=self.mainFrame.getActiveFit(),
-                itemIDs=[i.itemID for i in set.implants]))
+        self.callingWindow.addImplantSet(impSet)
 
 
 AddImplantSet.register()
diff --git a/gui/characterEditor.py b/gui/characterEditor.py
index 1de9c4ac3..17c27929d 100644
--- a/gui/characterEditor.py
+++ b/gui/characterEditor.py
@@ -708,6 +708,16 @@ class ImplantEditorView(BaseImplantEditorView):
 
         sChar.removeImplant(char.ID, implant)
 
+    def addImplantSet(self, impSet):
+        charEditor = self.Parent.Parent
+        char = charEditor.entityEditor.getActiveEntity()
+
+        sChar = Character.getInstance()
+        for implant in impSet.implants:
+            sChar.addImplant(char.ID, implant.item.ID)
+
+        wx.PostEvent(charEditor, GE.CharChanged())
+
     def spawnMenu(self, event):
         context = (("implantEditor",),)
         menu = ContextMenu.getMenu(self, None, None, *context)