From 86e12fc4378d583e144a69b34b569066286f50b6 Mon Sep 17 00:00:00 2001
From: Lipu Fei <lipu.fei815@gmail.com>
Date: Sat, 3 Mar 2018 15:21:29 +0100
Subject: [PATCH] Add getDefaultVariantNode() in VariantManager

---
 cura/Machines/VariantManager.py   | 25 ++++++++++++++++++++++++-
 cura/Settings/CuraStackBuilder.py | 26 +++++++-------------------
 2 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/cura/Machines/VariantManager.py b/cura/Machines/VariantManager.py
index f7d80100a8..6cb0a3d7b2 100644
--- a/cura/Machines/VariantManager.py
+++ b/cura/Machines/VariantManager.py
@@ -3,14 +3,18 @@
 
 from enum import Enum
 from collections import OrderedDict
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 from UM.Logger import Logger
 from UM.Settings.ContainerRegistry import ContainerRegistry
+from UM.Util import parseBool
 
 from cura.Machines.ContainerNode import ContainerNode
 from cura.Settings.GlobalStack import GlobalStack
 
+if TYPE_CHECKING:
+    from UM.Settings.DefinitionContainer import DefinitionContainer
+
 
 class VariantType(Enum):
     BUILD_PLATE = "buildplate"
@@ -86,3 +90,22 @@ class VariantManager:
                         variant_type: Optional["VariantType"] = VariantType.NOZZLE) -> dict:
         machine_definition_id = machine.definition.getId()
         return self._machine_to_variant_dict_map.get(machine_definition_id, {}).get(variant_type, {})
+
+    #
+    # Gets the default variant for the given machine definition.
+    #
+    def getDefaultVariantNode(self, machine_definition: "DefinitionContainer",
+                              variant_type: VariantType) -> Optional["ContainerNode"]:
+        machine_definition_id = machine_definition.getId()
+        preferred_variant_name = None
+        if variant_type == VariantType.BUILD_PLATE:
+            if parseBool(machine_definition.getMetaDataEntry("has_variant_buildplates", False)):
+                preferred_variant_name = machine_definition.getMetaDataEntry("preferred_variant_buildplate_name")
+        else:
+            if parseBool(machine_definition.getMetaDataEntry("has_variants", False)):
+                preferred_variant_name = machine_definition.getMetaDataEntry("preferred_variant_name")
+
+        node = None
+        if preferred_variant_name:
+            node = self.getVariantNode(machine_definition_id, preferred_variant_name, variant_type)
+        return node
diff --git a/cura/Settings/CuraStackBuilder.py b/cura/Settings/CuraStackBuilder.py
index 4d103d2537..efdfa677fc 100644
--- a/cura/Settings/CuraStackBuilder.py
+++ b/cura/Settings/CuraStackBuilder.py
@@ -41,29 +41,17 @@ class CuraStackBuilder:
 
         # get variant container for the global stack
         global_variant_container = application.empty_variant_container
-        if parseBool(machine_definition.getMetaDataEntry("has_variant_buildplates", False)):
-            global_variant_name = machine_definition.getMetaDataEntry("preferred_variant_buildplate_name")
-            if global_variant_name:
-                variant_node = variant_manager.getVariantNode(definition_id, global_variant_name,
-                                                              variant_type = VariantType.BUILD_PLATE)
-                if variant_node is None:
-                    raise RuntimeError("Cannot find buildplate variant with definition [%s] and variant name [%s]" %
-                                       (definition_id, global_variant_name))
-                global_variant_container = variant_node.getContainer()
+        global_variant_node = variant_manager.getDefaultVariantNode(machine_definition, VariantType.BUILD_PLATE)
+        if global_variant_node:
+            global_variant_container = global_variant_node.getContainer()
 
         # get variant container for extruders
         extruder_variant_container = application.empty_variant_container
-        # Only look for the preferred variant if this machine has variants
+        extruder_variant_node = variant_manager.getDefaultVariantNode(machine_definition, VariantType.NOZZLE)
         extruder_variant_name = None
-        if parseBool(machine_definition.getMetaDataEntry("has_variants", False)):
-            extruder_variant_name = machine_definition.getMetaDataEntry("preferred_variant_name")
-            if extruder_variant_name:
-                variant_node = variant_manager.getVariantNode(definition_id, extruder_variant_name)
-                # Sanity check. If you see this error, the related definition files should be fixed.
-                if variant_node is None:
-                    raise RuntimeError("Cannot find extruder variant with definition [%s] and variant name [%s]" %
-                                       (definition_id, extruder_variant_name))
-                extruder_variant_container = variant_node.getContainer()
+        if extruder_variant_node:
+            extruder_variant_container = extruder_variant_node.getContainer()
+            extruder_variant_name = extruder_variant_container.getName()
 
         generated_name = registry.createUniqueName("machine", "", name, machine_definition.getName())
         # Make sure the new name does not collide with any definition or (quality) profile