ByoungSooPark/cadquery
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Color implementation for Assembly

Browse Source
This commit is contained in:
adam-urbanczyk
2020-09-06 17:45:32 +02:00
parent fe850502f4
commit 60a664874f
5 changed files with 102 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
cadquery/__init__.py
View File

@ -28,7 +28,7 @@ from .selectors import (
Selector,
)
from .cq import CQ, Workplane
from .assembly import Assembly
from .assembly import Assembly, Color
from . import selectors
from . import plugins
@ -37,6 +37,7 @@ __all__ = [
"CQ",
"Workplane",
"Assembly",
"Color",
"plugins",
"selectors",
"Plane",

11
cadquery/assembly.py
View File

@ -5,6 +5,7 @@ from uuid import uuid1 as uuid
from .cq import Workplane
from .occ_impl.shapes import Shape
from .occ_impl.geom import Location
from .occ_impl.assembly import Color
from .occ_impl.exporters.assembly import exportAssembly, exportCAF
@ -26,6 +27,7 @@ class Assembly(object):
loc: Location
name: str
color: Optional[Color]
metadata: Mapping[str, Any]
obj: AssemblyObjects
@ -40,6 +42,7 @@ class Assembly(object):
obj: AssemblyObjects = None,
loc: Optional[Location] = None,
name: Optional[str] = None,
color: Optional[Color] = None,
):
"""
construct an assembly
@ -47,8 +50,10 @@ class Assembly(object):
:param obj: root object of the assembly (deafault: None)
:param loc: location of the root object (deafault: None, interpreted as identity transformation)
:param name: unique name of the root object (default: None, reasulting in an UUID being generated)
:param color: color of the added object (default: None)
:return: An Assembly object.
To create an empt assembly use::
assy = Assembly(None)
@ -63,6 +68,7 @@ class Assembly(object):
self.obj = obj
self.loc = loc if loc else Location()
self.name = name if name else str(uuid())
self.color = color if color else None
self.parent = None
self.children = []
@ -74,6 +80,7 @@ class Assembly(object):
obj: "Assembly",
loc: Optional[Location] = None,
name: Optional[str] = None,
color: Optional[Color] = None,
) -> "Assembly":
"""
add a subassembly to the current assembly.
@ -81,6 +88,7 @@ class Assembly(object):
:param obj: subassembly to be added
:param loc: location of the root object (deafault: None, resulting in the location stored in the subassembly being used)
:param name: unique name of the root object (default: None, resulting in the name stored in the subassembly being used)
:param color: color of the added object (default: None, resulting in the color stored in the subassembly being used)
"""
...
@ -90,6 +98,7 @@ class Assembly(object):
obj: AssemblyObjects,
loc: Optional[Location] = None,
name: Optional[str] = None,
color: Optional[Color] = None,
) -> "Assembly":
"""
add a subassembly to the current assembly with explicit location and name
@ -97,6 +106,7 @@ class Assembly(object):
:param obj: object to be added as a subassembly
:param loc: location of the root object (deafault: None, interpreted as identity transformation)
:param name: unique name of the root object (default: None, resulting in an UUID being generated)
:param color: color of the added object (default: None)
"""
...
@ -108,6 +118,7 @@ class Assembly(object):
arg.obj,
kwargs["loc"] if kwargs.get("loc") else arg.loc,
kwargs["name"] if kwargs.get("name") else arg.name,
kwargs["color"] if kwargs.get("color") else arg.color,
)
subassy.children.extend(arg.children)

50
cadquery/occ_impl/assembly.py
View File

@ -1,17 +1,49 @@
from typing import Iterable, Tuple, Dict
from typing import Iterable, Tuple, Dict, overload, Optional
from typing_extensions import Protocol
from OCP.TDocStd import TDocStd_Document
from OCP.TCollection import TCollection_ExtendedString
from OCP.XCAFDoc import XCAFDoc_DocumentTool
from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorType
from OCP.TDataStd import TDataStd_Name
from OCP.TDF import TDF_Label
from OCP.TopLoc import TopLoc_Location
from OCP.Quantity import Quantity_ColorRGBA
from .geom import Location
from .shapes import Shape, Compound
class Color(object):
wrapped: Quantity_ColorRGBA
@overload
def __init__(self, name: str):
...
@overload
def __init__(self, r: float, g: float, b: float, a: float = 0):
...
def __init__(self, *args, **kwargs):
if len(args) == 1:
self.wrapped = Quantity_ColorRGBA()
exists = Quantity_ColorRGBA.ColorFromName_s(args[0], self.wrapped)
if not exists:
raise ValueError(f"Unknown color name: {args[0]}")
elif len(args) == 3:
r, g, b = args
self.wrapped = Quantity_ColorRGBA(r, g, b, 1)
if kwargs.get("a"):
self.wrapped.SetAlpha(kwargs.get("a"))
elif len(args) == 4:
r, g, b, a = args
self.wrapped = Quantity_ColorRGBA(r, g, b, a)
else:
raise ValueError(f"Unsupported arguments: {args}, {kwargs}")
class AssemblyProtocol(Protocol):
@property
def loc(self) -> Location:
@ -21,6 +53,10 @@ class AssemblyProtocol(Protocol):
def name(self) -> str:
...
@property
def color(self) -> Optional[Color]:
...
@property
def shapes(self) -> Iterable[Shape]:
...
@ -33,17 +69,23 @@ class AssemblyProtocol(Protocol):
...
def setName(l: TDF_Label, name, tool):
def setName(l: TDF_Label, name: str, tool):
TDataStd_Name.Set_s(l, TCollection_ExtendedString(name))
def setColor(l: TDF_Label, color: Color, tool):
tool.SetColor(l, color.wrapped, XCAFDoc_ColorType.XCAFDoc_ColorSurf)
def toCAF(assy: AssemblyProtocol) -> Tuple[TDF_Label, TDocStd_Document]:
# prepare a doc
doc = TDocStd_Document(TCollection_ExtendedString("XmlOcaf"))
tool = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main())
tool.SetAutoNaming_s(False)
ctool = XCAFDoc_DocumentTool.ColorTool_s(doc.Main())
# add root
top = tool.NewShape()
@ -61,6 +103,8 @@ def toCAF(assy: AssemblyProtocol) -> Tuple[TDF_Label, TDocStd_Document]:
subassy = tool.NewShape()
tool.AddComponent(subassy, lab, TopLoc_Location())
setName(subassy, k, tool)
if v.color:
setColor(subassy, v.color, ctool)
subassys[k] = (subassy, v.loc)

2
cadquery/occ_impl/exporters/assembly.py
View File

@ -20,6 +20,8 @@ def exportAssembly(assy: AssemblyProtocol, path: str) -> bool:
session = XSControl_WorkSession()
writer = STEPCAFControl_Writer(session, False)
writer.SetNameMode(True)
writer.SetColorMode(True)
writer.SetLayerMode(True)
writer.Transfer(doc, STEPControl_StepModelType.STEPControl_AsIs)
status = writer.Write(path)

43
tests/test_assembly.py
View File

@ -8,7 +8,7 @@ from cadquery.occ_impl.exporters.assembly import exportAssembly, exportCAF
@pytest.fixture
def simple_assy():
b1 = cq.Solid.makeBox(1,1,1)
b1 = cq.Solid.makeBox(1, 1, 1)
b2 = cq.Workplane().box(1, 1, 2)
b3 = cq.Workplane().pushPoints([(0, 0), (-2, -5)]).box(1, 1, 3)
@ -24,16 +24,38 @@ def nested_assy():
b1 = cq.Workplane().box(1, 1, 1)
b2 = cq.Workplane().box(1, 1, 1)
b3 = cq.Workplane().pushPoints([(-2, 0), (2, 0)]).box(1, 1, .5)
b3 = cq.Workplane().pushPoints([(-2, 0), (2, 0)]).box(1, 1, 0.5)
assy = cq.Assembly(b1, loc=cq.Location(cq.Vector(0, 0, 0)), name="TOP")
assy2 = cq.Assembly(b2, loc=cq.Location(cq.Vector(0, 4, 0)), name="SECOND")
assy2.add(b3, loc=cq.Location(cq.Vector(0, 4, 0)), name="BOTTOM")
assy.add(assy2)
assy.add(assy2, color=cq.Color("green"))
return assy
def test_color():
c1 = cq.Color("red")
assert c1.wrapped.GetRGB().Red() == 1
assert c1.wrapped.Alpha() == 1
c2 = cq.Color(1, 0, 0)
assert c2.wrapped.GetRGB().Red() == 1
assert c2.wrapped.Alpha() == 1
c3 = cq.Color(1, 0, 0, 0.5)
assert c3.wrapped.GetRGB().Red() == 1
assert c3.wrapped.Alpha() == 0.5
with pytest.raises(ValueError):
cq.Color("?????")
with pytest.raises(ValueError):
cq.Color(1, 2, 3, 4, 5)
def test_assembly(simple_assy, nested_assy):
# basic checks
@ -72,25 +94,26 @@ def test_native_export(simple_assy):
# only sanity check for now
assert os.path.exists("assy.xml")
def test_save(simple_assy):
simple_assy.save('simple.step')
simple_assy.save("simple.step")
assert os.path.exists("simple.step")
simple_assy.save('simple.xml')
simple_assy.save("simple.xml")
assert os.path.exists("simple.xml")
simple_assy.save('simple.step')
simple_assy.save("simple.step")
assert os.path.exists("simple.step")
simple_assy.save('simple.stp','STEP')
simple_assy.save("simple.stp", "STEP")
assert os.path.exists("simple.stp")
simple_assy.save('simple.caf','XML')
simple_assy.save("simple.caf", "XML")
assert os.path.exists("simple.caf")
with pytest.raises(ValueError):
simple_assy.save('simple.dxf')
simple_assy.save("simple.dxf")
with pytest.raises(ValueError):
simple_assy.save('simple.step','DXF')
simple_assy.save("simple.step", "DXF")