Merge pull request #1206 from Balint-H:feature/mouse-spring-qol

PiperOrigin-RevId: 590741248
Change-Id: I0228a2e6a27fd869f34f3e96163cd29e5ac99f35

unity/Editor/Components/MjMouseSpring.cs

@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#if UNITY_EDITOR
 using System;
 using UnityEditor;
 using UnityEngine;
@@ -23,7 +22,7 @@ namespace Mujoco {
   // and left mouse drag will apply a force on the body.  Holding shift down will change the applied
   // force direction between World XZ plane and Y[camera-up].
 
-  [CustomEditor(typeof(MjBody))]
+  [CustomEditor(typeof(MjComponent), true)]
   public class MjMouseSpring : Editor {
     private bool _lastShiftKeyState = false;
 
@@ -102,7 +101,11 @@ namespace Mujoco {
         return;
       }
 
-      MjBody body = target as MjBody;
+      var targetObject = target as MjComponent;
+      MjBody body = targetObject.GetComponentInParent<MjBody>();
+      if(!body)
+        return;
+
       Vector3 bodyPosition =
           body != null ? body.transform.position : Vector3.zero;
       var scene = MjScene.Instance;
@@ -188,4 +191,3 @@ namespace Mujoco {
     }
   }
 }
-#endif

unity/Editor/Components/MjShapeComponentEditor.cs

@@ -22,7 +22,7 @@ namespace Mujoco {
 
 [CustomEditor(typeof(MjShapeComponent), true)]
 [CanEditMultipleObjects]
-public class MjShapeComponentEditor : Editor {
+public class MjShapeComponentEditor : MjMouseSpring {
 
   public override void OnInspectorGUI() {