diff --git a/onshape_api/urdf.py b/onshape_api/urdf.py index d2861ef..0afbfa2 100644 --- a/onshape_api/urdf.py +++ b/onshape_api/urdf.py @@ -44,7 +44,7 @@ VisualLink, ) from onshape_api.parse import MATE_JOINER, PARENT, RELATION_PARENT -from onshape_api.utilities.helpers import get_sanitized_name +from onshape_api.utilities.helpers import get_sanitized_name, make_unique_keys SCRIPT_DIR = os.path.dirname(__file__) @@ -486,7 +486,15 @@ def get_urdf_components( assets_map[child] = downloadable_link links.append(link) - links_map = {link.name: link for link in links} - joints_map = {joint.name: joint for joint in joints} + unique_joint_key_map = make_unique_keys([joint.name for joint in joints]) + unique_link_key_map = make_unique_keys([link.name for link in links]) + + for joint_key in unique_joint_key_map: + joints[unique_joint_key_map[joint_key]].name = joint_key + joints_map[joint_key] = joints[unique_joint_key_map[joint_key]] + + for link_key in unique_link_key_map: + links[unique_link_key_map[link_key]].name = link_key + links_map[link_key] = links[unique_link_key_map[link_key]] return links_map, joints_map, assets_map diff --git a/onshape_api/utilities/helpers.py b/onshape_api/utilities/helpers.py index a405c79..3a6bcea 100644 --- a/onshape_api/utilities/helpers.py +++ b/onshape_api/utilities/helpers.py @@ -204,6 +204,37 @@ def get_random_names(directory: str, count: int, filename: str = "words.txt") -> return random.sample(words, count) +def make_unique_keys(keys: list[str]) -> dict[str, int]: + """ + Make a list of keys unique by appending a number to duplicate keys and + return a mapping of unique keys to their original indices. + + Args: + keys: List of keys. + + Returns: + A dictionary mapping unique keys to their original indices. + + Examples: + >>> make_unique_keys(["a", "b", "a", "a"]) + {"a": 0, "b": 1, "a-1": 2, "a-2": 3} + """ + unique_key_map = {} + key_count = {} + + for index, key in enumerate(keys): + if key in key_count: + key_count[key] += 1 + unique_key = f"{key}-{key_count[key]}" + else: + key_count[key] = 0 + unique_key = key + + unique_key_map[unique_key] = index + + return unique_key_map + + def get_sanitized_name(name: str, replace_with: str = "-") -> str: """ Sanitize a name by removing special characters, preserving "-" and "_", and