Use defined object class in get_objects_in_tree (#102)

Author: wkentaro

pyrep/objects/cartesian_path.py

@@ -1,6 +1,6 @@
 from typing import Tuple, List
 from pyrep.backend import sim
-from pyrep.objects.object import Object
+from pyrep.objects.object import Object, object_type_to_class
 from pyrep.const import ObjectType
 
 
@@ -90,3 +90,6 @@ def _script_call(self, func: str, ints=(), floats=(), strings=(), bytes=''):
         return sim.simExtCallScriptFunction(
             func, sim.sim_scripttype_addonscript,
             list(ints), list(floats), list(strings), bytes)
+
+
+object_type_to_class[ObjectType.PATH] = CartesianPath

pyrep/objects/dummy.py

@@ -1,4 +1,4 @@
-from pyrep.objects.object import Object
+from pyrep.objects.object import Object, object_type_to_class
 from pyrep.const import ObjectType
 from pyrep.backend import sim
 
@@ -21,3 +21,6 @@ def create(size=0.01) -> 'Dummy':
 
     def _get_requested_type(self) -> ObjectType:
         return ObjectType.DUMMY
+
+
+object_type_to_class[ObjectType.DUMMY] = Dummy

pyrep/objects/force_sensor.py

@@ -1,6 +1,6 @@
 from typing import Tuple, List
 from pyrep.backend import sim
-from pyrep.objects.object import Object
+from pyrep.objects.object import Object, object_type_to_class
 from pyrep.const import ObjectType
 
 
@@ -20,3 +20,6 @@ def read(self) -> Tuple[List[float], List[float]]:
         """
         _, forces, torques = sim.simReadForceSensor(self._handle)
         return forces, torques
+
+
+object_type_to_class[ObjectType.FORCE_SENSOR] = ForceSensor

pyrep/objects/joint.py

@@ -1,7 +1,7 @@
 from typing import Tuple, List, Union
 from pyrep.backend import sim
 from pyrep.const import JointType, JointMode
-from pyrep.objects.object import Object
+from pyrep.objects.object import Object, object_type_to_class
 from pyrep.const import ObjectType
 
 
@@ -246,3 +246,6 @@ def set_joint_mode(self, value: JointMode) -> None:
         :param value: The new joint mode value.
         """
         sim.simSetJointMode(self._handle, value.value)
+
+
+object_type_to_class[ObjectType.JOINT] = Joint

pyrep/objects/object.py

@@ -6,6 +6,9 @@
 import numpy as np
 
 
+object_type_to_class = {}
+
+
 class Object(object):
     """Base class for V-REP scene objects that are used for building a scene.
 
@@ -598,7 +601,9 @@ def get_objects_in_tree(self, object_type=ObjectType.ALL, exclude_base=True,
             self._handle, object_type.value, options)
         objects = []
         for h in handles:
-            objects.append(Object(h))
+            object_type = ObjectType(sim.simGetObjectType(h))
+            cls = object_type_to_class.get(object_type, Object)
+            objects.append(cls(h))
         return objects
 
     def copy(self) -> 'Object':

pyrep/objects/proximity_sensor.py

@@ -1,7 +1,7 @@
 from math import sqrt
 
 from pyrep.backend import sim
-from pyrep.objects.object import Object
+from pyrep.objects.object import Object, object_type_to_class
 from pyrep.const import ObjectType
 
 
@@ -36,3 +36,6 @@ def is_detected(self, obj: Object) -> bool:
         state, point = sim.simCheckProximitySensor(
             self._handle, obj.get_handle())
         return state == 1
+
+
+object_type_to_class[ObjectType.PROXIMITY_SENSOR] = ProximitySensor

pyrep/objects/shape.py

@@ -1,6 +1,6 @@
 from typing import List, Tuple
 from pyrep.backend import sim
-from pyrep.objects.object import Object
+from pyrep.objects.object import Object, object_type_to_class
 from pyrep.const import ObjectType, PrimitiveShape, TextureMappingMode
 from pyrep.textures.texture import Texture
 import os
@@ -356,3 +356,6 @@ def ungroup(self) -> List['Shape']:
         """
         handles = sim.simUngroupShape(self.get_handle())
         return [Shape(handle) for handle in handles]
+
+
+object_type_to_class[ObjectType.SHAPE] = Shape

pyrep/objects/vision_sensor.py

@@ -1,7 +1,7 @@
 import math
 from typing import List, Union
 from pyrep.backend import sim
-from pyrep.objects.object import Object
+from pyrep.objects.object import Object, object_type_to_class
 import numpy as np
 from pyrep.const import ObjectType, PerspectiveMode, RenderMode
 
@@ -265,3 +265,6 @@ def set_far_clipping_plane(self, far_clipping: float) -> None:
         sim.simSetObjectFloatParameter(
             self._handle, sim.sim_visionfloatparam_far_clipping, far_clipping
         )
+
+
+object_type_to_class[ObjectType.VISION_SENSOR] = VisionSensor

tests/test_objects.py

@@ -142,8 +142,12 @@ def test_get_bounding_box(self):
 
     def test_get_objects_in_tree(self):
         dummys = [Dummy('nested_dummy%d' % i) for i in range(3)]
-        self.assertListEqual(dummys[0].get_objects_in_tree(
-            exclude_base=False, first_generation_only=False), dummys)
+
+        objects = dummys[0].get_objects_in_tree(
+            exclude_base=False, first_generation_only=False)
+        self.assertListEqual(objects, dummys)
+        for obj in objects:
+            self.assertIs(type(obj), Dummy)
 
         self.assertListEqual(
             dummys[0].get_objects_in_tree(