145 lines
3.7 KiB
Python
145 lines
3.7 KiB
Python
# encoding: utf-8
|
|
"""
|
|
"""
|
|
|
|
from dataclasses import dataclass as _dataclass, fields as __dataclass_fields
|
|
|
|
from comfy_api.latest import io as _io
|
|
|
|
from .__meta import category_old as _category_old
|
|
from .__typing import _t, T as _T, _A, _O, _U, _TypeVar
|
|
|
|
|
|
_DICT = _io.Custom("DICT")
|
|
_DICT_INPUT_OPTIONAL = _DICT.Input(
|
|
'dict',
|
|
optional=True,
|
|
tooltip="An (optional) Dictionary to work with."
|
|
)
|
|
_DICT_OUTPUT = _DICT.Output(
|
|
'DICT',
|
|
display_name='dict'
|
|
)
|
|
_KEY_INPUT_ADD = _io.String.Input(
|
|
'key',
|
|
tooltip="Key (name) of the item inserted into the dict.",
|
|
)
|
|
|
|
_T_Input = _t.TypeVar('T_Input', bound=_io.Input)
|
|
|
|
|
|
# ----------------------------------------------------------
|
|
|
|
|
|
@_dataclass
|
|
class _InputsConverter:
|
|
"""
|
|
Config-like class used to rebuild schema's inputs for old version of a node.
|
|
|
|
- Only the inputs with IDs in ``preserved`` set are kept (or all, if empty).
|
|
- Then, they're optionally renamed according to ``renames`` dict.
|
|
"""
|
|
preserved: _O[_t.Sequence[str]] = None
|
|
renames: _O[_t.Dict[str, str]] = None
|
|
|
|
def convert(self, inputs: _t.Sequence[_T_Input]) -> _t.List[_T_Input]:
|
|
preserved = self.preserved
|
|
renames = self.renames
|
|
if not(preserved or renames):
|
|
return list(inputs)
|
|
|
|
def _out_item_with_renames(inp: _T_Input) -> _T_Input:
|
|
_id = inp.id
|
|
if _id in renames:
|
|
return _input_override(inp, id=renames[_id])
|
|
return inp
|
|
|
|
def _out_item_as_is(inp: _T_Input) -> _T_Input:
|
|
return inp
|
|
|
|
out_item = _out_item_with_renames if renames else _out_item_as_is
|
|
|
|
if not preserved:
|
|
# We keep all the inputs in the original order:
|
|
return [out_item(x) for x in inputs]
|
|
|
|
# Instead, we filter and possibly reorder them:
|
|
orig_inputs_dict = {x.id: x for x in inputs}
|
|
return [out_item(orig_inputs_dict[_id]) for _id in preserved]
|
|
|
|
|
|
# ----------------------------------------------------------
|
|
|
|
|
|
def __dataclass_fields_dict(obj) -> _t.Dict[str, _t.Any]:
|
|
"""Return all the field values from a dataclass - as dict. No deep-copy."""
|
|
field_names = (field.name for field in __dataclass_fields(obj))
|
|
return {
|
|
attr: getattr(obj, attr)
|
|
for attr in field_names
|
|
}
|
|
|
|
|
|
def _dataclass_override(obj: _T, **overrides) -> _T:
|
|
"""Copy a dataclass instance with optional overrides."""
|
|
obj_dict = __dataclass_fields_dict(obj)
|
|
obj_dict.update(overrides)
|
|
obj_type = type(obj)
|
|
return obj_type(**obj_dict)
|
|
|
|
|
|
def _input_override(input_obj: _T_Input, **overrides) -> _T_Input:
|
|
"""Copy an ``io.Input`` instance with optional overrides."""
|
|
new_id = input_obj.id
|
|
if 'id' in overrides:
|
|
new_id = overrides.pop('id')
|
|
|
|
input_args = input_obj.as_dict()
|
|
if 'id' in input_args:
|
|
del input_args['id']
|
|
input_args.update(overrides)
|
|
|
|
input_type: _t.Type[_T_Input] = type(input_obj)
|
|
return input_type(new_id, **input_args)
|
|
|
|
|
|
def _schema_old_node(
|
|
schema: _io.Schema, node_id: str,
|
|
inputs_converter: _InputsConverter = None,
|
|
**overrides
|
|
):
|
|
"""Build a deprecated schema for old node-ID."""
|
|
if inputs_converter is not None:
|
|
inputs: _t.List[_T_Input] = overrides.get('inputs', schema.inputs)
|
|
inputs = inputs_converter.convert(inputs)
|
|
overrides['inputs'] = inputs
|
|
|
|
schema_dict = __dataclass_fields_dict(schema)
|
|
if 'node_id' in schema_dict:
|
|
del schema_dict['node_id']
|
|
name_old = schema_dict.pop('display_name')
|
|
|
|
schema_dict.update(dict(
|
|
category=_category_old,
|
|
display_name=f"{name_old} [OLD]",
|
|
is_deprecated=True,
|
|
is_dev_only=True,
|
|
))
|
|
schema_dict.update(overrides)
|
|
return _io.Schema(node_id, **schema_dict)
|
|
|
|
|
|
# noinspection PyAbstractClass
|
|
class _BaseNode(_io.ComfyNode):
|
|
"""Shared base class for all the nodes in pack.
|
|
Stores schema in internal class attribute.
|
|
"""
|
|
_schema: _io.Schema = None
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> _io.Schema:
|
|
return cls._schema
|
|
|
|
|
|
_T_BaseNode = _TypeVar('T_BaseNode', bound=_BaseNode)
|