Cosplay/nhf/build.py

314 lines
9.0 KiB
Python

"""
The NHF build system
Usage: For any parametric assembly, inherit the `Model` class, and mark the
output objects with the `@target` decorator. Each marked function should only
take `self` as an argument.
```python
class BuildScaffold(Model):
@target(name="obj1")
def o1(self):
return Cq.Solid.makeBox(10, 10, 10)
def o2(self):
return Cq.Solid.makeCylinder(10, 20)
```
"""
from enum import Enum
from pathlib import Path
from typing import Union, Optional
from functools import wraps
import traceback
from colorama import Fore, Style
import cadquery as Cq
import nhf.checks as NC
import nhf.utils
TOL=1e-6
class TargetKind(Enum):
STL = "stl",
DXF = "dxf",
def __init__(self, ext: str):
self.ext = ext
class Target:
"""
Marks a function's output for serialization
"""
def __init__(self,
method,
name: Optional[str] = None,
prototype: bool = False,
kind: TargetKind = TargetKind.STL,
**kwargs):
self._method = method
self.name = name
self.prototype = prototype
self.kind = kind
self.kwargs = kwargs
def __str__(self):
return f"<target {self.name}.{self.kind.ext} {self._method}>"
def __call__(self, obj, *args, **kwargs):
"""
Raw call function which passes arguments directly to `_method`
"""
return self._method(obj, *args, **kwargs)
@property
def file_name(self) -> Optional[str]:
"""
Output file name
"""
if self.name:
return f"{self.name}.{self.kind.ext}"
else:
return None
def write_to(self, obj, path: str) -> bool:
"""
Returns false if target is `None`
"""
x = self._method(obj)
if x is None:
return False
if self.kind == TargetKind.STL:
assert isinstance(x, Union[
Cq.Workplane, Cq.Shape, Cq.Compound, Cq.Assembly])
if isinstance(x, Cq.Workplane):
x = x.val()
if isinstance(x, Cq.Assembly):
x = x.toCompound().fuse(tol=TOL)
x.exportStl(path, **self.kwargs)
elif self.kind == TargetKind.DXF:
if isinstance(x, Cq.Sketch):
# https://github.com/CadQuery/cadquery/issues/1575
x = (
Cq.Workplane()
.add(x._faces)
.add(x.wires)
.add(x._edges)
)
assert isinstance(x, Cq.Workplane)
Cq.exporters.exportDXF(x, path, **self.kwargs)
else:
assert False, f"Invalid kind: {self.kind}"
return True
@classmethod
def methods(cls, subject):
"""
List of all methods of a class or objects annotated with this decorator.
"""
def g():
for name in dir(subject):
if name == 'target_names':
continue
method = getattr(subject, name)
if hasattr(method, '_target'):
yield method._target
return {method.name: method for method in g()}
def target(**deco_kwargs):
"""
Decorator for annotating a build output
"""
def f(method):
@wraps(method)
def wrapper(self, *args, **kwargs):
return method(self, *args, **kwargs)
wrapper._target = Target(method, **deco_kwargs)
return wrapper
return f
class Assembly:
"""
Marks a function's output for assembly property checking
"""
def __init__(self,
method,
collision_check: bool = True,
**kwargs):
self._method = method
self.name = method.__name__
self.collision_check = collision_check
self.kwargs = kwargs
def __str__(self):
return f"<assembly {self.name} {self._method}>"
def __call__(self, obj, *args, **kwargs):
"""
Raw call function which passes arguments directly to `_method`
"""
return self._method(obj, *args, **kwargs)
def check(self, obj):
x = self._method(obj)
assert isinstance(x, Cq.Assembly)
if self.collision_check:
intersections = NC.pairwise_intersection(x)
assert not intersections, f"In {self}, collision detected: {intersections}"
@classmethod
def methods(cls, subject):
"""
List of all methods of a class or objects annotated with this decorator.
"""
def g():
for name in dir(subject):
if name == 'target_names':
continue
method = getattr(subject, name)
if hasattr(method, '_assembly'):
yield method._assembly
return {method.name: method for method in g()}
def assembly(**deco_kwargs):
"""
Decorator for annotating an assembly output
"""
def f(method):
@wraps(method)
def wrapper(self, *args, **kwargs):
return method(self, *args, **kwargs)
wrapper._assembly = Assembly(method, **deco_kwargs)
return wrapper
return f
class Submodel:
"""
Marks a function's output as a submodel
"""
def __init__(self,
method,
name: str,
prototype: bool = False,
**kwargs):
self._method = method
self.name = name
self.prototype = prototype
self.kwargs = kwargs
def __str__(self):
return f"<submodel {self.name} {self._method}>"
def __call__(self, obj, *args, **kwargs):
"""
Raw call function which passes arguments directly to `_method`
"""
return self._method(obj, *args, **kwargs)
@property
def file_name(self):
"""
Output file name
"""
return self.name
def write_to(self, obj, path: str):
x = self._method(obj)
assert isinstance(x, Model), f"Unexpected type: {type(x)}"
x.build_all(path, prefix=False)
@classmethod
def methods(cls, subject):
"""
List of all methods of a class or objects annotated with this decorator.
"""
def g():
for name in dir(subject):
if name == 'target_names':
continue
method = getattr(subject, name)
if hasattr(method, '_submodel'):
yield method._submodel
return {method.name: method for method in g()}
def submodel(name, **deco_kwargs):
"""
Decorator for annotating a build output
"""
def f(method):
@wraps(method)
def wrapper(self, *args, **kwargs):
return method(self, *args, **kwargs)
wrapper._submodel = Submodel(method, name, **deco_kwargs)
return wrapper
return f
class Model:
"""
Base class for a parametric assembly
"""
def __init__(self, name: Optional[str] = None):
self.name = name
@property
def target_names(self) -> list[str]:
"""
List of decorated target functions
"""
return list(Target.methods(self).keys())
def check_all(self) -> int:
"""
Build all models and run all the checks. Return number of checks passed
"""
total = 0
for t in Target.methods(self).values():
result = t(self)
if result:
total += 1
for t in Assembly.methods(self).values():
t.check(self)
total += 1
return total
def build_all(
self,
output_dir: Union[Path, str] = "build",
prefix: bool = True,
verbose=1):
"""
Build all targets in this model and write the results to file
"""
output_dir = Path(output_dir)
if prefix:
output_dir = output_dir / self.name
targets = Target.methods(self)
for t in targets.values():
file_name = t.file_name
if file_name is None:
assert len(targets) == 1, "Only one anonymous target is permitted"
output_file = output_dir.with_suffix('.' + t.kind.ext)
else:
output_file = output_dir / file_name
if output_file.is_file():
if verbose >= 1:
print(f"{Fore.GREEN}Skipping{Style.RESET_ALL} {output_file}")
continue
output_file.parent.mkdir(exist_ok=True, parents=True)
if verbose >= 1:
print(f"{Fore.BLUE}Building{Style.RESET_ALL} {output_file}")
try:
flag = t.write_to(self, str(output_file))
if flag and verbose >= 1:
print(f"{Fore.GREEN}Built{Style.RESET_ALL} {output_file}")
except Exception as e:
print(f"{Fore.RED}Failed to build{Style.RESET_ALL} {output_file}: {e}")
traceback.print_exc()
for t in Submodel.methods(self).values():
d = output_dir / t.name
t.write_to(self, str(d))