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, 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))