blenderings/playground/playground.py

import bpy
from bpy.types import (
    Curve,
    LayerCollection,
    MaterialSlot,
    ViewLayer,
    Object,
    Mesh,
    Collection,
)
from pathlib import Path
from typing import Generic, TypeVar, Union
from mathutils import Color

# Tiny helper to make autocompletion more precise for objects
T = TypeVar("T")


class ObjectData(Generic[T], Object):
    data: T


# Default names

PARTS = ("controller", "curve", "fin", "maker", "screw", "spool")


class Slide:
    collection: LayerCollection
    controller: ObjectData[None]
    curve: ObjectData[Curve]
    fin: ObjectData[Curve]
    maker: ObjectData[Mesh]
    screw: ObjectData[Mesh]
    spool: ObjectData[Mesh]

    @staticmethod
    def from_curve(curve: ObjectData[Curve], blend_file_path: str) -> "Slide":
        slide = Slide(blend_file_path)
        bpy.data.objects.remove(slide.curve, do_unlink=True)
        slide.change_curve(curve)
        return slide

    @classmethod
    def rehydrate(Cls, c: Collection) -> Union["Slide", None]:
        slide = Cls()
        slide.__dict__.update(get_parts(c))
        return slide

    def __init__(self, blend_file_path: str = None):
        if blend_file_path:
            parts = load_slide(blend_file_path)
        else:
            parts = {"collection": None}
            parts.update({part: None for part in PARTS})
        self.__dict__.update(parts)
        if blend_file_path:
            self.__localize()

    def __localize(self):
        "Make sure node inputs are linked to themselves"
        self.change_curve(self.curve)
        self.maker.modifiers["AddSpool"]["Input_3"] = self.spool
        slide_body_material = self.maker.material_slots[0].material
        self.maker.modifiers["AddSpool"]["Input_7"] = slide_body_material
        self.fin.modifiers["AddFin"]["Input_3"] = self.spool.material_slots[0].material
        self.screw.hide_viewport = True
        self.screw.hide_render = True
        self.spool.hide_viewport = True
        self.spool.hide_render = True
        for part in PARTS[1:]:
            getattr(self, part).parent = self.controller

    @property
    def radius(self) -> float:
        return self.maker.modifiers["MakeTunnel"]["Input_2"]

    @radius.setter
    def radius(self, v: float):
        self.maker.modifiers["MakeTunnel"]["Input_2"] = v
        self.maker.modifiers["AddSpool"]["Input_8"] = v
        self.fin.modifiers["AddFin"]["Input_4"] = v

    @property
    def thickness(self) -> float:
        return self.maker.modifiers["Solidify"].thickness

    @thickness.setter
    def thickness(self, v: float):
        self.maker.modifiers["Solidify"].thickness = v
        self.maker.modifiers["AddSpool"]["Input_9"] = v
        self.fin.modifiers["AddFin"]["Input_5"] = v

    @property
    def spool_spacing(self) -> float:
        return self.maker.modifiers["AddSpool"]["Input_5"]

    @spool_spacing.setter
    def spool_spacing(self, v: float):
        self.maker.modifiers["AddSpool"]["Input_5"] = v

    @property
    def body_color(self) -> Color:
        c = tuple(
            self.maker.material_slots[0]
            .material.node_tree.nodes["Principled BSDF.001"]
            .inputs[0]
            .default_value
        )
        return Color(c[:3])

    @body_color.setter
    def body_color(self, c: Color):
        c = (*c, 1.0)
        self.maker.material_slots[0].material.node_tree.nodes[
            "Principled BSDF.001"
        ].inputs[0].default_value = c

    @property
    def secondary_color(self) -> Color:
        c = tuple(
            self.spool.material_slots[0]
            .material.node_tree.nodes["Principled BSDF.001"]
            .inputs[0]
            .default_value
        )
        return Color(c[:3])

    @secondary_color.setter
    def secondary_color(self, c: Color):
        c = (*c, 1.0)
        self.spool.material_slots[0].material.node_tree.nodes[
            "Principled BSDF.001"
        ].inputs[0].default_value = c

    def change_curve(self, c: ObjectData[Curve]):
        self.maker.modifiers["MakeTunnel"]["Input_5"] = c
        self.maker.modifiers["AddSpool"]["Input_6"] = c
        self.fin.modifiers["AddFin"]["Input_2"] = c
        c.parent = self.controller
        if not any(c == o for o in self.collection.collection.all_objects):
            self.collection.collection.objects.link(c)


def get_parts(c: Collection):
    return {
        part: next(obj for obj in c.all_objects if part in obj.name.lower())
        for part in PARTS
    }


def toplevel_collection_to_layer_collection(
    c: Collection, view_layer: ViewLayer
) -> LayerCollection:
    return next(cv for cv in view_layer.layer_collection.children if cv.collection == c)


def duplicate_collection_objects(c: Collection, name: str):
    c2 = bpy.data.collections.new(name)
    for obj in c.all_objects:
        c2.objects.link(obj.copy())
    return c2


def dup_obj_materials_inplace(obj: Object):
    for slot in obj.material_slots:
        slot: MaterialSlot
        slot.material = slot.material.copy()


def load_slide(file_path: str):
    bpy.ops.wm.append(
        filepath=f"{file_path}/Scene/SlideTemplate",
        directory=f"{file_path}/Scene/",
        filename="SlideTemplate",
    )
    scene = bpy.data.scenes["SlideTemplate"]
    template = scene.collection.children[0]

    c = duplicate_collection_objects(template, "Slide")
    bpy.context.scene.collection.children.link(c)
    bpy.data.scenes.remove(scene)
    c = toplevel_collection_to_layer_collection(c, bpy.context.scene.view_layers[0])
    parts = get_parts(c.collection)
    for part_obj in parts.values():
        dup_obj_materials_inplace(part_obj)

    return {"collection": c, **parts}


DEFAULT_SLIDE_BLEND_PATH = str((Path(__file__) / ".." / "playground.blend").resolve())

if __name__ == "__main__":
    slide = Slide(DEFAULT_SLIDE_BLEND_PATH)