Rigify: cleanup PyCharm warning highlights in base_rig.py

Also declare stage names for auto-completion.

3 files changed, 29 insertions, 18 deletions

diff --git a/rigify/base_rig.py b/rigify/base_rig.py
index a1f9e952..41430996 100644
--- a/rigify/base_rig.py
+++ b/rigify/base_rig.py
@@ -1,18 +1,17 @@
 # SPDX-License-Identifier: GPL-2.0-or-later
 
 import collections
+import typing
 
 from .utils.errors import RaiseErrorMixin
 from .utils.bones import BoneDict, BoneUtilityMixin
 from .utils.mechanism import MechanismUtilityMixin
 from .utils.metaclass import BaseStagedClass
 
-# Only export certain symbols via 'from base_rig import *'
-__all__ = ['BaseRig', 'stage']
 
-#=============================================
+##############################################
 # Base Rig
-#=============================================
+##############################################
 
 class GenerateCallbackHost(BaseStagedClass, define_stages=True):
     """
@@ -162,13 +161,13 @@ class BaseRig(GenerateCallbackHost, RaiseErrorMixin, BoneUtilityMixin, Mechanism
         # Collection of bone names for use in implementing the rig
         self.bones = BoneDict(
             # ORG bone names
-            org = self.find_org_bones(pose_bone),
+            org=self.find_org_bones(pose_bone),
             # Control bone names
-            ctrl = BoneDict(),
+            ctrl=BoneDict(),
             # MCH bone names
-            mch = BoneDict(),
+            mch=BoneDict(),
             # DEF bone names
-            deform = BoneDict(),
+            deform=BoneDict(),
         )
 
         # Data useful for complex rig interaction:
@@ -234,9 +233,9 @@ class BaseRig(GenerateCallbackHost, RaiseErrorMixin, BoneUtilityMixin, Mechanism
         """
 
 
-#=============================================
+##############################################
 # Rig Utility
-#=============================================
+##############################################
 
 
 class RigUtility(BoneUtilityMixin, MechanismUtilityMixin):
@@ -270,13 +269,21 @@ class RigComponent(LazyRigComponent):
         self.enable_component()
 
 
-#=============================================
+##############################################
 # Rig Stage Decorators
-#=============================================
+##############################################
 
+# Generate @stage.<...> decorators for all valid stages.
+@GenerateCallbackHost.stage_decorator_container
 class stage:
-    pass
-
-# Generate @stage.<...> decorators for all valid stages
-for name, decorator in GenerateCallbackHost.make_stage_decorators():
-    setattr(stage, name, decorator)
+    # Declare stages for auto-completion - doesn't affect execution.
+    initialize: typing.Callable
+    prepare_bones: typing.Callable
+    generate_bones: typing.Callable
+    parent_bones: typing.Callable
+    configure_bones: typing.Callable
+    preapply_bones: typing.Callable
+    apply_bones: typing.Callable
+    rig_bones: typing.Callable
+    generate_widgets: typing.Callable
+    finalize: typing.Callable
diff --git a/rigify/utils/metaclass.py b/rigify/utils/metaclass.py
index c8054d7f..1b10f385 100644
--- a/rigify/utils/metaclass.py
+++ b/rigify/utils/metaclass.py
@@ -114,6 +114,11 @@ class StagedMetaclass(type):
     def make_stage_decorators(self):
         return [(name, rigify_stage(name)) for name in self.rigify_stages]
 
+    def stage_decorator_container(self, cls):
+        for name, stage in self.make_stage_decorators():
+            setattr(cls, name, stage)
+        return cls
+
 
 class BaseStagedClass(object, metaclass=StagedMetaclass):
     rigify_sub_objects = tuple()
diff --git a/rigify/utils/switch_parent.py b/rigify/utils/switch_parent.py
index d4852c43..8d5fcc4e 100644
--- a/rigify/utils/switch_parent.py
+++ b/rigify/utils/switch_parent.py
@@ -14,7 +14,6 @@ from .mechanism import MechanismUtilityMixin
 from .rig import rig_is_child
 from .misc import map_list, map_apply, force_lazy
 
-from ..base_rig import *
 from ..base_generate import GeneratorPlugin
 
 from collections import defaultdict