From 7062c6f6772909ec573deb43481b588acaa856b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= <bkh16@ic.ac.uk>
Date: Fri, 17 Nov 2023 14:43:10 +0000
Subject: [PATCH] Enable mouse spring perturbations when clicking on a geom in
 scene view

---
 unity/Editor/Components/MjMouseSpring.cs          | 8 ++++++--
 unity/Editor/Components/MjShapeComponentEditor.cs | 2 +-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/unity/Editor/Components/MjMouseSpring.cs b/unity/Editor/Components/MjMouseSpring.cs
index b2264727..c502e551 100644
--- a/unity/Editor/Components/MjMouseSpring.cs
+++ b/unity/Editor/Components/MjMouseSpring.cs
@@ -23,7 +23,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 +102,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;
diff --git a/unity/Editor/Components/MjShapeComponentEditor.cs b/unity/Editor/Components/MjShapeComponentEditor.cs
index 41b8ba22..3d0d687c 100644
--- a/unity/Editor/Components/MjShapeComponentEditor.cs
+++ b/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() {