From 4faad280d8c45fc382c7dcc4d89439a6c43ddbd7 Mon Sep 17 00:00:00 2001
From: Tristan Strathearn <hello+tristan@hugemenace.co>
Date: Sun, 23 Jan 2022 11:19:10 +1000
Subject: [PATCH] fix: add poll method to operators to ensure all conditions
 are correct before they can be invoked

---
 bolt.py         |  5 +++++
 faux_bevel.py   |  6 ++++++
 sketch_bevel.py | 20 +++++++++++---------
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/bolt.py b/bolt.py
index 7f0bd21..0fada51 100644
--- a/bolt.py
+++ b/bolt.py
@@ -46,6 +46,11 @@ def execute(self, context):
         return self.handle_modal(context)
 
 
+    @classmethod
+    def poll(cls, context):
+        return context.mode == 'OBJECT'
+
+
     def handle_modal(self, context):
         self.add_object(context)
         self.add_screw_x_modifier()
diff --git a/faux_bevel.py b/faux_bevel.py
index 3c5a547..63c2f90 100644
--- a/faux_bevel.py
+++ b/faux_bevel.py
@@ -36,6 +36,12 @@ def execute(self, context):
         return self.handle_modal(context)
 
 
+    @classmethod
+    def poll(cls, context):
+        if context.mode == 'OBJECT':
+            return len(context.selected_objects) == 1
+
+
     def handle_modal(self, context):
         self.add_bevel_modifier(context)
         self.add_weighted_normal_modifer(context)
diff --git a/sketch_bevel.py b/sketch_bevel.py
index ab578a5..c525ef7 100644
--- a/sketch_bevel.py
+++ b/sketch_bevel.py
@@ -1,4 +1,5 @@
 import bpy
+import bmesh
 
 
 class ND_OT_sketch_bevel(bpy.types.Operator):
@@ -37,17 +38,18 @@ def invoke(self, context, event):
     def execute(self, context):
         return self.handle_modal(context)
 
+    @classmethod
+    def poll(cls, context):
+        if context.mode == 'EDIT_MESH':
+            mesh = bmesh.from_edit_mesh(context.object.data)
+            return len([vert for vert in mesh.verts if vert.select]) >= 1
 
     def handle_modal(self, context):
-        if context.object and context.object.type == 'MESH' and context.mode == 'EDIT_MESH' and context.object.data.total_vert_sel > 0:
-            self.add_vertex_group(context)
-            self.add_bevel_modifier(context)
-
-            context.window_manager.modal_handler_add(self)
-            return {'RUNNING_MODAL'}
-        else:
-            self.report({'WARNING'}, "No active object, could not finish")
-            return {'CANCELLED'}
+        self.add_vertex_group(context)
+        self.add_bevel_modifier(context)
+
+        context.window_manager.modal_handler_add(self)
+        return {'RUNNING_MODAL'}
 
 
     def add_vertex_group(self, context):