import logging
from collections import OrderedDict
from typing import Any
import bpy
from procfunc import compute_graph as cg
from procfunc import context
from procfunc import types as pt
from procfunc.nodes import geo
from procfunc.nodes import types as nt
from procfunc.nodes.geo import store_named_attribute
from procfunc.nodes.util.bpy_node_info import (
NodeGroupType,
)
from procfunc.ops._util import modify
from procfunc.ops.object import mesh_to_curve as ops_mesh_to_curve
from procfunc.ops.primitives.mesh import mesh_single_vertex
from procfunc.tracer import primitive as tracer_primitive
from procfunc.util import pytree
from procfunc.util.bpy_info import bpy_nocollide_data_name
from .construct_nodes import (
as_nodegroup,
construct_procnode_to_bpy,
)
from .realize import (
extract_geometry_singlekey,
nodegroup_to_output,
unwrap_and_drop,
)
logger = logging.getLogger(__name__)
[docs]
@tracer_primitive
def to_environment(
surface: nt.ProcNode[nt.Shader] | None = None,
volume: nt.ProcNode[nt.Shader] | None = None,
) -> pt.World:
# Get or create world
if bpy.context.scene.world is None:
world = bpy.data.worlds.new("World")
bpy.context.scene.world = world
else:
world = bpy.context.scene.world
# Enable use of nodes for the world
world.use_nodes = True
world.node_tree.nodes.clear()
outputs = unwrap_and_drop({"surface": surface, "volume": volume})
graph = cg.ComputeGraph(
inputs=pytree.PyTree({}),
outputs=pytree.PyTree(outputs),
name="to_environment",
metadata={},
)
body = as_nodegroup(graph, NodeGroupType.SHADER)
nodegroup_to_output(
world.node_tree, body, "ShaderNodeOutputWorld", list(outputs.keys())
)
return pt.World(world)
[docs]
@tracer_primitive
def to_light(
light: pt.LightObject,
surface: nt.ProcNode[nt.Shader],
) -> pt.LightObject:
"""Apply a shader node graph to a light's internal node tree."""
lamp_data = light.item().data
lamp_data.use_nodes = True
lamp_data.node_tree.nodes.clear()
outputs = unwrap_and_drop({"surface": surface})
graph = cg.ComputeGraph(
inputs=pytree.PyTree({}),
outputs=pytree.PyTree(outputs),
name="to_light",
metadata={},
)
body = as_nodegroup(graph, NodeGroupType.SHADER)
nodegroup_to_output(
lamp_data.node_tree, body, "ShaderNodeOutputLight", list(outputs.keys())
)
return light
[docs]
@tracer_primitive
def to_compositor(
results: dict[str, nt.ProcNode],
):
bpy.context.scene.use_nodes = True
nt = bpy.context.scene.node_tree
nt.nodes.clear()
cache = {}
for k, v in results.items():
construct_procnode_to_bpy(v.item(), nt, cache)
return bpy.context.scene
[docs]
@tracer_primitive
def to_mesh_object(
geometry: nt.ProcNode[pt.MeshObject],
_skip_apply: bool = False,
) -> pt.MeshObject:
graph = cg.ComputeGraph(
inputs=pytree.PyTree({}),
outputs=pytree.PyTree({"geometry": geometry.item()}),
name="to_mesh_object",
metadata={},
)
ng = as_nodegroup(graph, NodeGroupType.GEOMETRY)
ng_inouts = ng.interface.items_tree
item = ng_inouts.get("geometry")
if item is None:
raise ValueError(
f"Node group {ng.name} has no output geometry, available are {ng_inouts.keys()}"
)
if item.in_out != "OUTPUT":
raise ValueError(
f"Node group attempted to extract geometry but it wasnt an output! {item.in_out=}"
)
try:
obj_result = extract_geometry_singlekey(
ng,
"geometry",
attribute_keys=[],
realize=True,
_skip_apply=_skip_apply,
)
except RuntimeError as e:
if "does not contain a mesh" not in str(e):
raise
mode = context.globals.warn_mode_empty_geonodes
if mode == "throw":
raise
if mode == "warn":
logger.warning(
f"to_mesh_object produced no mesh geometry, returning empty mesh: {e}"
)
return mesh_single_vertex()
assert isinstance(obj_result, pt.MeshObject)
return obj_result
[docs]
@tracer_primitive
def to_mesh_object_with_attributes(
geometry: nt.ProcNode[pt.MeshObject],
attributes: dict[str, nt.ProcNode[nt.AnyDataVal]] | None = None,
) -> tuple[pt.MeshObject, dict[str, Any]]:
if attributes is None:
attributes = {}
for k, v in attributes.items():
geometry = store_named_attribute(geometry, name=k, value=v)
obj = to_mesh_object(geometry)
attr_dict = {k: obj.item().data.attributes[k] for k in attributes.keys()}
return obj, attr_dict
[docs]
@tracer_primitive
def to_curve_object(
geometry: nt.ProcNode[pt.CurveObject],
) -> pt.CurveObject:
"""
WARNING: currently discards any bezier config or knots, only extracts the points.Any
TODO: Need a better geonodes -> curve op from blender, or need to engineer this in via attribute extraction.
"""
obj = to_mesh_object(geo.curve_to_mesh(geometry, profile_curve=None))
return ops_mesh_to_curve(obj)
[docs]
@tracer_primitive
def to_objects_multi(
geometries: dict[str, nt.ProcNode[pt.MeshObject]],
attributes: dict[str, dict[str, nt.ProcNode[nt.AnyDataVal]]] | None = None,
) -> OrderedDict[str, pt.MeshObject]:
"""
Convert a nodegroup which has multiple output geometries into multiple realized objects.
If the objects should have any attributes, provide node definitions for them in the 'attributes' argument.
Args:
geometries: named output geometry nodes to be converted into objects
attributes: named data attributes which should be annotated on those objects. keys must be a subset of the keys of 'geometries'
input_obj: optional input object to use for the geometry nodegroups
Returns:
- result_objects: dict of {output_key: object}. output keys are the same as the keys of 'geometries'
"""
if attributes is None:
attributes = {}
extra_attr_keys = set(attributes.keys()) - set(geometries.keys())
if extra_attr_keys:
raise ValueError(
f"{to_objects_multi.__name__} got {attributes.keys()=} but these are not a subset of {geometries.keys()=} due to {extra_attr_keys=}"
)
attr_keys_dedup = {
kobj: {f"{kobj}_{kattr}": v for kattr, v in attrs.items()}
for kobj, attrs in attributes.items()
}
all_ng_outputs: dict[str, nt.ProcNode] = geometries.copy()
for kobj, attrs in attr_keys_dedup.items():
all_ng_outputs.update(attrs)
all_ng_outputs = unwrap_and_drop(all_ng_outputs) # type: ignore
graph = cg.ComputeGraph(
inputs=pytree.PyTree({}),
outputs=pytree.PyTree(all_ng_outputs),
name="to_objects_multi",
metadata={},
)
nodegroup = as_nodegroup(
graph,
NodeGroupType.GEOMETRY,
)
ng_inouts = nodegroup.interface.items_tree
result_objects = OrderedDict()
for k, v in geometries.items():
item = ng_inouts.get(k)
if item is None:
raise ValueError(
f"Node group {nodegroup.name} has no output {k}, available are {ng_inouts.keys()}"
)
if item.in_out != "OUTPUT":
raise ValueError(
f"Node group attempted to extract {k=} but it wasnt an output! {item.in_out=}"
)
result_objects[k] = extract_geometry_singlekey(
nodegroup,
k,
attribute_keys=attr_keys_dedup.get(k, {}).keys(),
)
return result_objects
[docs]
@tracer_primitive
def to_aliases(
geometry: nt.ProcNode[pt.MeshObject],
) -> list[pt.MeshObject]:
"""
Convert instanced geometry into aliases - separate objects sharing the same mesh data.
Uses depsgraph to extract instance transforms and creates one bpy.data.object per instance.
All aliases point directly to the original mesh data from the scene - no copying occurs.
Each alias has its own transform (position, rotation, scale) but shares mesh data.
Note: Requires geometry nodes that use actual bpy.data.collections with real objects
(via collection_info node), not just joined geometry. The instances must be visible
in the viewport (show_viewport=True) to be detected by the depsgraph.
Args:
geometry: node producing instanced geometry (e.g., from instance_on_points)
Returns:
list of MeshObject, one per instance, with each instance pointing to the original
mesh data from the scene
"""
graph = cg.ComputeGraph(
inputs=pytree.PyTree({}),
outputs=pytree.PyTree({"geometry": geometry.item()}),
name="to_aliases",
metadata={},
)
nodegroup = as_nodegroup(graph, NodeGroupType.GEOMETRY)
obj_with_modifier = modify(
mesh_single_vertex(), "NODES", node_group=nodegroup, _skip_apply=True
)
temp_obj = obj_with_modifier.item()
for mod in temp_obj.modifiers:
mod.show_viewport = True
bpy.context.view_layer.update()
depsgraph = bpy.context.evaluated_depsgraph_get()
eval_mesh_to_instances = {}
eval_mesh_to_copy = {}
for deps_instance in depsgraph.object_instances:
if not deps_instance.is_instance:
continue
obj = deps_instance.object
if obj.type != "MESH":
continue
if deps_instance.parent is None or deps_instance.parent.original != temp_obj:
continue
eval_mesh = obj.data
if eval_mesh is None:
continue
mesh_id = eval_mesh.as_pointer()
if mesh_id not in eval_mesh_to_instances:
eval_mesh_to_instances[mesh_id] = []
eval_mesh_to_copy[mesh_id] = eval_mesh.copy()
eval_mesh_to_instances[mesh_id].append(deps_instance.matrix_world.copy())
bpy.data.objects.remove(temp_obj, do_unlink=True)
result_objects = []
for mesh_id, matrices in eval_mesh_to_instances.items():
mesh_data = eval_mesh_to_copy[mesh_id]
for matrix in matrices:
alias_obj = bpy.data.objects.new(
bpy_nocollide_data_name("alias", bpy.data.objects), mesh_data
)
alias_obj.matrix_world = matrix
bpy.context.collection.objects.link(alias_obj)
result_objects.append(pt.MeshObject(alias_obj))
return result_objects