Cosplay/nhf/build.py

"""
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
from functools import wraps
from colorama import Fore, Style
import cadquery as Cq
import nhf.checks as NC

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: str,
                 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):
        """
        Output file name
        """
        return f"{self.name}.{self.kind.ext}"

    def write_to(self, obj, path: str):
        x = self._method(obj)
        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:
            assert isinstance(x, Cq.Workplane)
            Cq.exporters.exportDXF(x, path, **self.kwargs)
        else:
            assert False, f"Invalid kind: {self.kind}"

    @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(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._target = Target(method, name, **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 Model:
    """
    Base class for a parametric assembly
    """
    def __init__(self, name: str):
        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():
            t(self)
            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", verbose=1):
        """
        Build all targets in this model and write the results to file
        """
        output_dir = Path(output_dir)
        for t in Target.methods(self).values():
            output_file = output_dir / self.name / t.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:
                t.write_to(self, str(output_file))
                if 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}")