From 60a664874fc09044583dad18c7a89abdfec6cd0c Mon Sep 17 00:00:00 2001 From: adam-urbanczyk Date: Sun, 6 Sep 2020 17:45:32 +0200 Subject: [PATCH] Color implementation for Assembly --- cadquery/__init__.py | 3 +- cadquery/assembly.py | 11 +++++ cadquery/occ_impl/assembly.py | 50 ++++++++++++++++++++-- cadquery/occ_impl/exporters/assembly.py | 2 + tests/test_assembly.py | 57 +++++++++++++++++-------- 5 files changed, 102 insertions(+), 21 deletions(-) diff --git a/cadquery/__init__.py b/cadquery/__init__.py index 52e574ae..691f0326 100644 --- a/cadquery/__init__.py +++ b/cadquery/__init__.py @@ -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", diff --git a/cadquery/assembly.py b/cadquery/assembly.py index 6d9e2aa9..7e3b5b0f 100644 --- a/cadquery/assembly.py +++ b/cadquery/assembly.py @@ -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) diff --git a/cadquery/occ_impl/assembly.py b/cadquery/occ_impl/assembly.py index 2e8bd3e6..c4e7fca3 100644 --- a/cadquery/occ_impl/assembly.py +++ b/cadquery/occ_impl/assembly.py @@ -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) diff --git a/cadquery/occ_impl/exporters/assembly.py b/cadquery/occ_impl/exporters/assembly.py index f07f091e..88a0ae9e 100644 --- a/cadquery/occ_impl/exporters/assembly.py +++ b/cadquery/occ_impl/exporters/assembly.py @@ -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) diff --git a/tests/test_assembly.py b/tests/test_assembly.py index 39c9d7bc..d412d740 100644 --- a/tests/test_assembly.py +++ b/tests/test_assembly.py @@ -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') \ No newline at end of file + simple_assy.save("simple.step", "DXF")