Parse

This module contains functions that provide a way to traverse the assembly structure, extract information about parts, subassemblies, instances, and mates, and generate a hierarchical representation of the assembly.

build_rigid_subassembly_occurrence_map(rigid_subassemblies, id_to_name_map, parts) async

Asynchronously build a map of rigid subassembly occurrences.

Source code in onshape_robotics_toolkit\parse.py

async def build_rigid_subassembly_occurrence_map(
    rigid_subassemblies: dict[str, RootAssembly], id_to_name_map: dict[str, str], parts: dict[str, Part]
) -> dict[str, dict[str, Occurrence]]:
    """
    Asynchronously build a map of rigid subassembly occurrences.
    """
    occurrence_map: dict[str, dict[str, Occurrence]] = {}
    for assembly_key, rigid_subassembly in rigid_subassemblies.items():
        sub_occurrences: dict[str, Occurrence] = {}
        for occurrence in rigid_subassembly.occurrences:
            try:
                occurrence_path = [id_to_name_map[path] for path in occurrence.path]
                sub_occurrences[SUBASSEMBLY_JOINER.join(occurrence_path)] = occurrence
            except KeyError:
                LOGGER.warning(f"Occurrence path {occurrence.path} not found")

        # Populate parts data
        parts[assembly_key] = Part(
            isStandardContent=False,
            fullConfiguration=rigid_subassembly.fullConfiguration,
            configuration=rigid_subassembly.configuration,
            documentId=rigid_subassembly.documentId,
            elementId=rigid_subassembly.elementId,
            documentMicroversion=rigid_subassembly.documentMicroversion,
            documentVersion="",
            partId="",
            bodyType="",
            MassProperty=rigid_subassembly.MassProperty,
            isRigidAssembly=True,
            rigidAssemblyWorkspaceId=rigid_subassembly.documentMetaData.defaultWorkspace.id,
            rigidAssemblyToPartTF={},
        )
        occurrence_map[assembly_key] = sub_occurrences

    return occurrence_map

fetch_rigid_subassemblies_async(subassembly, key, client, rigid_subassembly_map) async

Fetch rigid subassemblies asynchronously.

Source code in onshape_robotics_toolkit\parse.py

async def fetch_rigid_subassemblies_async(
    subassembly: SubAssembly, key: str, client: Client, rigid_subassembly_map: dict[str, RootAssembly]
):
    """
    Fetch rigid subassemblies asynchronously.
    """
    try:
        rigid_subassembly_map[key] = await asyncio.to_thread(
            client.get_root_assembly,
            did=subassembly.documentId,
            wtype=WorkspaceType.M.value,
            wid=subassembly.documentMicroversion,
            eid=subassembly.elementId,
            with_mass_properties=True,
            log_response=False,
        )
    except Exception as e:
        LOGGER.error(f"Failed to fetch rigid subassembly for {key}: {e}")

get_instances(assembly, max_depth=0)

Optimized synchronous wrapper for get_instances.

Source code in onshape_robotics_toolkit\parse.py

def get_instances(
    assembly: Assembly, max_depth: int = 0
) -> tuple[dict[str, Union[PartInstance, AssemblyInstance]], dict[str, Occurrence], dict[str, str]]:
    """
    Optimized synchronous wrapper for `get_instances`.
    """
    instance_map: dict[str, Union[PartInstance, AssemblyInstance]] = {}
    id_to_name_map: dict[str, str] = {}
    asyncio.run(
        traverse_instances_async(
            assembly.rootAssembly,
            "",
            0,
            max_depth,
            assembly,
            id_to_name_map,
            instance_map,
        )
    )
    occurrence_map = get_occurrences(assembly, id_to_name_map, max_depth)
    return instance_map, occurrence_map, id_to_name_map

get_instances_sync(assembly, max_depth=0)

Get instances and their sanitized names from an Onshape assembly.

Parameters:

Name	Type	Description	Default
assembly	Assembly	The Onshape assembly object to use for extracting instances.	required
max_depth	int	Maximum depth to traverse in the assembly hierarchy. Default is 5	0

Returns:

Type	Description
dict[str, Union[PartInstance, AssemblyInstance]]	A tuple containing:
dict[str, Occurrence]	A dictionary mapping instance IDs to their corresponding instances.
dict[str, str]	A dictionary mapping instance IDs to their sanitized names.

Examples:

>>> assembly = Assembly(...)
>>> get_instances(assembly, max_depth=2)
(
    {
        "part1": PartInstance(...),
        "subassembly1": AssemblyInstance(...),
    },
    {
        "part1": "part1",
        "subassembly1": "subassembly1",
    }
)

Source code in onshape_robotics_toolkit\parse.py

def get_instances_sync(
    assembly: Assembly, max_depth: int = 0
) -> tuple[dict[str, Union[PartInstance, AssemblyInstance]], dict[str, Occurrence], dict[str, str]]:
    """
    Get instances and their sanitized names from an Onshape assembly.

    Args:
        assembly: The Onshape assembly object to use for extracting instances.
        max_depth: Maximum depth to traverse in the assembly hierarchy. Default is 5

    Returns:
        A tuple containing:
        - A dictionary mapping instance IDs to their corresponding instances.
        - A dictionary mapping instance IDs to their sanitized names.

    Examples:
        >>> assembly = Assembly(...)
        >>> get_instances(assembly, max_depth=2)
        (
            {
                "part1": PartInstance(...),
                "subassembly1": AssemblyInstance(...),
            },
            {
                "part1": "part1",
                "subassembly1": "subassembly1",
            }
        )
    """

    def traverse_instances(
        root: Union[RootAssembly, SubAssembly], prefix: str = "", current_depth: int = 0
    ) -> tuple[dict[str, Union[PartInstance, AssemblyInstance]], dict[str, str]]:
        """
        Traverse the assembly structure to get instances.

        Args:
            root: Root assembly or subassembly object to traverse.
            prefix: Prefix for the instance ID.
            current_depth: Current depth in the assembly hierarchy.

        Returns:
            A tuple containing:
            - A dictionary mapping instance IDs to their corresponding instances.
            - A dictionary mapping instance IDs to their sanitized names.
        """
        instance_map = {}
        id_to_name_map = {}

        # Stop traversing if the maximum depth is reached
        if current_depth >= max_depth:
            LOGGER.debug(f"Max depth {max_depth} reached. Stopping traversal at depth {current_depth}.")
            return instance_map, id_to_name_map

        for instance in root.instances:
            sanitized_name = get_sanitized_name(instance.name)
            LOGGER.debug(f"Parsing instance: {sanitized_name}")
            instance_id = f"{prefix}{SUBASSEMBLY_JOINER}{sanitized_name}" if prefix else sanitized_name
            id_to_name_map[instance.id] = sanitized_name
            instance_map[instance_id] = instance

            # Recursively process sub-assemblies if applicable
            if instance.type == InstanceType.ASSEMBLY:
                for sub_assembly in assembly.subAssemblies:
                    if sub_assembly.uid == instance.uid:
                        sub_instance_map, sub_id_to_name_map = traverse_instances(
                            sub_assembly, instance_id, current_depth + 1
                        )
                        instance_map.update(sub_instance_map)
                        id_to_name_map.update(sub_id_to_name_map)

        return instance_map, id_to_name_map

    instance_map, id_to_name_map = traverse_instances(assembly.rootAssembly)
    # return occurrences internally as it relies on max_depth
    occurrence_map = get_occurrences(assembly, id_to_name_map, max_depth)

    return instance_map, occurrence_map, id_to_name_map

get_mates_and_relations(assembly, subassemblies, rigid_subassemblies, id_to_name_map, parts)

Synchronous wrapper for get_mates_and_relations_async.

Source code in onshape_robotics_toolkit\parse.py

def get_mates_and_relations(
    assembly: Assembly,
    subassemblies: dict[str, SubAssembly],
    rigid_subassemblies: dict[str, RootAssembly],
    id_to_name_map: dict[str, str],
    parts: dict[str, Part],
) -> tuple[dict[str, MateFeatureData], dict[str, MateRelationFeatureData]]:
    """
    Synchronous wrapper for `get_mates_and_relations_async`.
    """
    return asyncio.run(
        get_mates_and_relations_async(assembly, subassemblies, rigid_subassemblies, id_to_name_map, parts)
    )

get_mates_and_relations_async(assembly, subassemblies, rigid_subassemblies, id_to_name_map, parts) async

Asynchronously get mates and relations.

Source code in onshape_robotics_toolkit\parse.py

async def get_mates_and_relations_async(
    assembly: Assembly,
    subassemblies: dict[str, SubAssembly],
    rigid_subassemblies: dict[str, RootAssembly],
    id_to_name_map: dict[str, str],
    parts: dict[str, Part],
) -> tuple[dict[str, MateFeatureData], dict[str, MateRelationFeatureData]]:
    """
    Asynchronously get mates and relations.
    """
    rigid_subassembly_occurrence_map = await build_rigid_subassembly_occurrence_map(
        rigid_subassemblies, id_to_name_map, parts
    )

    mates_map, relations_map = await process_features_async(
        assembly.rootAssembly.features,
        parts,
        id_to_name_map,
        rigid_subassembly_occurrence_map,
        rigid_subassemblies,
        None,
    )

    for key, subassembly in subassemblies.items():
        sub_mates, sub_relations = await process_features_async(
            subassembly.features, parts, id_to_name_map, rigid_subassembly_occurrence_map, rigid_subassemblies, key
        )
        mates_map.update(sub_mates)
        relations_map.update(sub_relations)

    return mates_map, relations_map

get_occurrence_name(occurrences, subassembly_prefix=None)

Get the mapping name for an occurrence path.

Parameters:

Name	Type	Description	Default
occurrences	list[str]	Occurrence path.	required
subassembly_prefix	Optional[str]	Prefix for the subassembly.	None

Returns:

Type	Description
str	The mapping name.

Examples:

>>> get_occurrence_name(["subassembly1", "part1"], "subassembly1")
"subassembly1-SUB-part1"

>>> get_occurrence_name(["part1"], "subassembly1")
"subassembly1-SUB-part1"

Source code in onshape_robotics_toolkit\parse.py

def get_occurrence_name(occurrences: list[str], subassembly_prefix: Optional[str] = None) -> str:
    """
    Get the mapping name for an occurrence path.

    Args:
        occurrences: Occurrence path.
        subassembly_prefix: Prefix for the subassembly.

    Returns:
        The mapping name.

    Examples:
        >>> get_occurrence_name(["subassembly1", "part1"], "subassembly1")
        "subassembly1-SUB-part1"

        >>> get_occurrence_name(["part1"], "subassembly1")
        "subassembly1-SUB-part1"
    """
    prefix = f"{subassembly_prefix}{SUBASSEMBLY_JOINER}" if subassembly_prefix else ""
    return f"{prefix}{SUBASSEMBLY_JOINER.join(occurrences)}"

get_occurrences(assembly, id_to_name_map, max_depth=0)

Optimized occurrences fetching using comprehensions.

Source code in onshape_robotics_toolkit\parse.py

def get_occurrences(assembly: Assembly, id_to_name_map: dict[str, str], max_depth: int = 0) -> dict[str, Occurrence]:
    """
    Optimized occurrences fetching using comprehensions.
    """
    return {
        SUBASSEMBLY_JOINER.join([
            id_to_name_map[path] for path in occurrence.path if path in id_to_name_map
        ]): occurrence
        for occurrence in assembly.rootAssembly.occurrences
        if len(occurrence.path) <= max_depth + 1
    }

get_parts(assembly, rigid_subassemblies, client, instances)

Get parts of an Onshape assembly.

Parameters:

Name	Type	Description	Default
assembly	Assembly	The Onshape assembly object to use for extracting parts.	required
rigid_subassemblies	dict[str, RootAssembly]	Mapping of instance IDs to rigid subassemblies.	required
client	Client	The Onshape client object to use for sending API requests.	required
instances	dict[str, Union[PartInstance, AssemblyInstance]]	Mapping of instance IDs to their corresponding instances.	required

Returns:

Type	Description
dict[str, Part]	A dictionary mapping part IDs to their corresponding part objects.

Source code in onshape_robotics_toolkit\parse.py

def get_parts(
    assembly: Assembly,
    rigid_subassemblies: dict[str, RootAssembly],
    client: Client,
    instances: dict[str, Union[PartInstance, AssemblyInstance]],
) -> dict[str, Part]:
    """
    Get parts of an Onshape assembly.

    Args:
        assembly: The Onshape assembly object to use for extracting parts.
        rigid_subassemblies: Mapping of instance IDs to rigid subassemblies.
        client: The Onshape client object to use for sending API requests.
        instances: Mapping of instance IDs to their corresponding instances.

    Returns:
        A dictionary mapping part IDs to their corresponding part objects.
    """
    return asyncio.run(_get_parts_async(assembly, rigid_subassemblies, client, instances))

get_subassemblies(assembly, client, instances)

Synchronous wrapper for get_subassemblies_async.

Source code in onshape_robotics_toolkit\parse.py

def get_subassemblies(
    assembly: Assembly,
    client: Client,
    instances: dict[str, Union[PartInstance, AssemblyInstance]],
) -> tuple[dict[str, SubAssembly], dict[str, RootAssembly]]:
    """
    Synchronous wrapper for `get_subassemblies_async`.
    """
    return asyncio.run(get_subassemblies_async(assembly, client, instances))

get_subassemblies_async(assembly, client, instance_map) async

Asynchronously fetch subassemblies.

Source code in onshape_robotics_toolkit\parse.py

async def get_subassemblies_async(
    assembly: Assembly,
    client: Client,
    instance_map: dict[str, Union[PartInstance, AssemblyInstance]],
) -> tuple[dict[str, SubAssembly], dict[str, RootAssembly]]:
    """
    Asynchronously fetch subassemblies.
    """
    subassembly_map: dict[str, SubAssembly] = {}
    rigid_subassembly_map: dict[str, RootAssembly] = {}

    # Group by UID
    subassembly_instance_map = {}
    rigid_subassembly_instance_map = {}

    for instance_key, instance in instance_map.items():
        if instance.type == InstanceType.ASSEMBLY:
            if instance.isRigid:
                rigid_subassembly_instance_map.setdefault(instance.uid, []).append(instance_key)
            else:
                subassembly_instance_map.setdefault(instance.uid, []).append(instance_key)

    # Process subassemblies concurrently
    tasks = []
    for subassembly in assembly.subAssemblies:
        uid = subassembly.uid
        if uid in subassembly_instance_map:
            is_rigid = len(subassembly.features) == 0 or all(
                feature.featureType == AssemblyFeatureType.MATEGROUP for feature in subassembly.features
            )
            for key in subassembly_instance_map[uid]:
                if is_rigid:
                    tasks.append(fetch_rigid_subassemblies_async(subassembly, key, client, rigid_subassembly_map))
                else:
                    subassembly_map[key] = subassembly

        elif uid in rigid_subassembly_instance_map:
            for key in rigid_subassembly_instance_map[uid]:
                tasks.append(fetch_rigid_subassemblies_async(subassembly, key, client, rigid_subassembly_map))

    await asyncio.gather(*tasks)
    return subassembly_map, rigid_subassembly_map

join_mate_occurrences(parent, child, prefix=None)

Join two occurrence paths with a mate joiner.

Parameters:

Name	Type	Description	Default
parent	list[str]	Occurrence path of the parent entity.	required
child	list[str]	Occurrence path of the child entity.	required
prefix	Optional[str]	Prefix to add to the occurrence path.	None

Returns:

Type	Description
str	The joined occurrence path.

Examples:

>>> join_mate_occurrences(["subassembly1", "part1"], ["subassembly2"])
"subassembly1-SUB-part1-MATE-subassembly2"

>>> join_mate_occurrences(["part1"], ["part2"])
"part1-MATE-part2"

Source code in onshape_robotics_toolkit\parse.py

def join_mate_occurrences(parent: list[str], child: list[str], prefix: Optional[str] = None) -> str:
    """
    Join two occurrence paths with a mate joiner.

    Args:
        parent: Occurrence path of the parent entity.
        child: Occurrence path of the child entity.
        prefix: Prefix to add to the occurrence path.

    Returns:
        The joined occurrence path.

    Examples:
        >>> join_mate_occurrences(["subassembly1", "part1"], ["subassembly2"])
        "subassembly1-SUB-part1-MATE-subassembly2"

        >>> join_mate_occurrences(["part1"], ["part2"])
        "part1-MATE-part2"
    """
    parent_occurrence = get_occurrence_name(parent, prefix)
    child_occurrence = get_occurrence_name(child, prefix)
    return f"{parent_occurrence}{MATE_JOINER}{child_occurrence}"

process_features_async(features, parts, id_to_name_map, rigid_subassembly_occurrence_map, rigid_subassemblies, subassembly_prefix) async

Process assembly features asynchronously.

Source code in onshape_robotics_toolkit\parse.py

async def process_features_async(  # noqa: C901
    features: list[AssemblyFeature],
    parts: dict[str, Part],
    id_to_name_map: dict[str, str],
    rigid_subassembly_occurrence_map: dict[str, dict[str, Occurrence]],
    rigid_subassemblies: dict[str, RootAssembly],
    subassembly_prefix: Optional[str],
) -> tuple[dict[str, MateFeatureData], dict[str, MateRelationFeatureData]]:
    """
    Process assembly features asynchronously.
    """
    mates_map: dict[str, MateFeatureData] = {}
    relations_map: dict[str, MateRelationFeatureData] = {}

    for feature in features:
        feature.featureData.id = feature.id

        if feature.suppressed:
            continue

        if feature.featureType == AssemblyFeatureType.MATE:
            if len(feature.featureData.matedEntities) < 2:
                LOGGER.warning(f"Invalid mate feature: {feature}")
                continue

            try:
                child_occurrences = [
                    id_to_name_map[path] for path in feature.featureData.matedEntities[CHILD].matedOccurrence
                ]
                parent_occurrences = [
                    id_to_name_map[path] for path in feature.featureData.matedEntities[PARENT].matedOccurrence
                ]
            except KeyError as e:
                LOGGER.warning(e)
                LOGGER.warning(f"Key not found in {id_to_name_map.keys()}")
                continue

            # Handle rigid subassemblies
            if parent_occurrences[0] in rigid_subassemblies:
                _occurrence = rigid_subassembly_occurrence_map[parent_occurrences[0]].get(parent_occurrences[1])
                if _occurrence:
                    parent_parentCS = MatedCS.from_tf(np.matrix(_occurrence.transform).reshape(4, 4))
                    parts[parent_occurrences[0]].rigidAssemblyToPartTF[parent_occurrences[1]] = parent_parentCS.part_tf
                    feature.featureData.matedEntities[PARENT].parentCS = parent_parentCS
                parent_occurrences = [parent_occurrences[0]]

            if child_occurrences[0] in rigid_subassemblies:
                _occurrence = rigid_subassembly_occurrence_map[child_occurrences[0]].get(child_occurrences[1])
                if _occurrence:
                    child_parentCS = MatedCS.from_tf(np.matrix(_occurrence.transform).reshape(4, 4))
                    parts[child_occurrences[0]].rigidAssemblyToPartTF[child_occurrences[1]] = child_parentCS.part_tf
                    feature.featureData.matedEntities[CHILD].parentCS = child_parentCS
                child_occurrences = [child_occurrences[0]]

            mates_map[
                join_mate_occurrences(
                    parent=parent_occurrences,
                    child=child_occurrences,
                    prefix=subassembly_prefix,
                )
            ] = feature.featureData

        elif feature.featureType == AssemblyFeatureType.MATERELATION:
            if feature.featureData.relationType == RelationType.SCREW:
                child_joint_id = feature.featureData.mates[0].featureId
            else:
                child_joint_id = feature.featureData.mates[RELATION_CHILD].featureId

            relations_map[child_joint_id] = feature.featureData

    return mates_map, relations_map

traverse_instances_async(root, prefix, current_depth, max_depth, assembly, id_to_name_map, instance_map) async

Asynchronously traverse the assembly structure to get instances.

Source code in onshape_robotics_toolkit\parse.py

async def traverse_instances_async(
    root: Union[RootAssembly, SubAssembly],
    prefix: str,
    current_depth: int,
    max_depth: int,
    assembly: Assembly,
    id_to_name_map: dict[str, str],
    instance_map: dict[str, Union[PartInstance, AssemblyInstance]],
) -> None:
    """
    Asynchronously traverse the assembly structure to get instances.
    """
    isRigid = False
    if current_depth >= max_depth:
        LOGGER.debug(
            f"Max depth {max_depth} reached. Assuming all sub-assemblies to be rigid at depth {current_depth}."
        )
        isRigid = True

    for instance in root.instances:
        sanitized_name = get_sanitized_name(instance.name)
        LOGGER.debug(f"Parsing instance: {sanitized_name}")
        instance_id = f"{prefix}{SUBASSEMBLY_JOINER}{sanitized_name}" if prefix else sanitized_name
        id_to_name_map[instance.id] = sanitized_name
        instance_map[instance_id] = instance

        if instance.type == InstanceType.ASSEMBLY:
            instance_map[instance_id].isRigid = isRigid

        # Handle subassemblies concurrently
        if instance.type == InstanceType.ASSEMBLY:
            tasks = [
                traverse_instances_async(
                    sub_assembly, instance_id, current_depth + 1, max_depth, assembly, id_to_name_map, instance_map
                )
                for sub_assembly in assembly.subAssemblies
                if sub_assembly.uid == instance.uid
            ]
            await asyncio.gather(*tasks)