node_wrangler/operators.py

@@ -2339,21 +2339,22 @@ class NWAlignNodes(Operator, NWBase):
     bl_options = {'REGISTER', 'UNDO'}
     margin: IntProperty(name='Margin', default=50, description='The amount of space between nodes')
 
+    @classmethod
+    def poll(cls, context):
+        selection = [node for node in context.selected_nodes if node.type != 'FRAME']
+        if len(selection) < 2:
+            return False
+
+        return nw_check(context)
+
     def execute(self, context):
-        nodes, links = get_nodes_links(context)
+        selection = [node for node in context.selected_nodes if node.type != 'FRAME']
+        active_node = context.active_node
         margin = self.margin
 
-        selection = []
-        for node in nodes:
-            if node.select and node.type != 'FRAME':
-                selection.append(node)
-
-        # If no nodes are selected, align all nodes
         active_loc = None
-        if not selection:
-            selection = nodes
-        elif nodes.active in selection:
-            active_loc = copy(nodes.active.location)  # make a copy, not a reference
+        if active_node in selection:
+            active_loc = copy(active_node.location)  # make a copy, not a reference
 
         # Check if nodes should be laid out horizontally or vertically
         # use dimension to get center of node, not corner
@@ -2388,7 +2389,7 @@ class NWAlignNodes(Operator, NWBase):
 
         # If active node is selected, center nodes around it
         if active_loc is not None:
-            active_loc_diff = active_loc - nodes.active.location
+            active_loc_diff = active_loc - active_node.location
             for node in selection:
                 node.location += active_loc_diff
         else:  # Position nodes centered around where they used to be