issue #22: [fix] avoid TypeError by branching the processing based on the variable's type

Author: siruku6

pybullet_tools/pr2_primitives.py

@@ -5,6 +5,7 @@
 import random
 import time
 from itertools import islice, count
+from typing import Tuple, Optional, Union
 
 import numpy as np
 
@@ -193,15 +194,23 @@ def create_trajectory(robot, joints, path):
 ##################################################
 
 class GripperCommand(Command):
-    def __init__(self, robot, arm, position, teleport=False):
+    def __init__(self, robot, arm, position: Union[float, Tuple[float, ...]], teleport=False):
         self.robot = robot
         self.arm = arm
-        self.position = position
+        self.position: Union[float, Tuple[float, ...]] = position
         self.teleport = teleport
     def apply(self, state, **kwargs):
         joints = get_gripper_joints(self.robot, self.arm)
         start_conf = get_joint_positions(self.robot, joints)
-        end_conf = [self.position] * len(joints)
+
+        # NOTE: position can be a single float (applied to all joints) or an array-like object of floats (one per joint)
+        if isinstance(self.position, float):
+            end_conf = [self.position] * len(joints)
+        elif isinstance(self.position, (list, tuple, np.ndarray)) and len(self.position) == len(joints):
+            end_conf = self.position
+        else:
+            raise ValueError(f'Invalid gripper position: {self.position}')
+
         if self.teleport:
             path = [start_conf, end_conf]
         else:
@@ -328,7 +337,7 @@ def fn(body):
                           for g in get_side_grasps(body, grasp_length=GRASP_LENGTH))
         filtered_grasps = []
         for grasp in grasps:
-            grasp_width = compute_grasp_width(problem.robot, arm, body, grasp.value) if collisions else 0.0
+            grasp_width: Optional[Union[float, Tuple[float, ...]]] = compute_grasp_width(problem.robot, arm, body, grasp.value) if collisions else 0.0
             if grasp_width is not None:
                 grasp.grasp_width = grasp_width
                 filtered_grasps.append(grasp)

pybullet_tools/pr2_utils.py

@@ -4,6 +4,7 @@
 import re
 from collections import namedtuple
 from itertools import combinations
+from typing import List, Optional, Tuple
 
 import numpy as np
 
@@ -17,7 +18,7 @@
     movable_from_joints, quat_from_axis_angle, LockRenderer, Euler, get_links, get_link_name, \
     get_extend_fn, get_moving_links, link_pairs_collision, get_link_subtree, \
     clone_body, get_all_links, pairwise_collision, tform_point, get_camera_matrix, ray_from_pixel, pixel_from_ray, dimensions_from_camera_matrix, \
-    wrap_angle, TRANSPARENT, PI, OOBB, pixel_from_point, set_all_color, wait_if_gui
+    wrap_angle, TRANSPARENT, PI, OOBB, pixel_from_point, set_all_color, wait_if_gui, POSE
 
 # TODO: restrict number of pr2 rotations to prevent from wrapping too many times
 
@@ -291,7 +292,7 @@ def close_arm(robot, arm):
 SIDE_HEIGHT_OFFSET = 0.03 # z distance from top of object
 
 def get_top_grasps(body, under=False, tool_pose=TOOL_POSE, body_pose=unit_pose(),
-                   max_width=MAX_GRASP_WIDTH, grasp_length=GRASP_LENGTH):
+                   max_width=MAX_GRASP_WIDTH, grasp_length=GRASP_LENGTH) -> List[POSE]:
     # TODO: rename the box grasps
     center, (w, l, h) = approximate_as_prism(body, body_pose=body_pose)
     reflect_z = Pose(euler=[0, math.pi, 0])
@@ -757,7 +758,9 @@ def get_base_extend_fn(robot):
 
 #####################################
 
-def close_until_collision(robot, gripper_joints, bodies=[], open_conf=None, closed_conf=None, num_steps=25, **kwargs):
+def close_until_collision(
+    robot, gripper_joints, bodies=[], open_conf: List[float] = None, closed_conf: List[float] = None, num_steps=25, **kwargs
+) -> Optional[Tuple[float, ...]]:
     if not gripper_joints:
         return None
     if open_conf is None:
@@ -766,7 +769,7 @@ def close_until_collision(robot, gripper_joints, bodies=[], open_conf=None, clos
         closed_conf = [get_min_limit(robot, joint) for joint in gripper_joints]
     resolutions = np.abs(np.array(open_conf) - np.array(closed_conf)) / num_steps
     extend_fn = get_extend_fn(robot, gripper_joints, resolutions=resolutions)
-    close_path = [open_conf] + list(extend_fn(open_conf, closed_conf))
+    close_path: List[Tuple[float, ...]] = [open_conf] + list(extend_fn(open_conf, closed_conf))
     collision_links = frozenset(get_moving_links(robot, gripper_joints))
 
     for i, conf in enumerate(close_path):
@@ -778,7 +781,7 @@ def close_until_collision(robot, gripper_joints, bodies=[], open_conf=None, clos
     return close_path[-1]
     #return None # False
 
-def compute_grasp_width(robot, arm, body, grasp_pose, **kwargs):
+def compute_grasp_width(robot, arm, body, grasp_pose, **kwargs) -> Optional[Tuple[float, ...]]:
     tool_link = get_gripper_link(robot, arm)
     tool_pose = get_link_pose(robot, tool_link)
     body_pose = multiply(tool_pose, grasp_pose)

pybullet_tools/utils.py

@@ -26,10 +26,15 @@
 from itertools import product, combinations, count, cycle, islice
 from multiprocessing import TimeoutError
 from contextlib import contextmanager
+from typing import List, Tuple
 
 from pybullet_utils.transformations import quaternion_from_matrix, unit_vector, euler_from_quaternion, quaternion_slerp, \
     random_quaternion, quaternion_about_axis
 
+VEC3 = Tuple[float, float, float]
+QUATERNION = Tuple[float, float, float, float]
+POSE = Tuple[VEC3, QUATERNION]
+
 def join_paths(*paths):
     return os.path.abspath(os.path.join(*paths))
 
@@ -118,8 +123,9 @@ def get_parent_dir(file): # __file__
 
 inf_generator = count # count | lambda: iter(int, 1)
 
-List = lambda *args: list(args)
-Tuple = lambda *args: tuple(args)
+# Disable these lines because they prevent developers from using type hints
+# List = lambda *args: list(args)
+# Tuple = lambda *args: tuple(args)
 
 def empty_sequence():
     return iter([])
@@ -1619,26 +1625,26 @@ def Point(x=0., y=0., z=0.):
 def Euler(roll=0., pitch=0., yaw=0.):
     return np.array([roll, pitch, yaw])
 
-def Pose(point=None, euler=None):
+def Pose(point=None, euler=None) -> POSE:
     point = Point() if point is None else point
     euler = Euler() if euler is None else euler
     return point, quat_from_euler(euler)
 
 def Pose2d(x=0., y=0., yaw=0.):
     return np.array([x, y, yaw])
 
-def invert(pose):
+def invert(pose: POSE) -> POSE:
     point, quat = pose
     return p.invertTransform(point, quat)
 
-def multiply(*poses):
-    pose = poses[0]
+def multiply(*poses) -> POSE:
+    pose: POSE = poses[0]
     for next_pose in poses[1:]:
         pose = p.multiplyTransforms(pose[0], pose[1], *next_pose)
     return pose
 
-def invert_quat(quat):
-    pose = (unit_point(), quat)
+def invert_quat(quat: QUATERNION) -> QUATERNION:
+    pose: POSE = (unit_point(), quat)
     return quat_from_pose(invert(pose))
 
 def multiply_quats(*quats):
@@ -2162,17 +2168,17 @@ def get_joint_limits(body, joint):
 
 get_joint_interval = get_joint_limits # TODO: get box limits?
 
-def get_min_limit(body, joint):
+def get_min_limit(body, joint) -> float:
     # TODO: rename to min_position
     return get_joint_limits(body, joint)[0]
 
-def get_min_limits(body, joints):
+def get_min_limits(body, joints) -> List[float]:
     return [get_min_limit(body, joint) for joint in joints]
 
-def get_max_limit(body, joint):
+def get_max_limit(body, joint) -> float:
     return get_joint_limits(body, joint)[1]
 
-def get_max_limits(body, joints):
+def get_max_limits(body, joints) -> List[float]:
     return [get_max_limit(body, joint) for joint in joints]
 
 def get_joint_intervals(body, joints):
@@ -4626,7 +4632,7 @@ def end_effector_from_body(body_pose, grasp_pose):
 def approach_from_grasp(approach_pose, end_effector_pose):
     return multiply(approach_pose, end_effector_pose)
 
-def get_grasp_pose(constraint):
+def get_grasp_pose(constraint) -> POSE:
     """
     Grasps are parent_from_child
     """