2024-03-11 16:35:41 +01:00
|
|
|
import typing as typ
|
|
|
|
import typing_extensions as typx
|
|
|
|
|
|
|
|
import pydantic as pyd
|
|
|
|
import sympy as sp
|
|
|
|
import sympy.physics.units as spu
|
|
|
|
|
|
|
|
import bpy
|
|
|
|
|
|
|
|
from ...utils import extra_sympy_units as spuex
|
|
|
|
from . import contracts as ct
|
|
|
|
from .contracts import SocketType as ST
|
|
|
|
from . import sockets as sck
|
|
|
|
|
|
|
|
# TODO: Caching?
|
|
|
|
# TODO: Move the manual labor stuff to contracts
|
|
|
|
|
|
|
|
BLSocketType = str ## A Blender-Defined Socket Type
|
|
|
|
BLSocketSize = int
|
|
|
|
DescType = str
|
|
|
|
Unit = typ.Any ## Type of a valid unit
|
|
|
|
|
|
|
|
####################
|
|
|
|
# - Socket to SocketDef
|
|
|
|
####################
|
|
|
|
SOCKET_DEFS = {
|
|
|
|
socket_type: getattr(
|
|
|
|
sck,
|
2024-03-20 12:56:50 +01:00
|
|
|
socket_type.value.removesuffix('SocketType') + 'SocketDef',
|
2024-03-11 16:35:41 +01:00
|
|
|
)
|
|
|
|
for socket_type in ST
|
2024-03-20 12:56:50 +01:00
|
|
|
if hasattr(sck, socket_type.value.removesuffix('SocketType') + 'SocketDef')
|
2024-03-11 16:35:41 +01:00
|
|
|
}
|
|
|
|
## TODO: Bit of a hack. Is it robust enough?
|
|
|
|
|
|
|
|
for socket_type in ST:
|
|
|
|
if not hasattr(
|
|
|
|
sck,
|
2024-03-20 12:56:50 +01:00
|
|
|
socket_type.value.removesuffix('SocketType') + 'SocketDef',
|
2024-03-11 16:35:41 +01:00
|
|
|
):
|
2024-03-20 12:56:50 +01:00
|
|
|
print('Missing SocketDef for', socket_type.value)
|
2024-03-11 16:35:41 +01:00
|
|
|
|
|
|
|
|
|
|
|
####################
|
|
|
|
# - BL Socket Size Parser
|
|
|
|
####################
|
|
|
|
BL_SOCKET_3D_TYPE_PREFIXES = {
|
2024-03-20 12:56:50 +01:00
|
|
|
'NodeSocketVector',
|
|
|
|
'NodeSocketRotation',
|
2024-03-11 16:35:41 +01:00
|
|
|
}
|
|
|
|
BL_SOCKET_4D_TYPE_PREFIXES = {
|
2024-03-20 12:56:50 +01:00
|
|
|
'NodeSocketColor',
|
2024-03-11 16:35:41 +01:00
|
|
|
}
|
2024-03-20 12:56:50 +01:00
|
|
|
|
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
def size_from_bl_interface_socket(
|
2024-03-20 12:56:50 +01:00
|
|
|
bl_interface_socket: bpy.types.NodeTreeInterfaceSocket,
|
2024-03-11 16:35:41 +01:00
|
|
|
) -> typx.Literal[1, 2, 3, 4]:
|
|
|
|
"""Parses the `size`, aka. number of elements, contained within the `default_value` of a Blender interface socket.
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
Since there are no 2D sockets in Blender, the user can specify "2D" in the Blender socket's description to "promise" that only the first two values will be used.
|
|
|
|
When this is done, the third value is left entirely untouched by this entire system.
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
A hard-coded set of NodeSocket<Type> prefixes are used to determine which interface sockets are, in fact, 3D.
|
|
|
|
- For 3D sockets, a hard-coded list of Blender node socket types is used.
|
|
|
|
- Else, it is a 1D socket type.
|
|
|
|
"""
|
2024-03-20 12:56:50 +01:00
|
|
|
if bl_interface_socket.description.startswith('2D'):
|
|
|
|
return 2
|
2024-03-11 16:35:41 +01:00
|
|
|
if any(
|
|
|
|
bl_interface_socket.socket_type.startswith(bl_socket_3d_type_prefix)
|
|
|
|
for bl_socket_3d_type_prefix in BL_SOCKET_3D_TYPE_PREFIXES
|
|
|
|
):
|
|
|
|
return 3
|
|
|
|
if any(
|
|
|
|
bl_interface_socket.socket_type.startswith(bl_socket_4d_type_prefix)
|
|
|
|
for bl_socket_4d_type_prefix in BL_SOCKET_4D_TYPE_PREFIXES
|
|
|
|
):
|
|
|
|
return 4
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
return 1
|
|
|
|
|
|
|
|
|
|
|
|
####################
|
|
|
|
# - BL Socket Type / Unit Parser
|
|
|
|
####################
|
|
|
|
def parse_bl_interface_socket(
|
|
|
|
bl_interface_socket: bpy.types.NodeTreeInterfaceSocket,
|
|
|
|
) -> tuple[ST, sp.Expr | None]:
|
|
|
|
"""Parse a Blender interface socket by parsing its description, falling back to any direct type links.
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
Arguments:
|
|
|
|
bl_interface_socket: An interface socket associated with the global input to a node tree.
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
Returns:
|
|
|
|
The type of a corresponding MaxwellSimSocket, as well as a unit (if a particular unit was requested by the Blender interface socket).
|
|
|
|
"""
|
|
|
|
size = size_from_bl_interface_socket(bl_interface_socket)
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
# Determine Direct Socket Type
|
|
|
|
if (
|
|
|
|
direct_socket_type := ct.BL_SOCKET_DIRECT_TYPE_MAP.get(
|
|
|
|
(bl_interface_socket.socket_type, size)
|
|
|
|
)
|
|
|
|
) is None:
|
|
|
|
msg = "Blender interface socket has no mapping among 'MaxwellSimSocket's."
|
|
|
|
raise ValueError(msg)
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
# (Maybe) Return Direct Socket Type
|
|
|
|
## When there's no description, that's it; return.
|
|
|
|
if not ct.BL_SOCKET_DESCR_ANNOT_STRING in bl_interface_socket.description:
|
|
|
|
return (direct_socket_type, None)
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
# Parse Description for Socket Type
|
|
|
|
tokens = (
|
|
|
|
_tokens
|
2024-03-20 12:56:50 +01:00
|
|
|
if (_tokens := bl_interface_socket.description.split(' '))[0] != '2D'
|
2024-03-11 16:35:41 +01:00
|
|
|
else _tokens[1:]
|
2024-03-20 12:56:50 +01:00
|
|
|
) ## Don't include the "2D" token, if defined.
|
2024-03-11 16:35:41 +01:00
|
|
|
if (
|
|
|
|
socket_type := ct.BL_SOCKET_DESCR_TYPE_MAP.get(
|
|
|
|
(tokens[0], bl_interface_socket.socket_type, size)
|
|
|
|
)
|
|
|
|
) is None:
|
2024-03-20 12:56:50 +01:00
|
|
|
return (
|
|
|
|
direct_socket_type,
|
|
|
|
None,
|
|
|
|
) ## Description doesn't map to anything
|
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
# Determine Socket Unit (to use instead of "unit system")
|
|
|
|
## This is entirely OPTIONAL
|
|
|
|
socket_unit = None
|
|
|
|
if socket_type in ct.SOCKET_UNITS:
|
|
|
|
## Case: Unit is User-Defined
|
2024-03-20 12:56:50 +01:00
|
|
|
if len(tokens) > 1 and '(' in tokens[1] and ')' in tokens[1]:
|
2024-03-11 16:35:41 +01:00
|
|
|
# Compute (<unit_str>) as Unit Token
|
2024-03-20 12:56:50 +01:00
|
|
|
unit_token = tokens[1].removeprefix('(').removesuffix(')')
|
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
# Compare Unit Token to Valid Sympy-Printed Units
|
2024-03-20 12:56:50 +01:00
|
|
|
socket_unit = (
|
|
|
|
_socket_unit
|
|
|
|
if (
|
|
|
|
_socket_unit := [
|
|
|
|
unit
|
|
|
|
for unit in ct.SOCKET_UNITS[socket_type][
|
|
|
|
'values'
|
|
|
|
].values()
|
|
|
|
if str(unit) == unit_token
|
|
|
|
]
|
|
|
|
)
|
|
|
|
else ct.SOCKET_UNITS[socket_type]['values'][
|
|
|
|
ct.SOCKET_UNITS[socket_type]['default']
|
|
|
|
]
|
|
|
|
)
|
2024-03-11 16:35:41 +01:00
|
|
|
## TODO: Enforce abbreviated sympy printing here, not globally
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
return (socket_type, socket_unit)
|
|
|
|
|
|
|
|
|
|
|
|
####################
|
|
|
|
# - BL Socket Interface Definition
|
|
|
|
####################
|
|
|
|
def socket_def_from_bl_interface_socket(
|
|
|
|
bl_interface_socket: bpy.types.NodeTreeInterfaceSocket,
|
|
|
|
):
|
2024-03-20 12:56:50 +01:00
|
|
|
"""Computes an appropriate (no-arg) SocketDef from the given `bl_interface_socket`, by parsing it."""
|
|
|
|
return SOCKET_DEFS[parse_bl_interface_socket(bl_interface_socket)[0]]
|
2024-03-11 16:35:41 +01:00
|
|
|
|
|
|
|
|
|
|
|
####################
|
|
|
|
# - Extract Default Interface Socket Value
|
|
|
|
####################
|
|
|
|
def value_from_bl(
|
|
|
|
bl_interface_socket: bpy.types.NodeTreeInterfaceSocket,
|
|
|
|
unit_system: dict | None = None,
|
|
|
|
) -> typ.Any:
|
|
|
|
"""Reads the value of any Blender socket, and writes its `default_value` to the `value` of any `MaxwellSimSocket`.
|
|
|
|
- If the size of the Blender socket is >1, then `value` is written to as a `sympy.Matrix`.
|
|
|
|
- If a unit system is given, then the Blender socket is matched to a `MaxwellSimSocket`, which is used to lookup an appropriate unit in the given `unit_system`.
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
"""
|
|
|
|
## TODO: Consider sympy.S()'ing the default_value
|
|
|
|
parsed_bl_socket_value = {
|
|
|
|
1: lambda: bl_interface_socket.default_value,
|
|
|
|
2: lambda: sp.Matrix(tuple(bl_interface_socket.default_value)[:2]),
|
|
|
|
3: lambda: sp.Matrix(tuple(bl_interface_socket.default_value)),
|
|
|
|
4: lambda: sp.Matrix(tuple(bl_interface_socket.default_value)),
|
|
|
|
}[size_from_bl_interface_socket(bl_interface_socket)]()
|
|
|
|
## The 'lambda' delays construction until size is determined
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
socket_type, unit = parse_bl_interface_socket(bl_interface_socket)
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
# Add Unit to Parsed (if relevant)
|
|
|
|
if unit is not None:
|
|
|
|
parsed_bl_socket_value *= unit
|
|
|
|
elif unit_system is not None:
|
|
|
|
parsed_bl_socket_value *= unit_system[socket_type]
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
return parsed_bl_socket_value
|
|
|
|
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
####################
|
|
|
|
# - Convert to Blender-Compatible Value
|
|
|
|
####################
|
|
|
|
def make_scalar_bl_compat(scalar: typ.Any) -> typ.Any:
|
|
|
|
"""Blender doesn't accept ex. Sympy numbers as values.
|
|
|
|
Therefore, we need to do some conforming.
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
Currently hard-coded; this is probably best.
|
|
|
|
"""
|
|
|
|
if isinstance(scalar, sp.Integer):
|
|
|
|
return int(scalar)
|
|
|
|
elif isinstance(scalar, sp.Float):
|
|
|
|
return float(scalar)
|
|
|
|
elif isinstance(scalar, sp.Rational):
|
|
|
|
return float(scalar)
|
|
|
|
elif isinstance(scalar, sp.Expr):
|
|
|
|
return float(scalar.n())
|
|
|
|
## TODO: More?
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
return scalar
|
|
|
|
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
def value_to_bl(
|
|
|
|
bl_interface_socket: bpy.types.NodeSocket,
|
|
|
|
value: typ.Any,
|
|
|
|
unit_system: dict | None = None,
|
|
|
|
) -> typ.Any:
|
|
|
|
socket_type, unit = parse_bl_interface_socket(bl_interface_socket)
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
# Set Socket
|
|
|
|
if unit is not None:
|
|
|
|
bl_socket_value = spu.convert_to(value, unit) / unit
|
2024-03-20 12:56:50 +01:00
|
|
|
elif unit_system is not None and socket_type in unit_system:
|
|
|
|
bl_socket_value = (
|
|
|
|
spu.convert_to(value, unit_system[socket_type])
|
|
|
|
/ unit_system[socket_type]
|
|
|
|
)
|
2024-03-11 16:35:41 +01:00
|
|
|
else:
|
|
|
|
bl_socket_value = value
|
2024-03-20 12:56:50 +01:00
|
|
|
|
2024-03-11 16:35:41 +01:00
|
|
|
return {
|
|
|
|
1: lambda: make_scalar_bl_compat(bl_socket_value),
|
2024-03-20 12:56:50 +01:00
|
|
|
2: lambda: tuple(
|
|
|
|
[
|
|
|
|
make_scalar_bl_compat(bl_socket_value[0]),
|
|
|
|
make_scalar_bl_compat(bl_socket_value[1]),
|
|
|
|
bl_interface_socket.default_value[2],
|
|
|
|
## Don't touch (unused) 3rd bl_socket coordinate
|
|
|
|
]
|
|
|
|
),
|
|
|
|
3: lambda: tuple(
|
|
|
|
[make_scalar_bl_compat(el) for el in bl_socket_value]
|
|
|
|
),
|
|
|
|
4: lambda: tuple(
|
|
|
|
[make_scalar_bl_compat(el) for el in bl_socket_value]
|
|
|
|
),
|
2024-03-11 16:35:41 +01:00
|
|
|
}[size_from_bl_interface_socket(bl_interface_socket)]()
|
|
|
|
## The 'lambda' delays construction until size is determined
|