AI로 자동 설계 프로그램 만들기 - 2단 Shaft

코드

"""
FreeCAD Part Design - 원기둥 2개 생성/수정 GUI 매크로
- 원기둥 1: 원점(X=0) 시작, X축 방향
- 원기둥 2: 동일한 XY_Plane, 로컬 원점을 X=L1 으로 오프셋
- FreeCAD 1.x (PySide6) 대응
"""

import FreeCAD
import FreeCADGui
import Part
import Sketcher

try:
    from PySide6 import QtWidgets, QtCore
except ImportError:
    from PySide2 import QtWidgets, QtCore

# ─────────────────────────────────────────────────────────────────────────────
# 전역 상태
# ─────────────────────────────────────────────────────────────────────────────
_state = {
    "doc":      None,
    "body":     None,
    "xy_plane": None,
    # 원기둥 1
    "sketch1":  None,
    "rev1":     None,
    "idx1_len": -1,
    "idx1_rad": -1,
    "L1":       0.0,   # 원기둥 1 현재 길이 (원기둥 2 오프셋 계산용)
    # 원기둥 2
    "sketch2":  None,
    "rev2":     None,
    "idx2_len": -1,
    "idx2_rad": -1,
    "created1": False,
    "created2": False,
}

# ─────────────────────────────────────────────────────────────────────────────
# XY_Plane 탐색
# ─────────────────────────────────────────────────────────────────────────────
def _find_xy_plane(doc):
    for obj in doc.Objects:
        if obj.Name == "XY_Plane" or obj.Label == "XY_Plane":
            return obj
    origin = doc.getObject("Origin")
    if origin and hasattr(origin, "OriginFeatures"):
        for child in origin.OriginFeatures:
            if "XY" in child.Label:
                return child
    return None

# ─────────────────────────────────────────────────────────────────────────────
# 스케치 + Revolution 공통 빌더
# ─────────────────────────────────────────────────────────────────────────────
def _build_sketch_revolution(body, doc, name_suffix, R, L, x_offset, xy_plane):
    """
    XY_Plane 에 단면 사각형을 그리고 Revolution(X축)을 만든다.
    x_offset : 스케치 로컬 원점의 월드 X 좌표 (MapPathParameter 방식 대신
               AttachmentOffset 으로 평행이동)
    """
    sketch = body.newObject("Sketcher::SketchObject", f"Sketch{name_suffix}")
    doc.recompute()

    # ── 동일 XY_Plane 부착 ────────────────────────────────────────────
    try:
        sketch.AttachmentSupport = (xy_plane, [""])
    except Exception:
        try:
            sketch.Support = (xy_plane, [""])
        except Exception:
            pass
    sketch.MapMode = "FlatFace"

    # ── X 방향 오프셋 적용 (AttachmentOffset) ─────────────────────────
    if x_offset != 0.0:
        sketch.AttachmentOffset = FreeCAD.Placement(
            FreeCAD.Vector(x_offset, 0, 0),
            FreeCAD.Rotation(0, 0, 0)
        )

    doc.recompute()

    # ── 단면 사각형 (로컬 좌표) ───────────────────────────────────────
    # 0: (0,0)→(L,0)   회전축(X)
    # 1: (L,0)→(L,R)   끝단
    # 2: (L,R)→(0,R)   외경
    # 3: (0,R)→(0,0)   시작단
    sketch.addGeometry(Part.LineSegment(FreeCAD.Vector(0, 0, 0), FreeCAD.Vector(L,  0, 0)), False)
    sketch.addGeometry(Part.LineSegment(FreeCAD.Vector(L, 0, 0), FreeCAD.Vector(L,  R, 0)), False)
    sketch.addGeometry(Part.LineSegment(FreeCAD.Vector(L, R, 0), FreeCAD.Vector(0,  R, 0)), False)
    sketch.addGeometry(Part.LineSegment(FreeCAD.Vector(0, R, 0), FreeCAD.Vector(0,  0, 0)), False)

    # ── 기하 구속 ─────────────────────────────────────────────────────
    sketch.addConstraint(Sketcher.Constraint("Coincident", 0, 2, 1, 1))   # 0
    sketch.addConstraint(Sketcher.Constraint("Coincident", 1, 2, 2, 1))   # 1
    sketch.addConstraint(Sketcher.Constraint("Coincident", 2, 2, 3, 1))   # 2
    sketch.addConstraint(Sketcher.Constraint("Coincident", 3, 2, 0, 1))   # 3
    sketch.addConstraint(Sketcher.Constraint("Coincident", 0, 1, -1, 1))  # 4 로컬 원점 고정
    sketch.addConstraint(Sketcher.Constraint("Horizontal", 0))             # 5
    sketch.addConstraint(Sketcher.Constraint("Horizontal", 2))             # 6
    sketch.addConstraint(Sketcher.Constraint("Vertical",   1))             # 7
    sketch.addConstraint(Sketcher.Constraint("Vertical",   3))             # 8

    # ── 치수 구속 ─────────────────────────────────────────────────────
    idx_len = sketch.ConstraintCount          # 9
    sketch.addConstraint(Sketcher.Constraint("DistanceX", 0, 1, 0, 2, L))

    idx_rad = sketch.ConstraintCount          # 10
    sketch.addConstraint(Sketcher.Constraint("DistanceY", 1, 1, 1, 2, R))

    doc.recompute()

    # ── Revolution ────────────────────────────────────────────────────
    revolution = body.newObject("PartDesign::Revolution", f"Revolution{name_suffix}")
    revolution.Profile       = sketch
    revolution.ReferenceAxis = (sketch, ["H_Axis"])
    revolution.Angle         = 360.0
    try:
        revolution.Reversed = False
    except Exception:
        pass
    doc.recompute()

    return sketch, revolution, idx_len, idx_rad


# ─────────────────────────────────────────────────────────────────────────────
# 원기둥 1 생성
# ─────────────────────────────────────────────────────────────────────────────
def _build_cylinder1(diameter: float, length: float):
    R, L = diameter / 2.0, length

    doc  = FreeCAD.newDocument("TwoCylinders")
    body = doc.addObject("PartDesign::Body", "Body")
    doc.recompute()

    xy_plane = _find_xy_plane(doc)

    sk, rv, idx_len, idx_rad = _build_sketch_revolution(
        body, doc, "1", R, L, x_offset=0.0, xy_plane=xy_plane
    )

    _state.update({
        "doc": doc, "body": body, "xy_plane": xy_plane,
        "sketch1": sk, "rev1": rv,
        "idx1_len": idx_len, "idx1_rad": idx_rad,
        "L1": length,
        "created1": True,
    })
    _view_fit()
    FreeCAD.Console.PrintMessage(f"[Cyl-1] 생성 | L={L}, R={R}\n")


# ─────────────────────────────────────────────────────────────────────────────
# 원기둥 2 생성  (동일 XY_Plane, x_offset = L1)
# ─────────────────────────────────────────────────────────────────────────────
def _build_cylinder2(diameter: float, length: float):
    if not _state["created1"]:
        raise RuntimeError("원기둥 1을 먼저 생성하세요.")

    R, L     = diameter / 2.0, length
    doc      = _state["doc"]
    body     = _state["body"]
    xy_plane = _state["xy_plane"]
    L1       = _state["L1"]          # 원기둥 2 시작 위치 = 원기둥 1 길이

    sk, rv, idx_len, idx_rad = _build_sketch_revolution(
        body, doc, "2", R, L, x_offset=L1, xy_plane=xy_plane
    )

    _state.update({
        "sketch2": sk, "rev2": rv,
        "idx2_len": idx_len, "idx2_rad": idx_rad,
        "created2": True,
    })
    _view_fit()
    FreeCAD.Console.PrintMessage(f"[Cyl-2] 생성 | L={L}, R={R}, offset={L1}\n")


# ─────────────────────────────────────────────────────────────────────────────
# 수정
# ─────────────────────────────────────────────────────────────────────────────
def _update_cylinder(cyl_no: int, new_diameter: float, new_length: float):
    R, L = new_diameter / 2.0, new_length
    doc  = _state["doc"]

    if cyl_no == 1:
        sketch  = _state["sketch1"]
        idx_len = _state["idx1_len"]
        idx_rad = _state["idx1_rad"]
    else:
        sketch  = _state["sketch2"]
        idx_len = _state["idx2_len"]
        idx_rad = _state["idx2_rad"]

    sketch.setDatum(idx_len, FreeCAD.Units.Quantity(f"{L} mm"))
    sketch.setDatum(idx_rad, FreeCAD.Units.Quantity(f"{R} mm"))

    # 원기둥 1 길이가 바뀌면 원기둥 2 의 AttachmentOffset 도 갱신
    if cyl_no == 1:
        _state["L1"] = new_length
        if _state["created2"]:
            sk2 = _state["sketch2"]
            sk2.AttachmentOffset = FreeCAD.Placement(
                FreeCAD.Vector(new_length, 0, 0),
                FreeCAD.Rotation(0, 0, 0)
            )
            FreeCAD.Console.PrintMessage(
                f"[Cyl-2] offset 갱신 → X={new_length}\n"
            )

    doc.recompute()
    _view_fit()
    FreeCAD.Console.PrintMessage(
        f"[Cyl-{cyl_no}] 수정 완료 | L={L}, R={R}\n"
    )


def _view_fit():
    try:
        FreeCADGui.activeDocument().activeView().viewIsometric()
        FreeCADGui.SendMsgToActiveView("ViewFit")
    except Exception:
        pass


# ─────────────────────────────────────────────────────────────────────────────
# GUI
# ─────────────────────────────────────────────────────────────────────────────
class TwoCylinderDialog(QtWidgets.QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("원기둥 2개 생성 / 수정")
        self.setFixedWidth(460)
        try:
            self.setWindowFlags(self.windowFlags() | QtCore.Qt.WindowType.WindowStaysOnTopHint)
        except AttributeError:
            self.setWindowFlags(self.windowFlags() | QtCore.Qt.WindowStaysOnTopHint)
        self._build_ui()

    def _build_ui(self):
        root = QtWidgets.QVBoxLayout(self)
        root.setSpacing(10)
        root.setContentsMargins(16, 16, 16, 16)

        title = QtWidgets.QLabel("🔩  FreeCAD 원기둥 × 2  (X축 직렬 · 동일 XY 평면)")
        self._center(title)
        title.setStyleSheet("font-size:13px; font-weight:bold; color:#4FC3F7;")
        root.addWidget(title)

        desc = QtWidgets.QLabel("두 스케치 모두 XY_Plane 사용 · 원기둥 2는 AttachmentOffset(X=L1) 적용")
        self._center(desc)
        desc.setStyleSheet("font-size:10px; color:#78909C;")
        root.addWidget(desc)
        root.addWidget(self._hline())

        # ── 원기둥 1 ──────────────────────────────────────────────────
        root.addWidget(self._sec("● 원기둥 1  (원점 X=0 시작)", "#4FC3F7"))
        root.addLayout(self._input_row(1))
        root.addLayout(self._btn_row(1))
        root.addWidget(self._hline())

        # ── 원기둥 2 ──────────────────────────────────────────────────
        root.addWidget(self._sec("● 원기둥 2  (X = L1 위치 시작, 동일 평면)", "#FFB74D"))
        root.addLayout(self._input_row(2))
        root.addLayout(self._btn_row(2))
        root.addWidget(self._hline())

        # ── 상태 ──────────────────────────────────────────────────────
        self.lbl_status = QtWidgets.QLabel("💡 원기둥 1을 먼저 생성하세요.")
        self.lbl_status.setWordWrap(True)
        self.lbl_status.setStyleSheet("color:#B0BEC5; font-size:11px;")
        root.addWidget(self.lbl_status)

        btn_close = QtWidgets.QPushButton("닫   기")
        btn_close.setFixedHeight(32)
        btn_close.setStyleSheet(self._btn("#455A64", "#607D8B"))
        btn_close.clicked.connect(self.close)
        root.addWidget(btn_close)

        self.setStyleSheet(
            "QDialog        { background:#1A2535; color:#ECEFF1; }"
            "QLabel          { color:#ECEFF1; }"
            "QDoubleSpinBox  { background:#243447; color:#ECEFF1;"
            "                  border:1px solid #455A64; border-radius:4px; padding:4px; }"
            "QDoubleSpinBox::up-button, QDoubleSpinBox::down-button { background:#37474F; }"
        )

    def _input_row(self, n):
        row = QtWidgets.QHBoxLayout()
        diam   = self._spin(1, 2000, 80 if n == 1 else 60,  "mm")
        length = self._spin(1, 5000, 150 if n == 1 else 100, "mm")
        row.addWidget(QtWidgets.QLabel("직경:"))
        row.addWidget(diam)
        row.addSpacing(8)
        row.addWidget(QtWidgets.QLabel("길이:"))
        row.addWidget(length)
        if n == 1:
            self.inp1_diam, self.inp1_length = diam, length
        else:
            self.inp2_diam, self.inp2_length = diam, length
        return row

    def _btn_row(self, n):
        row = QtWidgets.QHBoxLayout()
        if n == 1:
            bc = self._make_btn("▶  원기둥 1  생성", "#1565C0", "#1E88E5", enabled=True)
            be = self._make_btn("✏  원기둥 1  수정", "#2E7D32", "#43A047", enabled=False)
            bc.clicked.connect(self._on_create1)
            be.clicked.connect(self._on_edit1)
            self.btn_create1, self.btn_edit1 = bc, be
        else:
            bc = self._make_btn("▶  원기둥 2  생성", "#6A1B9A", "#8E24AA", enabled=False)
            be = self._make_btn("✏  원기둥 2  수정", "#E65100", "#FB8C00", enabled=False)
            bc.clicked.connect(self._on_create2)
            be.clicked.connect(self._on_edit2)
            self.btn_create2, self.btn_edit2 = bc, be
        row.addWidget(bc)
        row.addWidget(be)
        return row

    # ── 이벤트 ────────────────────────────────────────────────────────
    def _on_create1(self):
        d, l = self.inp1_diam.value(), self.inp1_length.value()
        try:
            _build_cylinder1(d, l)
            self.btn_create1.setEnabled(False)
            self.btn_edit1.setEnabled(True)
            self.btn_create2.setEnabled(True)
            self._msg(f"✅ 원기둥 1 생성  Ø{d} × L{l} mm", "#A5D6A7")
        except Exception as e:
            self._msg(f"❌ 원기둥 1 생성 실패: {e}", "#EF9A9A")

    def _on_edit1(self):
        d, l = self.inp1_diam.value(), self.inp1_length.value()
        try:
            _update_cylinder(1, d, l)
            self._msg(f"✏ 원기둥 1 수정  Ø{d} × L{l} mm  "
                      f"{'→ 원기둥 2 offset 갱신' if _state['created2'] else ''}",
                      "#FFF59D")
        except Exception as e:
            self._msg(f"❌ 원기둥 1 수정 실패: {e}", "#EF9A9A")

    def _on_create2(self):
        d, l = self.inp2_diam.value(), self.inp2_length.value()
        try:
            _build_cylinder2(d, l)
            self.btn_create2.setEnabled(False)
            self.btn_edit2.setEnabled(True)
            self._msg(f"✅ 원기둥 2 생성  Ø{d} × L{l} mm  (offset=L1={_state['L1']} mm)",
                      "#FFCC80")
        except Exception as e:
            self._msg(f"❌ 원기둥 2 생성 실패: {e}", "#EF9A9A")

    def _on_edit2(self):
        d, l = self.inp2_diam.value(), self.inp2_length.value()
        try:
            _update_cylinder(2, d, l)
            self._msg(f"✏ 원기둥 2 수정  Ø{d} × L{l} mm", "#FFF59D")
        except Exception as e:
            self._msg(f"❌ 원기둥 2 수정 실패: {e}", "#EF9A9A")

    # ── 헬퍼 ──────────────────────────────────────────────────────────
    def _msg(self, text, color="#B0BEC5"):
        self.lbl_status.setText(text)
        self.lbl_status.setStyleSheet(f"color:{color}; font-size:11px;")

    def _make_btn(self, label, base, hover, enabled=True):
        b = QtWidgets.QPushButton(label)
        b.setFixedHeight(38)
        b.setEnabled(enabled)
        b.setStyleSheet(self._btn(base, hover))
        return b

    @staticmethod
    def _sec(text, color):
        lbl = QtWidgets.QLabel(text)
        lbl.setStyleSheet(f"color:{color}; font-weight:bold; font-size:12px; padding:2px 0;")
        return lbl

    @staticmethod
    def _center(lbl):
        try:
            lbl.setAlignment(QtCore.Qt.AlignmentFlag.AlignCenter)
        except AttributeError:
            lbl.setAlignment(QtCore.Qt.AlignCenter)

    @staticmethod
    def _spin(mn, mx, val, suf):
        sb = QtWidgets.QDoubleSpinBox()
        sb.setRange(mn, mx)
        sb.setValue(val)
        sb.setSuffix(f"  {suf}")
        sb.setDecimals(2)
        sb.setSingleStep(1.0)
        sb.setFixedHeight(30)
        return sb

    @staticmethod
    def _hline():
        f = QtWidgets.QFrame()
        try:
            f.setFrameShape(QtWidgets.QFrame.Shape.HLine)
        except AttributeError:
            f.setFrameShape(QtWidgets.QFrame.HLine)
        f.setStyleSheet("color:#2E3F50;")
        return f

    @staticmethod
    def _btn(base, hover):
        return (f"QPushButton          {{ background:{base}; color:white; border:none;"
                f"                        border-radius:5px; font-size:12px; font-weight:bold; }}"
                f"QPushButton:hover    {{ background:{hover}; }}"
                f"QPushButton:pressed  {{ background:{base}; }}"
                f"QPushButton:disabled {{ background:#37474F; color:#546E7A; }}")


# ─────────────────────────────────────────────────────────────────────────────
# 진입점
# ─────────────────────────────────────────────────────────────────────────────
def main():
    for w in QtWidgets.QApplication.topLevelWidgets():
        if isinstance(w, TwoCylinderDialog):
            w.raise_()
            w.activateWindow()
            return
    dlg = TwoCylinderDialog(FreeCADGui.getMainWindow())
    dlg.show()

main()

 

프로그램 실행

 

by korealionkk@gmail.com