oscillode/blender_maxwell/node_trees/maxwell_sim_nodes/sockets/base.py

220 lines
5.4 KiB
Python
Raw Normal View History

import typing as typ
import bpy
import sympy as sp
import sympy.physics.units as spu
from .. import contracts
class BLSocket(bpy.types.NodeSocket):
"""A base type for nodes that greatly simplifies the implementation of
reliable, powerful nodes.
Should be used together with `contracts.BLSocketProtocol`.
"""
def __init_subclass__(cls, **kwargs: typ.Any):
super().__init_subclass__(**kwargs) ## Yucky superclass setup.
# Set bl_idname
cls.bl_idname = cls.socket_type.value
cls.socket_color = contracts.SocketType_to_color[
cls.socket_type.value
]
# Configure Use of Units
if (
hasattr(cls, "use_units")
and cls.socket_type in contracts.SocketType_to_units
):
# Set Unit Properties
cls.__annotations__["raw_unit"] = bpy.props.EnumProperty(
name="Unit",
description="Choose a unit",
items=[
(unit_name, str(unit_value), str(unit_value))
for unit_name, unit_value in contracts.SocketType_to_units[
cls.socket_type
]["values"].items()
],
default=contracts.SocketType_to_units[
cls.socket_type
]["default"],
update=lambda self, context: self._update_unit(),
)
cls.__annotations__["raw_unit_previous"] = bpy.props.StringProperty(
default=contracts.SocketType_to_units[
cls.socket_type
]["default"]
)
# Declare Node Property: 'preset' EnumProperty
if hasattr(cls, "draw_preview"):
cls.__annotations__["preview_active"] = bpy.props.BoolProperty(
name="Preview",
description="Preview the socket value",
default=False,
)
####################
# - Internal Methods
####################
@property
def units(self) -> dict[str, sp.Expr]:
return contracts.SocketType_to_units[
self.socket_type
]["values"]
@property
def unit(self) -> sp.Expr:
return contracts.SocketType_to_units[
self.socket_type
]["values"][self.raw_unit]
@unit.setter
def unit(self, value) -> sp.Expr:
raw_unit_name = [
raw_unit_name
for raw_unit_name, unit_value in contracts.SocketType_to_units[
self.socket_type
]["values"].items()
if value == unit_value
][0]
self.raw_unit = raw_unit_name
@property
def _unit_previous(self) -> sp.Expr:
return contracts.SocketType_to_units[
self.socket_type
]["values"][self.raw_unit_previous]
@_unit_previous.setter
def _unit_previous(self, value) -> sp.Expr:
raw_unit_name = [
raw_unit_name
for raw_unit_name, unit_value in contracts.SocketType_to_units[
self.socket_type
]["values"].items()
if value == unit_value
][0]
self.raw_unit_previous = raw_unit_name
def value_as_unit(self, value) -> typ.Any:
"""Return the given value expresse as the current internal unit,
without the unit.
"""
if hasattr(self, "raw_value") and hasattr(self, "unit"):
# (Guard) Value Compatibility
if not self.is_compatible(value):
msg = f"Tried setting socket ({self}) to incompatible value ({value}) of type {type(value)}"
raise ValueError(msg)
# Return Converted Unit
return spu.convert_to(
value, self.unit
) / self.unit
else:
raise ValueError("Tried to get 'raw_value_as_unit', but class has no 'raw_value'")
def _update_unit(self) -> None:
"""Convert (if needed) the `raw_value` property, to use the unit
set in the `unit` property.
If the `raw_value` property isn't set, this only sets "unit_previous".
Run right after setting the `unit` property, in order to synchronize
the value with the new unit.
"""
if hasattr(self, "raw_value") and hasattr(self, "unit"):
self.raw_value = spu.convert_to(
self.raw_value * self._unit_previous,
self.unit,
) / self.unit
self._unit_previous = self.unit
####################
# - Methods
####################
def is_compatible(self, value: typ.Any) -> bool:
if not hasattr(self, "compatible_types"):
return True
for compatible_type, checks in self.compatible_types.items():
if (
compatible_type is typ.Any or
isinstance(value, compatible_type)
):
return all(check(self, value) for check in checks)
return False
####################
# - UI
####################
@classmethod
def draw_color_simple(cls) -> contracts.BlenderColorRGB:
return cls.socket_color
def draw(
self,
context: bpy.types.Context,
layout: bpy.types.UILayout,
node: bpy.types.Node,
text: str,
) -> None:
if self.is_output:
self.draw_output(context, layout, node, text)
else:
self.draw_input(context, layout, node, text)
def draw_input(
self,
context: bpy.types.Context,
layout: bpy.types.UILayout,
node: bpy.types.Node,
text: str,
) -> None:
if self.is_linked:
layout.label(text=text)
return
# Column
col = layout.column(align=True)
# Row: Label & Preview Toggle
label_col_row = col.row(align=True)
if hasattr(self, "draw_label_row"):
self.draw_label_row(label_col_row, text)
else:
label_col_row.label(text=text)
if hasattr(self, "draw_preview"):
label_col_row.prop(
self,
"preview_active",
toggle=True,
text="",
icon="SEQ_PREVIEW",
)
# Row: Preview (in Box)
if hasattr(self, "draw_preview"):
if self.preview_active:
col_box = col.box()
self.draw_preview(col_box)
# Row(s): Value
if hasattr(self, "draw_value"):
self.draw_value(col)
def draw_output(
self,
context: bpy.types.Context,
layout: bpy.types.UILayout,
node: bpy.types.Node,
text: str,
) -> None:
layout.label(text=text)