Cosplay/nhf/utils.py

"""
Marking utilities for `Cq.Workplane`

Adds the functions to `Cq.Workplane`:
1. `tagPoint`
2. `tagPlane`
"""
import math
import functools
import cadquery as Cq
from nhf import Role
from typing import Union, Tuple, cast

# Bug fixes
def _subloc(self, name: str) -> Tuple[Cq.Location, str]:
    """
       Calculate relative location of an object in a subassembly.

       Returns the relative positions as well as the name of the top assembly.
       """

    rv = Cq.Location()
    obj = self.objects[name]
    name_out = name

    if obj not in self.children and obj is not self:
        locs = []
        while not obj.parent is self:
            locs.append(obj.loc)
            obj = cast(Cq.Assembly, obj.parent)
            name_out = obj.name

        rv = functools.reduce(lambda l1, l2: l2 * l1, locs)

    return (rv, name_out)
Cq.Assembly._subloc = _subloc


def tagPoint(self, tag: str):
    """
    Adds a vertex that can be used in `Point` constraints.
    """
    vertex = Cq.Vertex.makeVertex(0, 0, 0)
    self.eachpoint(vertex.moved, useLocalCoordinates=True).tag(tag)

Cq.Workplane.tagPoint = tagPoint

def tagPlane(self, tag: str,
             direction: Union[str, Cq.Vector, Tuple[float, float, float]] = '+Z'):
    """
    Adds a phantom `Cq.Edge` in the given location which can be referenced in a
    `Axis`, `Point`, or `Plane` constraint.
    """
    if isinstance(direction, str):
        x, y, z = 0, 0, 0
        assert len(direction) == 2
        sign, axis = direction
        if axis in ('z', 'Z'):
            z = 1
        elif axis in ('y', 'Y'):
            y = 1
        elif axis in ('x', 'X'):
            x = 1
        else:
            assert False, "Axis must be one of x,y,z"
        if sign == '+':
            sign = 1
        elif sign == '-':
            sign = -1
        else:
            assert False, "Sign must be one of +/-"
        v = Cq.Vector(x, y, z) * sign
    else:
        v = Cq.Vector(direction)
    edge = Cq.Edge.makeLine(v * (-1), v)
    return self.eachpoint(edge.located, useLocalCoordinates=True).tag(tag)

Cq.Workplane.tagPlane = tagPlane

def make_sphere(r: float = 2) -> Cq.Solid:
    """
    Makes a full sphere. The default function makes a hemisphere
    """
    return Cq.Solid.makeSphere(r, angleDegrees1=-90)
def make_arrow(size: float = 2) -> Cq.Workplane:
    cone = Cq.Solid.makeCone(
        radius1 = size,
        radius2 = 0,
        height=size)
    result = (
        Cq.Workplane("XY")
        .cylinder(radius=size / 2, height=size, centered=(True, True, False))
        .union(cone.located(Cq.Location((0, 0, size))))
    )
    result.faces("<Z").tag("dir_rev")
    return result

def to_marker_name(tag: str) -> str:
    return tag.replace("?", "__T").replace("/", "__Z") + "_marker"

def mark_point(self: Cq.Assembly,
               tag: str,
               size: float = 2,
               color: Cq.Color = Role.MARKER.color) -> Cq.Assembly:
    """
    Adds a marker to make a point visible
    """
    name = to_marker_name(tag)
    return (
        self
        .add(make_sphere(size), name=name, color=color)
        .constrain(tag, name, "Point")
    )

Cq.Assembly.markPoint = mark_point

def mark_plane(self: Cq.Assembly,
               tag: str,
               size: float = 2,
               color: Cq.Color = Role.MARKER.color) -> Cq.Assembly:
    """
    Adds a marker to make a plane visible
    """
    name = to_marker_name(tag)
    return (
        self
        .add(make_arrow(size), name=name, color=color)
        .constrain(tag, f"{name}?dir_rev", "Plane", param=180)
    )

Cq.Assembly.markPlane = mark_plane


def extrude_with_markers(sketch: Cq.Sketch,
                         thickness: float,
                         tags: list[Tuple[str, Tuple[float, float], float]],
                         reverse: bool = False):
    """
    Extrudes a sketch and place tags on the sketch for mating.

    Each tag is of the format `(name, (x, y), angle)`, where the angle is
    specifies in degrees counterclockwise from +X. Two marks are generated for
    each `name`, "{name}" for the location (with normal) and "{name}_dir" for
    the directrix specified by the angle.

    This simulates a process of laser cutting and bonding (for wood and acrylic)
    """
    result = (
        Cq.Workplane('XY')
        .placeSketch(sketch)
        .extrude(thickness)
    )
    plane = result.faces("<Z" if reverse else ">Z").workplane()
    sign = -1 if reverse else 1
    for tag, (px, py), angle in tags:
        theta = sign * math.radians(angle)
        direction = (math.cos(theta), math.sin(theta), 0)
        plane.moveTo(px, sign * py).tagPlane(tag)
        plane.moveTo(px, sign * py).tagPlane(f"{tag}_dir", direction)
    return result