Initial upload/instantiation of custom nodes

pyworkflow/pyworkflow/node_factory.py

@@ -14,7 +14,7 @@ def node_factory(node_info):
     elif node_type == 'FlowNode':
         new_node = flow_node(node_key, node_info)
     else:
-        new_node = None
+        new_node = custom_node(node_type, node_key, node_info)
 
     return new_node
 
@@ -46,3 +46,16 @@ def manipulation_node(node_key, node_info):
         return FilterNode(node_info)
     else:
         return None
+
+
+def custom_node(filename, node_key, node_info):
+    try:
+        package = __import__('custom_nodes.' + filename)
+        module = getattr(package, filename)
+        my_class = getattr(module, node_key)
+        instance = my_class(node_info)
+
+        return instance
+    except Exception as e:
+        print(str(e))
+        return None

pyworkflow/pyworkflow/workflow.py

@@ -19,9 +19,14 @@ class Workflow:
     def __init__(self, name="Untitled", root_dir=None, graph=nx.DiGraph(), flow_vars=nx.Graph()):
         self._name = name
         self._root_dir = WorkflowUtils.set_root_dir(root_dir)
+        self._custom_node_dir = WorkflowUtils.set_custom_nodes_dir()
         self._graph = graph
         self._flow_vars = flow_vars
 
+    @property
+    def custom_node_dir(self):
+        return self._custom_node_dir
+
     @property
     def graph(self):
         return self._graph
@@ -404,6 +409,15 @@ def execute_workflow(workflow_location):
 
 
 class WorkflowUtils:
+    @staticmethod
+    def set_custom_nodes_dir():
+        custom_node_dir = os.path.join(os.getcwd(), '../pyworkflow/custom_nodes')
+
+        if not os.path.exists(custom_node_dir):
+            os.makedirs(custom_node_dir)
+
+        return custom_node_dir
+
     @staticmethod
     def set_root_dir(root_dir):
         if root_dir is None:

vp/vp/views.py

@@ -2,6 +2,11 @@
 from rest_framework.decorators import api_view
 from drf_yasg.utils import swagger_auto_schema
 from pyworkflow import Node
+from modulefinder import ModuleFinder
+
+import os
+import inspect
+import sys
 
 
 @swagger_auto_schema(method='get', responses={200:'JSON response with data'})
@@ -39,27 +44,84 @@ def retrieve_nodes_for_user(request):
     """
     data = dict()
 
-    # Iterate through node 'types'
+    # Iterate through installed Nodes
     for parent in Node.__subclasses__():
         key = getattr(parent, "display_name", parent.__name__)
         data[key] = list()
 
         # Iterate through node 'keys'
         for child in parent.__subclasses__():
-            # TODO: check attribute-scope is handled correctly
-            child_node = {
-                'name': child.name,
-                'node_key': child.__name__,
-                'node_type': parent.__name__,
-                'num_in': child.num_in,
-                'num_out': child.num_out,
-                'color': child.color or parent.color,
-                'doc': child.__doc__,
-                'options': {k: v.get_value() for k, v in child.options.items()},
-                'option_types': child.option_types,
-                'download_result':  getattr(child, "download_result", False)
-            }
-
-            data[key].append(child_node)
-
-    return JsonResponse(data)
+            node = extract_node_info(parent, child)
+            data[key].append(node)
+
+    # Check for any installed Custom Nodes
+    # TODO: Workflow loading excluded in middleware for this route
+    #       Should probably have a way to access the 'custom_node` dir dynamically
+    custom_node_path = os.path.join(os.getcwd(), '../pyworkflow/custom_nodes')
+    data['CustomNode'] = import_custom_node(custom_node_path)
+
+    return JsonResponse(data)
+
+
+def check_missing_packages(node_path):
+    finder = ModuleFinder(node_path)
+    finder.run_script(node_path)
+
+    uninstalled = list()
+    for missing_package in finder.badmodules.keys():
+        if missing_package not in sys.modules:
+            uninstalled.append(missing_package)
+
+    return uninstalled
+
+
+def extract_node_info(parent, child):
+    # TODO: check attribute(s) accessing is handled correctly
+    return {
+        'name': child.name,
+        'node_key': child.__name__,
+        'node_type': str(parent),
+        'num_in': child.num_in,
+        'num_out': child.num_out,
+        'color': child.color or parent.color or 'black',
+        'doc': child.__doc__,
+        'options': {k: v.get_value() for k, v in child.options.items()},
+        'option_types': child.option_types,
+        'download_result': getattr(child, "download_result", False)
+    }
+
+
+def import_custom_node(root_path):
+    # Get list of files in path
+    try:
+        files = os.listdir(root_path)
+    except OSError as e:
+        return None
+
+    data = list()
+    for file in files:
+        # Check file is not a dir
+        node_path = os.path.join(root_path, file)
+        if not os.path.isfile(node_path):
+            continue
+
+        node, ext = os.path.splitext(file)
+
+        try:
+            package = __import__('custom_nodes.' + node)
+            module = getattr(package, node)
+        except ModuleNotFoundError:
+            data.append({
+                "name": node,
+                "missing_packages": check_missing_packages(node_path)
+            })
+            continue
+
+        for name, klass in inspect.getmembers(module):
+            if inspect.isclass(klass) and klass.__module__.startswith('custom_nodes.'):
+                custom_node = extract_node_info(node, klass)
+                data.append(custom_node)
+
+    return data
+
+

vp/workflow/urls.py

@@ -10,5 +10,6 @@
     path('execute/<str:node_id>/successors', views.get_successors, name='get node successors'),
     path('globals', views.global_vars, name="retrieve global variables"),
     path('upload', views.upload_file, name='upload file'),
-    path('download', views.download_file, name='download file')
+    path('download', views.download_file, name='download file'),
+    path('custom_node', views.add_custom_node, name='add custom node')
 ]

vp/workflow/views.py

@@ -1,12 +1,14 @@
 import os
 import json
+import sys
 
 from django.http import JsonResponse, HttpResponse
 from django.conf import settings
 from rest_framework.decorators import api_view
 from pyworkflow import Workflow, WorkflowException
 from drf_yasg.utils import swagger_auto_schema
 
+from modulefinder import ModuleFinder
 
 @swagger_auto_schema(method='post',
                      operation_summary='Create a new workflow.',
@@ -216,6 +218,36 @@ def get_successors(request, node_id):
     return JsonResponse(order, safe=False)
 
 
+@swagger_auto_schema(method='post',
+                     operation_summary='Uploads a custom node to server.',
+                     operation_description='Uploads a custom node to server location.',
+                     responses={
+                         200: 'File uploaded',
+                         404: 'No specified file'
+                     })
+@api_view(['POST'])
+def add_custom_node(request):
+    node_file = request.FILES.get('file')
+
+    if node_file is None:
+        return JsonResponse("Empty content", status=404)
+
+    to_open = os.path.join(request.pyworkflow.custom_node_dir, node_file.name)
+
+    try:
+        with open(to_open, 'wb') as f:
+            f.write(node_file.read())
+    except OSError as e:
+        return JsonResponse({'message': str(e)}, status=500)
+
+    node_file.close()
+
+    return JsonResponse({
+        "filename": to_open,
+        "missing_packages": check_missing_packages(to_open)
+    }, status=201, safe=False)
+
+
 @swagger_auto_schema(method='post',
                      operation_summary='Uploads a file to server.',
                      operation_description='Uploads a new file to server location.',
@@ -275,3 +307,14 @@ def download_file(request):
     except WorkflowException as e:
         return JsonResponse({e.action: e.reason}, status=500)
 
+
+def check_missing_packages(node_path):
+    finder = ModuleFinder(node_path)
+    finder.run_script(node_path)
+
+    uninstalled = list()
+    for missing_package in finder.badmodules.keys():
+        if missing_package not in sys.modules:
+            uninstalled.append(missing_package)
+
+    return uninstalled