import logging
from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Generic, NamedTuple, TypeAlias, TypeVar
import bpy
import mathutils
from procfunc.util.pytree import register_pytree_container
T = TypeVar("T")
logger = logging.getLogger(__name__)
__all__ = [
"Asset",
"BlenderAsset",
"ObjectType",
"Object",
"CameraObject",
"MeshObject",
"CurveObject",
"EmptyObject",
"ArmatureObject",
"HairObject",
"LatticeObject",
"LightObject",
"LightProbeObject",
"MetaObject",
"Material",
"Texture",
"Image",
"Collection",
"VolumeObject",
"PointCloudObject",
"World",
"ValueRange",
]
INVALIDATED_USAGE_ID = -100
@dataclass
class AssetUsageTable:
counts: dict[int, int] = field(default_factory=lambda: defaultdict(int))
_global_usage_table = AssetUsageTable()
[docs]
class Asset(Generic[T]):
[docs]
def item(self) -> T:
raise NotImplementedError("Subclasses of Asset must implement item()")
[docs]
class BlenderAsset(Asset, Generic[T]):
"""
A pythonic wrapper around a blender object,material,texture, or other bpy.data.stuffgoeshere member.
The underlying bpy datablock is owned by Blender and is not deleted when this wrapper is garbage collected.
"""
[docs]
def __init__(self, item: T):
if isinstance(item, Asset):
raise ValueError(
f"Attempted to wrap {item=} in a second {self.__class__.__name__} wrapper. This is almost certainly not intended"
)
self._item = item
_global_usage_table.counts[id(item)] += 1
if logger.isEnabledFor(logging.DEBUG):
logger.debug(
f"Asset ref {id(self)} created for {self._item.name}. remaining refs: {_global_usage_table.counts[id(self._item)]}"
)
# self._dependencies: list[Asset] = []
[docs]
def item(self) -> T:
if _global_usage_table.counts[id(self._item)] == INVALIDATED_USAGE_ID:
raise ValueError(
"Cant access .item() since asset has been explicitly invalidated by a previous operation"
)
return self._item
[docs]
class ObjectType(Enum):
ARMATURE = "ARMATURE"
CAMERA = "CAMERA"
CURVE = "CURVE"
EMPTY = "EMPTY"
FONT = "FONT"
HAIR = "HAIR"
LATTICE = "LATTICE"
LIGHT = "LIGHT"
LIGHT_PROBE = "LIGHT_PROBE"
MESH = "MESH"
META = "META"
POINTCLOUD = "POINTCLOUD"
SURFACE = "SURFACE"
VOLUME = "VOLUME"
# GPENCIL = "GPENCIL"
# GREASEPENCIL = "GREASEPENCIL"
# SPEAKER = "SPEAKER"
Vector: TypeAlias = mathutils.Vector
Color: TypeAlias = mathutils.Color
Euler: TypeAlias = mathutils.Euler
Quaternion: TypeAlias = mathutils.Quaternion
Matrix: TypeAlias = mathutils.Matrix
BVHTree: TypeAlias = mathutils.bvhtree.BVHTree
NodeGroup: TypeAlias = bpy.types.NodeGroup
Scene: TypeAlias = bpy.types.Scene
ViewLayer: TypeAlias = bpy.types.ViewLayer
[docs]
class Object(BlenderAsset[bpy.types.Object]):
[docs]
def __init__(self, obj: bpy.types.Object):
assert isinstance(obj, bpy.types.Object)
super().__init__(obj)
def __repr__(self):
return f"pf.{self.__class__.__name__}(bpy.data.objects[{self._item.name!r}])"
[docs]
def clone(self):
new_obj = self._item.copy()
new_obj.data = self._item.data.copy()
bpy.context.collection.objects.link(new_obj)
return self.__class__(new_obj)
[docs]
class CameraObject(Object):
[docs]
def __init__(self, obj: bpy.types.Object):
assert obj.type == ObjectType.CAMERA.value
super().__init__(obj)
[docs]
class MeshObject(Object):
[docs]
def __init__(self, obj: bpy.types.Object):
assert obj.type == ObjectType.MESH.value
super().__init__(obj)
if len(obj.children) > 0:
raise ValueError(
f"MeshObject {obj.name} had children {obj.children}, but this is not allowed for {self.__class__.__name__}"
)
if obj.parent is not None:
logger.warning(f"MeshObject {obj.name} had a parent {obj.parent}")
[docs]
class CurveObject(Object):
[docs]
def __init__(self, obj: bpy.types.Object):
assert obj.type == ObjectType.CURVE.value
super().__init__(obj)
[docs]
class EmptyObject(Object):
[docs]
def __init__(self, obj: bpy.types.Object):
assert obj.type == ObjectType.EMPTY.value
super().__init__(obj)
[docs]
class ArmatureObject(Object):
[docs]
def __init__(self, obj: bpy.types.Object):
assert obj.type == ObjectType.ARMATURE.value
super().__init__(obj)
[docs]
class HairObject(Object):
[docs]
def __init__(self, obj: bpy.types.Object):
assert obj.type == ObjectType.HAIR.value
super().__init__(obj)
[docs]
class LatticeObject(Object):
[docs]
def __init__(self, obj: bpy.types.Object):
assert obj.type == ObjectType.LATTICE.value
super().__init__(obj)
[docs]
class LightObject(Object):
[docs]
def __init__(self, obj: bpy.types.Object):
assert obj.type == ObjectType.LIGHT.value
super().__init__(obj)
[docs]
class LightProbeObject(Object):
[docs]
def __init__(self, obj: bpy.types.Object):
assert obj.type == ObjectType.LIGHT_PROBE.value
super().__init__(obj)
def is_zero_displacement(value: Any) -> bool:
"""True for a constant zero-vector displacement node.
Such a displacement has no effect, so the material build can leave the
output socket disconnected and skip emitting the subgraph entirely.
"""
from procfunc import compute_graph as cg
from procfunc.nodes import types as nt
if isinstance(value, nt.ProcNode):
value = value.item()
if not isinstance(value, cg.ConstantNode):
return False
return isinstance(value.value, (mathutils.Vector, tuple, list)) and tuple(
value.value
) == (0.0, 0.0, 0.0)
[docs]
@dataclass
class Material:
surface: Any = None
displacement: Any = None
volume: Any = None
_bpy_material: Any = field(default=None, init=False, repr=False)
[docs]
def item(self) -> bpy.types.Material:
if self._bpy_material is None:
from procfunc.nodes.execute.realize import build_bpy_material
self._bpy_material = build_bpy_material(
surface=self.surface,
displacement=self.displacement,
volume=self.volume,
)
return self._bpy_material
[docs]
class Texture(BlenderAsset[bpy.types.Texture]):
[docs]
def __init__(self, tex: bpy.types.Texture):
assert isinstance(tex, bpy.types.Texture)
super().__init__(tex)
[docs]
class Image(BlenderAsset[bpy.types.Image]):
[docs]
def __init__(self, img: bpy.types.Image):
assert isinstance(img, bpy.types.Image)
super().__init__(img)
[docs]
class Collection:
[docs]
def __init__(
self, objects: "list[Object] | bpy.types.Collection", name: str = "collection"
):
if isinstance(objects, bpy.types.Collection):
self._collection = objects
self._objects = [Object(o) for o in objects.objects]
else:
self._objects = objects
self._collection = bpy.data.collections.new(name=name)
for obj in objects:
self._collection.objects.link(obj.item())
def __repr__(self):
return f"pf.Collection(bpy.data.collections[{self._collection.name!r}])"
[docs]
def map(
self,
fn: Callable[[Object], Object],
skip_none: bool = True,
) -> "Collection":
objs = []
for obj in self._objects:
result = fn(obj)
if result is None and skip_none:
continue
objs.append(result)
return Collection(objs, name=self._collection.name + "_" + fn.__name__)
def __len__(self):
return len(self._objects)
def __iter__(self):
return iter(self._objects)
[docs]
def item(self) -> bpy.types.Collection:
return self._collection
[docs]
class VolumeObject(BlenderAsset[bpy.types.Volume]):
[docs]
def __init__(self, vol: bpy.types.Volume):
assert isinstance(vol, bpy.types.Volume)
super().__init__(vol)
[docs]
class PointCloudObject(BlenderAsset[bpy.types.PointCloud]):
[docs]
def __init__(self, vol: bpy.types.PointCloud):
assert isinstance(vol, bpy.types.PointCloud)
super().__init__(vol)
[docs]
class World(BlenderAsset[bpy.types.World]):
[docs]
def __init__(self, world: bpy.types.World):
assert isinstance(world, bpy.types.World)
super().__init__(world)
def __del__(self):
pass # no cleanup for these
# TODO Font, PointCloud, Surface, Volume, Scene? World? ViewLayer?
TRangeType = TypeVar("TRangeType")
[docs]
class ValueRange(NamedTuple, Generic[TRangeType]):
min: TRangeType | None
max: TRangeType | None
register_pytree_container(
Material,
flatten_func=lambda m: ([m.surface, m.displacement, m.volume], None),
unflatten_func=lambda vals, _: Material(
surface=vals[0], displacement=vals[1], volume=vals[2]
),
names_func=lambda m: ["surface", "displacement", "volume"],
)