[processing] class name is not needed anymore to load scripts

qgis · Feb 5, 2018 · cd7bb16 · cd7bb16
1 parent 5ad3af5
commit cd7bb16
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 27 deletions.
diff --git a/python/plugins/processing/script/DeleteScriptAction.py b/python/plugins/processing/script/DeleteScriptAction.py
@@ -52,5 +52,11 @@ def execute(self):
                                      QMessageBox.No)
         if reply == QMessageBox.Yes:
             filePath = ScriptUtils.findAlgorithmSource(self.itemData.__class__.__name__)
-            os.remove(filePath)
-            QgsApplication.processingRegistry().providerById("script").refreshAlgorithms()
+            if filePath is not None:
+                os.remove(filePath)
+                QgsApplication.processingRegistry().providerById("script").refreshAlgorithms()
+            else:
+                QgsMessageBox.warning(None,
+                                      self.tr("File not found"),
+                                      self.tr("Can not find corresponding script file.")
+                                      )
diff --git a/python/plugins/processing/script/EditScriptAction.py b/python/plugins/processing/script/EditScriptAction.py
@@ -46,5 +46,11 @@ def isEnabled(self):
 
     def execute(self):
         filePath = ScriptUtils.findAlgorithmSource(self.itemData.__class__.__name__)
-        dlg = ScriptEditorDialog(filePath, iface.mainWindow())
-        dlg.show()
+        if filePath is not None:
+            dlg = ScriptEditorDialog(filePath, iface.mainWindow())
+            dlg.show()
+        else:
+            QgsMessageBox.warning(None,
+                                  self.tr("File not found"),
+                                  self.tr("Can not find corresponding script file.")
+                                  )
diff --git a/python/plugins/processing/script/ScriptAlgorithmProvider.py b/python/plugins/processing/script/ScriptAlgorithmProvider.py
@@ -104,9 +104,9 @@ def loadAlgorithms(self):
             items = os.scandir(folder)
             for entry in items:
                 if entry.name.lower().endswith(".py") and entry.is_file():
-                    algName = os.path.splitext(entry.name)[0]
+                    moduleName = os.path.splitext(entry.name)[0]
                     filePath = os.path.abspath(os.path.join(folder, entry.name))
-                    alg = ScriptUtils.loadAlgorithm(algName, filePath)
+                    alg = ScriptUtils.loadAlgorithm(moduleName, filePath)
                     if alg is not None:
                         self.algs.append(alg)
 

diff --git a/python/plugins/processing/script/ScriptEditorDialog.py b/python/plugins/processing/script/ScriptEditorDialog.py
@@ -28,15 +28,16 @@
 import os
 import codecs
 import inspect
+import traceback
 
 from qgis.PyQt import uic
 from qgis.PyQt.QtCore import Qt
 from qgis.PyQt.QtGui import QCursor
 from qgis.PyQt.QtWidgets import (QMessageBox,
                                  QFileDialog)
 
-from qgis.gui import QgsGui
-from qgis.core import QgsApplication, QgsSettings
+from qgis.gui import QgsGui, QgsErrorDialog
+from qgis.core import QgsApplication, QgsSettings, QgsError, QgsProcessingAlgorithm
 from qgis.utils import iface, OverrideCursor
 
 from processing.gui.AlgorithmDialog import AlgorithmDialog
@@ -162,19 +163,21 @@ def saveAs(self):
         self.saveScript(True)
 
     def saveScript(self, saveAs):
+        newPath = None
         if self.filePath is None or saveAs:
             scriptDir = ScriptUtils.scriptsFolders()[0]
             newPath, _ = QFileDialog.getSaveFileName(self,
                                                      self.tr("Save script"),
                                                      scriptDir,
                                                      self.tr("Script files (*.py *.PY)"))
 
-        if newPath:
-            if not newPath.lower().endswith(".py"):
-                newPath += ".py"
+            if newPath:
+                if not newPath.lower().endswith(".py"):
+                    newPath += ".py"
 
-            self.filePath = newPath
+                self.filePath = newPath
 
+        if self.filePath:
             text = self.editor.text()
             try:
                 with codecs.open(self.filePath, "w", encoding="utf-8") as f:
@@ -194,10 +197,28 @@ def setHasChanged(self, hasChanged):
 
     def runAlgorithm(self):
         d = {}
-        exec(self.editor.text(), d)
+        try:
+            exec(self.editor.text(), d)
+        except Exception as e:
+            error = QgsError(traceback.format_exc(), "Processing")
+            QgsErrorDialog.show(error,
+                                self.tr("Execution error")
+                                )
+            return
+
+        alg = None
+        for k, v in d.items():
+            if inspect.isclass(v) and issubclass(v, QgsProcessingAlgorithm) and v.__name__ != "QgsProcessingAlgorithm":
+                alg = v()
+                break
+
+        if alg is None:
+            QMessageBox.warning(self,
+                                self.tr("No script found"),
+                                self.tr("Seems there is no valid script in the file.")
+                                )
+            return
 
-        className = d["__all__"][0]
-        alg = d[className]()
         alg.setProvider(QgsApplication.processingRegistry().providerById("script"))
         alg.initAlgorithm()
 

diff --git a/python/plugins/processing/script/ScriptUtils.py b/python/plugins/processing/script/ScriptUtils.py
@@ -36,6 +36,7 @@
 from processing.core.ProcessingConfig import ProcessingConfig
 from processing.tools.system import mkdir, userFolder
 
+scriptsRegistry = dict()
 
 SCRIPTS_FOLDERS = "SCRIPTS_FOLDERS"
 
@@ -55,26 +56,26 @@ def scriptsFolders():
 
 
 def loadAlgorithm(moduleName, filePath):
+    global scriptsRegistry
+
     try:
         spec = importlib.util.spec_from_file_location(moduleName, filePath)
         module = importlib.util.module_from_spec(spec)
         spec.loader.exec_module(module)
-        className = module.__all__[0]
-        obj = getattr(module, className)
-        return obj()
+        for x in dir(module):
+            obj = getattr(module, x)
+            if inspect.isclass(obj) and issubclass(obj, QgsProcessingAlgorithm) and obj.__name__ != "QgsProcessingAlgorithm":
+                scriptsRegistry[x] = filePath
+                return obj()
     except ImportError as e:
         QgsMessageLog.logMessage("Could not import script algorithm '{}' from '{}'\n{}".format(moduleName, filePath, str(e)),
                                  "Processing",
                                  QgsMessageLog.CRITICAL)
 
 
 def findAlgorithmSource(className):
-    fileName = "{}.py".format(className)
-    folders = scriptsFolders()
-    for folder in folders:
-        items = os.scandir(folder)
-        for entry in items:
-            if entry.is_file() and entry.name == fileName:
-                return os.path.abspath(os.path.join(folder, fileName))
-
-    return None
+    global scriptsRegistry
+    try:
+        return scriptsRegistry[className]
+    except:
+        return None