PostProcessingPlugin.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313
  1. # Copyright (c) 2018 Jaime van Kessel, Ultimaker B.V.
  2. # The PostProcessingPlugin is released under the terms of the AGPLv3 or higher.
  3. from PyQt5.QtCore import QObject, pyqtProperty, pyqtSignal, pyqtSlot
  4. from typing import Dict, Type, TYPE_CHECKING, List, Optional, cast
  5. from UM.PluginRegistry import PluginRegistry
  6. from UM.Resources import Resources
  7. from UM.Application import Application
  8. from UM.Extension import Extension
  9. from UM.Logger import Logger
  10. import configparser # The script lists are stored in metadata as serialised config files.
  11. import io # To allow configparser to write to a string.
  12. import os.path
  13. import pkgutil
  14. import sys
  15. import importlib.util
  16. from UM.i18n import i18nCatalog
  17. from cura.CuraApplication import CuraApplication
  18. i18n_catalog = i18nCatalog("cura")
  19. if TYPE_CHECKING:
  20. from .Script import Script
  21. ## The post processing plugin is an Extension type plugin that enables pre-written scripts to post process generated
  22. # g-code files.
  23. class PostProcessingPlugin(QObject, Extension):
  24. def __init__(self, parent = None) -> None:
  25. QObject.__init__(self, parent)
  26. Extension.__init__(self)
  27. self.setMenuName(i18n_catalog.i18nc("@item:inmenu", "Post Processing"))
  28. self.addMenuItem(i18n_catalog.i18nc("@item:inmenu", "Modify G-Code"), self.showPopup)
  29. self._view = None
  30. # Loaded scripts are all scripts that can be used
  31. self._loaded_scripts = {} # type: Dict[str, Type[Script]]
  32. self._script_labels = {} # type: Dict[str, str]
  33. # Script list contains instances of scripts in loaded_scripts.
  34. # There can be duplicates, which will be executed in sequence.
  35. self._script_list = [] # type: List[Script]
  36. self._selected_script_index = -1
  37. Application.getInstance().getOutputDeviceManager().writeStarted.connect(self.execute)
  38. Application.getInstance().globalContainerStackChanged.connect(self._onGlobalContainerStackChanged) # When the current printer changes, update the list of scripts.
  39. CuraApplication.getInstance().mainWindowChanged.connect(self._createView) # When the main window is created, create the view so that we can display the post-processing icon if necessary.
  40. selectedIndexChanged = pyqtSignal()
  41. @pyqtProperty(str, notify = selectedIndexChanged)
  42. def selectedScriptDefinitionId(self) -> Optional[str]:
  43. try:
  44. return self._script_list[self._selected_script_index].getDefinitionId()
  45. except IndexError:
  46. return ""
  47. @pyqtProperty(str, notify=selectedIndexChanged)
  48. def selectedScriptStackId(self) -> Optional[str]:
  49. try:
  50. return self._script_list[self._selected_script_index].getStackId()
  51. except IndexError:
  52. return ""
  53. ## Execute all post-processing scripts on the gcode.
  54. def execute(self, output_device) -> None:
  55. scene = Application.getInstance().getController().getScene()
  56. # If the scene does not have a gcode, do nothing
  57. if not hasattr(scene, "gcode_dict"):
  58. return
  59. gcode_dict = getattr(scene, "gcode_dict")
  60. if not gcode_dict:
  61. return
  62. # get gcode list for the active build plate
  63. active_build_plate_id = CuraApplication.getInstance().getMultiBuildPlateModel().activeBuildPlate
  64. gcode_list = gcode_dict[active_build_plate_id]
  65. if not gcode_list:
  66. return
  67. if ";POSTPROCESSED" not in gcode_list[0]:
  68. for script in self._script_list:
  69. try:
  70. gcode_list = script.execute(gcode_list)
  71. except Exception:
  72. Logger.logException("e", "Exception in post-processing script.")
  73. if len(self._script_list): # Add comment to g-code if any changes were made.
  74. gcode_list[0] += ";POSTPROCESSED\n"
  75. gcode_dict[active_build_plate_id] = gcode_list
  76. setattr(scene, "gcode_dict", gcode_dict)
  77. else:
  78. Logger.log("e", "Already post processed")
  79. @pyqtSlot(int)
  80. def setSelectedScriptIndex(self, index: int) -> None:
  81. if self._selected_script_index != index:
  82. self._selected_script_index = index
  83. self.selectedIndexChanged.emit()
  84. @pyqtProperty(int, notify = selectedIndexChanged)
  85. def selectedScriptIndex(self) -> int:
  86. return self._selected_script_index
  87. @pyqtSlot(int, int)
  88. def moveScript(self, index: int, new_index: int) -> None:
  89. if new_index < 0 or new_index > len(self._script_list) - 1:
  90. return # nothing needs to be done
  91. else:
  92. # Magical switch code.
  93. self._script_list[new_index], self._script_list[index] = self._script_list[index], self._script_list[new_index]
  94. self.scriptListChanged.emit()
  95. self.selectedIndexChanged.emit() #Ensure that settings are updated
  96. self._propertyChanged()
  97. ## Remove a script from the active script list by index.
  98. @pyqtSlot(int)
  99. def removeScriptByIndex(self, index: int) -> None:
  100. self._script_list.pop(index)
  101. if len(self._script_list) - 1 < self._selected_script_index:
  102. self._selected_script_index = len(self._script_list) - 1
  103. self.scriptListChanged.emit()
  104. self.selectedIndexChanged.emit() # Ensure that settings are updated
  105. self._propertyChanged()
  106. ## Load all scripts from all paths where scripts can be found.
  107. #
  108. # This should probably only be done on init.
  109. def loadAllScripts(self) -> None:
  110. if self._loaded_scripts: # Already loaded.
  111. return
  112. # The PostProcessingPlugin path is for built-in scripts.
  113. # The Resources path is where the user should store custom scripts.
  114. # The Preferences path is legacy, where the user may previously have stored scripts.
  115. for root in [PluginRegistry.getInstance().getPluginPath("PostProcessingPlugin"), Resources.getStoragePath(Resources.Resources), Resources.getStoragePath(Resources.Preferences)]:
  116. if root is None:
  117. continue
  118. path = os.path.join(root, "scripts")
  119. if not os.path.isdir(path):
  120. try:
  121. os.makedirs(path)
  122. except OSError:
  123. Logger.log("w", "Unable to create a folder for scripts: " + path)
  124. continue
  125. self.loadScripts(path)
  126. ## Load all scripts from provided path.
  127. # This should probably only be done on init.
  128. # \param path Path to check for scripts.
  129. def loadScripts(self, path: str) -> None:
  130. ## Load all scripts in the scripts folders
  131. scripts = pkgutil.iter_modules(path = [path])
  132. for loader, script_name, ispkg in scripts:
  133. # Iterate over all scripts.
  134. if script_name not in sys.modules:
  135. try:
  136. spec = importlib.util.spec_from_file_location(__name__ + "." + script_name, os.path.join(path, script_name + ".py"))
  137. loaded_script = importlib.util.module_from_spec(spec)
  138. if spec.loader is None:
  139. continue
  140. spec.loader.exec_module(loaded_script)
  141. sys.modules[script_name] = loaded_script #TODO: This could be a security risk. Overwrite any module with a user-provided name?
  142. loaded_class = getattr(loaded_script, script_name)
  143. temp_object = loaded_class()
  144. Logger.log("d", "Begin loading of script: %s", script_name)
  145. try:
  146. setting_data = temp_object.getSettingData()
  147. if "name" in setting_data and "key" in setting_data:
  148. self._script_labels[setting_data["key"]] = setting_data["name"]
  149. self._loaded_scripts[setting_data["key"]] = loaded_class
  150. else:
  151. Logger.log("w", "Script %s.py has no name or key", script_name)
  152. self._script_labels[script_name] = script_name
  153. self._loaded_scripts[script_name] = loaded_class
  154. except AttributeError:
  155. Logger.log("e", "Script %s.py is not a recognised script type. Ensure it inherits Script", script_name)
  156. except NotImplementedError:
  157. Logger.log("e", "Script %s.py has no implemented settings", script_name)
  158. except Exception as e:
  159. Logger.logException("e", "Exception occurred while loading post processing plugin: {error_msg}".format(error_msg = str(e)))
  160. loadedScriptListChanged = pyqtSignal()
  161. @pyqtProperty("QVariantList", notify = loadedScriptListChanged)
  162. def loadedScriptList(self) -> List[str]:
  163. return sorted(list(self._loaded_scripts.keys()))
  164. @pyqtSlot(str, result = str)
  165. def getScriptLabelByKey(self, key: str) -> Optional[str]:
  166. return self._script_labels.get(key)
  167. scriptListChanged = pyqtSignal()
  168. @pyqtProperty("QStringList", notify = scriptListChanged)
  169. def scriptList(self) -> List[str]:
  170. script_list = [script.getSettingData()["key"] for script in self._script_list]
  171. return script_list
  172. @pyqtSlot(str)
  173. def addScriptToList(self, key: str) -> None:
  174. Logger.log("d", "Adding script %s to list.", key)
  175. new_script = self._loaded_scripts[key]()
  176. new_script.initialize()
  177. self._script_list.append(new_script)
  178. self.setSelectedScriptIndex(len(self._script_list) - 1)
  179. self.scriptListChanged.emit()
  180. self._propertyChanged()
  181. ## When the global container stack is changed, swap out the list of active
  182. # scripts.
  183. def _onGlobalContainerStackChanged(self) -> None:
  184. self.loadAllScripts()
  185. new_stack = Application.getInstance().getGlobalContainerStack()
  186. if new_stack is None:
  187. return
  188. self._script_list.clear()
  189. if not new_stack.getMetaDataEntry("post_processing_scripts"): # Missing or empty.
  190. self.scriptListChanged.emit() # Even emit this if it didn't change. We want it to write the empty list to the stack's metadata.
  191. return
  192. self._script_list.clear()
  193. scripts_list_strs = new_stack.getMetaDataEntry("post_processing_scripts")
  194. for script_str in scripts_list_strs.split("\n"): # Encoded config files should never contain three newlines in a row. At most 2, just before section headers.
  195. if not script_str: # There were no scripts in this one (or a corrupt file caused more than 3 consecutive newlines here).
  196. continue
  197. script_str = script_str.replace(r"\\\n", "\n").replace(r"\\\\", "\\\\") # Unescape escape sequences.
  198. script_parser = configparser.ConfigParser(interpolation = None)
  199. script_parser.optionxform = str # type: ignore # Don't transform the setting keys as they are case-sensitive.
  200. script_parser.read_string(script_str)
  201. for script_name, settings in script_parser.items(): # There should only be one, really! Otherwise we can't guarantee the order or allow multiple uses of the same script.
  202. if script_name == "DEFAULT": # ConfigParser always has a DEFAULT section, but we don't fill it. Ignore this one.
  203. continue
  204. if script_name not in self._loaded_scripts: # Don't know this post-processing plug-in.
  205. Logger.log("e", "Unknown post-processing script {script_name} was encountered in this global stack.".format(script_name = script_name))
  206. continue
  207. new_script = self._loaded_scripts[script_name]()
  208. new_script.initialize()
  209. for setting_key, setting_value in settings.items(): # Put all setting values into the script.
  210. if new_script._instance is not None:
  211. new_script._instance.setProperty(setting_key, "value", setting_value)
  212. self._script_list.append(new_script)
  213. self.setSelectedScriptIndex(0)
  214. self.scriptListChanged.emit()
  215. @pyqtSlot()
  216. def writeScriptsToStack(self) -> None:
  217. script_list_strs = [] # type: List[str]
  218. for script in self._script_list:
  219. parser = configparser.ConfigParser(interpolation = None) # We'll encode the script as a config with one section. The section header is the key and its values are the settings.
  220. parser.optionxform = str # type: ignore # Don't transform the setting keys as they are case-sensitive.
  221. script_name = script.getSettingData()["key"]
  222. parser.add_section(script_name)
  223. for key in script.getSettingData()["settings"]:
  224. value = script.getSettingValueByKey(key)
  225. parser[script_name][key] = str(value)
  226. serialized = io.StringIO() # ConfigParser can only write to streams. Fine.
  227. parser.write(serialized)
  228. serialized.seek(0)
  229. script_str = serialized.read()
  230. script_str = script_str.replace("\\\\", r"\\\\").replace("\n", r"\\\n") # Escape newlines because configparser sees those as section delimiters.
  231. script_list_strs.append(script_str)
  232. script_list_string = "\n".join(script_list_strs) # ConfigParser should never output three newlines in a row when serialised, so it's a safe delimiter.
  233. global_stack = Application.getInstance().getGlobalContainerStack()
  234. if global_stack is None:
  235. return
  236. if "post_processing_scripts" not in global_stack.getMetaData():
  237. global_stack.setMetaDataEntry("post_processing_scripts", "")
  238. global_stack.setMetaDataEntry("post_processing_scripts", script_list_string)
  239. ## Creates the view used by show popup. The view is saved because of the fairly aggressive garbage collection.
  240. def _createView(self) -> None:
  241. Logger.log("d", "Creating post processing plugin view.")
  242. self.loadAllScripts()
  243. # Create the plugin dialog component
  244. path = os.path.join(cast(str, PluginRegistry.getInstance().getPluginPath("PostProcessingPlugin")), "PostProcessingPlugin.qml")
  245. self._view = CuraApplication.getInstance().createQmlComponent(path, {"manager": self})
  246. if self._view is None:
  247. Logger.log("e", "Not creating PostProcessing button near save button because the QML component failed to be created.")
  248. return
  249. Logger.log("d", "Post processing view created.")
  250. # Create the save button component
  251. CuraApplication.getInstance().addAdditionalComponent("saveButton", self._view.findChild(QObject, "postProcessingSaveAreaButton"))
  252. ## Show the (GUI) popup of the post processing plugin.
  253. def showPopup(self) -> None:
  254. if self._view is None:
  255. self._createView()
  256. if self._view is None:
  257. Logger.log("e", "Not creating PostProcessing window since the QML component failed to be created.")
  258. return
  259. self._view.show()
  260. ## Property changed: trigger re-slice
  261. # To do this we use the global container stack propertyChanged.
  262. # Re-slicing is necessary for setting changes in this plugin, because the changes
  263. # are applied only once per "fresh" gcode
  264. def _propertyChanged(self) -> None:
  265. global_container_stack = Application.getInstance().getGlobalContainerStack()
  266. if global_container_stack is not None:
  267. global_container_stack.propertyChanged.emit("post_processing_plugin", "value")