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

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()
            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}")