blenderings/playground/playground.py

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


SCENE = bpy.data.scenes[0]

# 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]

    def __init__(self, blend_file_path: str):
        self.__dict__.update(load_slide(blend_file_path))
        self.__localize()

    def __localize(self):
        "Make sure node inputs are linked to themselves"
        self.maker.modifiers["MakeTunnel"]["Input_5"] = self.curve
        self.maker.modifiers["AddSpool"]["Input_6"] = 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_2"] = self.curve
        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 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")
    SCENE.collection.children.link(c)
    bpy.data.scenes.remove(scene)
    c = toplevel_collection_to_layer_collection(c, 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}


if __name__ == "__main__":
    file_name = "playground.blend"
    blend_file_path = str((Path(__file__) / ".." / file_name).resolve())
    slide = Slide(blend_file_path)