screw_generator.py

import cadquery as cq
import math

def create_realistic_thread(length, major_dia, pitch=1.0):
    thread_height = 0.6134 * pitch  # Metric standard thread profile height

    # Create helical path for thread
    turns = length / pitch
    helix = cq.Workplane("XY").helix(pitch, length, major_dia / 2 - thread_height / 2)

    # Triangular thread profile (60°)
    profile = cq.Workplane("XZ").polyline([
        (0, 0),
        (pitch / 2, thread_height),
        (pitch, 0)
    ]).close()

    # Sweep profile along helix
    thread = profile.sweep(helix, multisection=True)
    return thread

def generate_screw(thread_diameter: float, length: float, head_type: str):
    shaft_radius = thread_diameter / 2 * 0.95
    shaft = cq.Workplane("XY").circle(shaft_radius).extrude(length)

    # Head
    if head_type == "hex":
        head = cq.Workplane("XY").polygon(6, thread_diameter * 1.5).extrude(thread_diameter * 0.6)
    elif head_type == "countersunk":
        head = cq.Workplane("XY").circle(thread_diameter * 0.75).extrude(thread_diameter * 0.4)
        head = head.edges("|Z").chamfer(thread_diameter * 0.25)
    elif head_type == "pan":
        head = cq.Workplane("XY").circle(thread_diameter).extrude(thread_diameter * 0.4)
    else:
        raise ValueError("Invalid head type")

    # Create thread
    pitch = 1.5  # Default pitch
    thread = create_realistic_thread(length, thread_diameter, pitch)

    # Cut thread from shaft
    screw = head.union(shaft).cut(thread)

    # Increase export quality
    cq.exporters.export(screw, "screw.stl", tolerance=0.01, angularTolerance=0.05)
    return "screw.stl"