Added load_plugin_safe method to return calc node if plugin is not available

aiida/backends/djsite/db/models.py

@@ -189,7 +189,7 @@ def get_aiida_class(self):
         """
         from aiida.orm.node import Node
         from aiida.common.old_pluginloader import from_type_to_pluginclassname
-        from aiida.common.pluginloader import load_plugin
+        from aiida.common.pluginloader import load_plugin_safe
         from aiida.common import aiidalogger
 
         try:
@@ -199,8 +199,9 @@ def get_aiida_class(self):
                                  "not valid: '{}'".format(self.pk, self.type))
 
         try:
-            PluginClass = load_plugin(Node, 'aiida.orm', pluginclassname)
+            PluginClass = load_plugin_safe(Node, 'aiida.orm', pluginclassname)
         except MissingPluginError:
+
             aiidalogger.error("Unable to find plugin for type '{}' (node= {}), "
                               "will use base Node class".format(self.type, self.pk))
             PluginClass = Node

aiida/backends/sqlalchemy/models/node.py

@@ -28,7 +28,6 @@
 from aiida.backends.sqlalchemy.models.utils import uuid_func
 
 from aiida.common import aiidalogger
-from aiida.common.pluginloader import load_plugin
 from aiida.common.exceptions import DbContentError, MissingPluginError
 from aiida.common.datastructures import calc_states, _sorted_datastates, sort_states
 
@@ -154,6 +153,7 @@ def get_aiida_class(self):
         """
         from aiida.common.old_pluginloader import from_type_to_pluginclassname
         from aiida.orm.node import Node
+        from aiida.common.pluginloader import load_plugin_safe
 
         try:
             pluginclassname = from_type_to_pluginclassname(self.type)
@@ -162,7 +162,7 @@ def get_aiida_class(self):
                                  "not valid: '{}'".format(self.pk, self.type))
 
         try:
-            PluginClass = load_plugin(Node, 'aiida.orm', pluginclassname)
+            PluginClass = load_plugin_safe(Node, 'aiida.orm', pluginclassname)
         except MissingPluginError:
             aiidalogger.error("Unable to find plugin for type '{}' (node= {}), "
                               "will use base Node class".format(self.type, self.pk))

aiida/backends/tests/nodes.py

@@ -1449,6 +1449,49 @@ def test_load_node(self):
             with self.assertRaises(NotExistent):
                 load_node(spec, parent_class=ArrayData)
 
+    def test_load_plugin_safe(self):
+        from aiida.orm import (JobCalculation, CalculationFactory, DataFactory)
+
+        ###### for calculation
+        calc_params = {
+            'computer': self.computer,
+            'resources': {'num_machines': 1, 'num_mpiprocs_per_machine': 1}
+        }
+
+        TemplateReplacerCalc = CalculationFactory('simpleplugins.templatereplacer')
+        testcalc = TemplateReplacerCalc(**calc_params).store()
+        jobcalc = JobCalculation(**calc_params).store()
+
+        # compare if plugin exist
+        obj = testcalc.dbnode.get_aiida_class()
+        self.assertEqual(type(testcalc), type(obj))
+
+        # change node type and save in database again
+        testcalc.dbnode.type = "calculation.job.simpleplugins_tmp.templatereplacer.TemplatereplacerCalculation."
+        testcalc.dbnode.save()
+
+        # changed node should return job calc as its plugin is not exist
+        obj = testcalc.dbnode.get_aiida_class()
+        self.assertEqual(type(jobcalc), type(obj))
+
+        ####### for data
+        KpointsData = DataFactory('array.kpoints')
+        kpoint = KpointsData().store()
+        Data = DataFactory("Data")
+        data = Data().store()
+
+        # compare if plugin exist
+        obj = kpoint.dbnode.get_aiida_class()
+        self.assertEqual(type(kpoint), type(obj))
+
+        # change node type and save in database again
+        kpoint.dbnode.type = "data.array.kpoints_tmp.KpointsData."
+        kpoint.dbnode.save()
+
+        # changed node should return data node as its plugin is not exist
+        obj = kpoint.dbnode.get_aiida_class()
+        self.assertEqual(type(data), type(obj))
+
 
 class TestSubNodesAndLinks(AiidaTestCase):
 

aiida/common/pluginloader.py

@@ -125,6 +125,48 @@ def get_plugin(category, name):
 
     return plugin
 
+def load_plugin_safe(base_class, plugins_module, plugin_type):
+    """
+    It is a copy of load_plugin function to return closely related node class
+    if plugin is not available. We are duplicating load_plugin function to not break
+    its default behaviour.
+
+    params: Look at the docstring of aiida.common.old_pluginloader.load_plugin for more Info
+    :return: The plugin class
+    """
+
+    try:
+        PluginClass = load_plugin(base_class, plugins_module, plugin_type)
+    except MissingPluginError:
+        nodeParts = plugin_type.partition(".")
+        baseNodeType = nodeParts[0]
+
+        ## data node: temporarily returning base data node.
+        # In future its better to check the closest available plugin and return it.
+        # For example if type is "aiida.orm.data.array.kpoints_tmp.KpointsData"
+        # it should return array data node and not base data node
+        if baseNodeType == "data":
+            PluginClass = load_plugin(base_class, plugins_module, 'data.Data')
+
+        ## code node
+        elif baseNodeType == "code":
+            PluginClass = load_plugin(base_class, plugins_module, 'code.Code')
+
+        ## calculation node: for calculation currently we are hardcoding cases
+        elif baseNodeType == "calculation":
+            subNodeParts = nodeParts[2].partition(".")
+            subNodeType = subNodeParts[0]
+            if subNodeType == "job":
+                PluginClass = load_plugin(base_class, plugins_module, 'calculation.job.JobCalculation')
+            elif subNodeType == "inline":
+                PluginClass = load_plugin(base_class, plugins_module, 'calculation.inline.InlineCalculation')
+            elif subNodeType == "work":
+                PluginClass = load_plugin(base_class, plugins_module, 'calculation.work.WorkCalculation')
+            else:
+                PluginClass = load_plugin(base_class, plugins_module, 'calculation.Calculation')
+
+    return PluginClass
+
 
 def load_plugin(base_class, plugins_module, plugin_type):
     """