import inspect
import logging
import re
from pathlib import Path
from typing import Callable, List, get_type_hints, Dict, Any, Type
from simstack.core.simstack_result import SimstackResult
from simstack.core.context import context
from simstack.tables.node_children import update_node_children
from simstack.models import Parameters
from simstack.models.models import NodeModel, ModelMapping, DataMapping
from simstack.tables.table_builder import TableBuilderBase
from simstack.util.db import Database
from simstack.util.docstring_parser import DocstringParser
from simstack.util.importer import import_class_by_name
logger = logging.getLogger("NodeTable")
[docs]
def is_node_function(func: Callable[..., Any]) -> bool:
"""Check if a function is marked as a node using the @node decorator."""
return hasattr(func, "_is_node") and getattr(func, "_is_node", False) is True
[docs]
class CreateNodeTable(TableBuilderBase):
"""
Helper class to build the node table without passing around many parameters.
Usage:
creator = CreateNodeTable(database)
await creator.make_node_table()
"""
@property
def logger(self) -> logging.Logger:
return logger
[docs]
async def build(self, *args, **kwargs) -> None:
if not context.initialized:
await context.initialize()
await context.refresh_mappings(models=True, nodes=False)
await super().build(*args, **kwargs)
await context.refresh_mappings(models=False, nodes=True)
async def _process_module(self, module: Any, drops: str) -> None:
await self._register_nodes_from_module(module, drops)
def _discover_module_functions(
self, module: Any
) -> List[tuple[str, Callable[..., Any]]]:
"""
Return (name, func) for functions that are defined in `module` (not imported).
"""
functions: List[tuple[str, Callable[..., Any]]] = inspect.getmembers(
module, inspect.isfunction
)
module_name = module.__name__
return [
(func_name, func)
for func_name, func in functions
if func.__module__ == module_name
]
def _build_inputs(
self,
sig: inspect.Signature,
type_hints: Dict[str, Any],
doc_params: Dict[str, Any] | None,
) -> List[Dict[str, Any]]:
inputs: List[Dict[str, Any]] = []
for param_name, param in sig.parameters.items():
if param_name == "self": # Skip self parameter for methods
continue
try:
param_info: Dict[str, Any] = {
"name": param_name,
"type": type_hints.get(param_name, param.annotation.__name__),
"type_str": str(
type_hints.get(
param_name,
param.annotation.__name__
if param.annotation != inspect.Parameter.empty
else "Any",
)
),
}
except AttributeError as e:
logger.error(f"Could not parse type for {param_name}: {e}")
param_info = {}
if doc_params and param_name in doc_params:
param_info["description"] = doc_params[param_name].get("description")
if param.default != inspect.Parameter.empty:
param_info["default"] = param.default
inputs.append(param_info)
return inputs
def _parse_generic_type(self, type_str: str) -> tuple[str | None, str | None]:
"""
Parse generic types like List[X] or Dict[str,X] and extract the inner type.
Returns (wrapper, inner_type) where wrapper is 'List' or 'Dict' and inner_type is the model type.
"""
list_match = re.match(r"List\[(.*?)\]", type_str)
if list_match:
return "List", list_match.group(1)
dict_match = re.match(r"Dict\[str,\s*(.*?)\]", type_str)
if dict_match:
return "Dict", dict_match.group(1)
return None, None
@staticmethod
def _normalize_docstring_type(type_str: str) -> str:
"""Strip top-level docstring qualifiers like ', optional' from a parsed type string."""
normalized = (type_str or "").strip()
bracket_depth = 0
for index, char in enumerate(normalized):
if char == "[":
bracket_depth += 1
elif char == "]" and bracket_depth > 0:
bracket_depth -= 1
elif char == "," and bracket_depth == 0:
return normalized[:index].strip()
return normalized
async def _build_outputs(
self,
func_name: str,
type_hints: Dict[str, Any],
parser: DocstringParser,
drops: str,
) -> List[DataMapping]:
outputs: List[Dict[str, Any]] = []
return_type = type_hints.get("return", None)
doc_returns = parser.returns()
if return_type and return_type != type(None): # Check for actual return type
output_info: Dict[str, Any] = {
"name": "result",
"type_str": str(return_type),
"type": return_type,
}
if doc_returns:
output_info["description"] = doc_returns.get("description")
outputs.append(output_info)
returns_simstack_result = any(
output["type"] == SimstackResult for output in outputs
)
result_mappings = []
db = context.db
if returns_simstack_result:
if len(outputs) > 1:
logger.warning(
f"Node {func_name} returns more than one output, one of which is SimstackResult"
)
else:
doc_simstack_result = parser.simstack_results()
if doc_simstack_result is None:
logger.warning(f"The docstring of {func_name} does not defines its SimstackResult outputs")
else:
for name, data in doc_simstack_result.items():
output_mapping = None
output_type = self._normalize_docstring_type(data["type"])
try:
wrapper, inner_type_str = self._parse_generic_type(
output_type
)
if wrapper:
if inner_type_str is None:
raise ValueError(
f"Could not parse inner type for '{output_type}'"
)
# Handle List[type] or Dict[str,type]
inner_mapping = None
if "." in inner_type_str:
inner_mapping = inner_type_str
else:
try:
inner_model = await import_class_by_name(
inner_type_str, db
)
except (ValueError, LookupError):
inner_model = None
if inner_model is not None:
inner_mapping = self.get_class_mapping(
inner_model, drops
)
if inner_mapping is not None:
if wrapper == "List":
output_mapping = f"List[{inner_mapping}]"
else: # Dict
output_mapping = f"Dict[str,{inner_mapping}]"
elif "." in output_type:
# It's a full mapping path
output_mapping = output_type
else:
# It's a single class name
try:
output_model = await import_class_by_name(
output_type, db
)
output_mapping = self.get_class_mapping(
output_model, drops
)
except (ValueError, LookupError) as e:
logger.error(
f"Could not parse '{data['type']}' to mapping: {e}"
)
output_mapping = None
if output_mapping is not None:
result_mappings.append(
DataMapping(
name=name,
mapping=output_mapping,
description=data.get("description"),
)
)
except (ValueError, LookupError) as e:
logger.error(
f"Could not parse '{data['type']}' to mapping: {e}"
)
else: # not a SimstackResult
for output in outputs:
try:
output_type = output["type"]
if isinstance(output_type, str):
output_mapping = output_type
if " | None" in output_mapping:
output_mapping = output_mapping.replace(" | None", "")
else:
# Handle Optional[T] / T | None
from typing import get_args, get_origin, Union
import types
origin = get_origin(output_type)
if origin is types.UnionType or origin is Union:
args = get_args(output_type)
if type(None) in args:
# It's an Optional, take the first non-None argument
output_type = next(arg for arg in args if arg is not type(None))
output_mapping = self.get_class_mapping(output_type, drops)
result_mappings.append(
DataMapping(
name=output["name"],
mapping=output_mapping,
description=output.get("description"),
)
)
except ValueError:
logger.error(f"Could not parse '{output['type']}' to mapping in {func_name}")
return result_mappings
def _extract_default_parameters(self, func: Callable[..., Any]) -> Parameters:
"""
Best-effort extraction of node Parameters from either a direct attribute
or from closure variables. Always returns a non-None Parameters().
"""
parameters = Parameters()
if hasattr(func, "_node_parameters"):
return func._node_parameters
closures = inspect.getclosurevars(func)
for name, values in closures._asdict().items():
if name == "nonlocals":
continue
if isinstance(values, dict):
kwargs_node = values.get("kwargs_node", None)
if kwargs_node and "parameters" in kwargs_node:
return kwargs_node["parameters"]
return parameters
async def _delete_existing_node_model_if_needed(
self,
node_name: str,
function_mapping: str,
) -> tuple[bool, bool]:
"""
If a NodeModel with the same name exists, delete it (and its pickle if present),
but preserve the 'favorite' flag for the new entry.
Returns:
(should_skip, existing_favorite)
"""
try:
existing_model = await self.db.find_one(
NodeModel, NodeModel.name == node_name
)
except Exception as e:
logger.error(f"Error finding existing NodeModel {node_name}: {e}")
return False, False
if not existing_model:
return False, False
if function_mapping != existing_model.function_mapping:
logger.error(
f"NodeModel '{node_name}' already exists in the database\n"
f" DB Mapping: {existing_model.function_mapping}\n"
f" New Mapping: {function_mapping}\n"
f"New Mapping will overwrite DB Mapping."
)
existing_favorite = getattr(existing_model, "favorite", False)
if getattr(existing_model, "pickle_function", None):
try:
await self.db.delete(existing_model.pickle_function)
except Exception as e:
logger.error(f"Error deleting FunctionPickle for {node_name}: {e}")
try:
await self.db.delete(existing_model)
except Exception as e:
logger.error(f"Error deleting existing NodeModel {node_name}: {e}")
return False, existing_favorite
async def _resolve_input_mappings(
self,
node_name: str,
inputs: List[Dict[str, Any]],
drops: str,
) -> List[str]:
"""
Convert input python types to ModelMapping.mapping strings and validate their existence.
"""
input_mappings: List[str] = []
if not inputs:
return input_mappings
try:
for specific_input in inputs:
if (
specific_input.get("type")
and hasattr(specific_input["type"], "__module__")
and hasattr(specific_input["type"], "__name__")
):
input_mapping = (
specific_input["type"].__module__
+ "."
+ specific_input["type"].__name__
)
if drops and input_mapping.startswith(drops + "."):
input_mapping = input_mapping[len(drops) + 1 :]
input_mapping_found = await self.db.find_one(
ModelMapping, ModelMapping.mapping == input_mapping
)
if not input_mapping_found and input_mapping:
logger.error(
f"Processing node: {node_name} model {input_mapping} not found in db!"
)
input_mappings.append(input_mapping)
except Exception as e:
logger.error(f"Error getting input mapping: {e}")
return input_mappings
[docs]
def get_class_mapping(self, type: Type, drops: str = "") -> str:
"""Return the class mapping for a given type, optionally dropping a prefix."""
if hasattr(type, "__module__") and hasattr(type, "__name__"):
mapping = type.__module__ + "." + type.__name__
if drops and mapping.startswith(drops + "."):
mapping = mapping[len(drops) + 1 :]
return mapping
else:
raise ValueError(f"Could not parse '{type}' to mapping")
async def _register_nodes_from_module(self, module: Any, drops: str) -> None:
"""
Core logic to discover node functions in a module and (re)create NodeModel entries.
Heuristic:
- All top-level callables (functions) whose names do not start with '_'
- Only functions actually defined in this module.
"""
functions = self._discover_module_functions(module)
for func_name, func in functions:
if not is_node_function(func):
continue
sig = inspect.signature(func)
parser = DocstringParser(inspect.getdoc(func))
doc_description = parser.description()
doc_params = parser.params()
type_hints = get_type_hints(func)
inputs = self._build_inputs(sig, type_hints, doc_params)
parameters = self._extract_default_parameters(func)
node_name = getattr(func, "_node_name", func_name)
node_description = getattr(func, "_node_description", doc_description or "")
if not inputs:
logger.warning(
f"{node_name} has no inputs -- this means the node will be executed only once."
)
input_mappings = await self._resolve_input_mappings(
node_name, inputs, drops
)
function_mapping = module.__name__ + "." + func_name
try:
(
should_skip,
existing_favorite,
) = await self._delete_existing_node_model_if_needed(
node_name, function_mapping
)
if should_skip:
continue
data_mappings = []
for data_input, input_mapping in zip(inputs, input_mappings):
if data_input.get("type") and hasattr(
data_input["type"], "__name__"
):
data_input_mapping = self.get_class_mapping(
data_input["type"], drops
)
if data_input_mapping != input_mapping:
self.logger.error(
f"Type mismatch for input '{data_input['name']}': expected '{data_input['type'].__name__}', got '{input_mapping}'"
)
else:
self.logger.error(
f"No type specified for input '{data_input['name']}'"
)
data_mappings.append(
DataMapping(
name=data_input["name"],
mapping=input_mapping,
description=data_input.get("description"),
)
)
result_mappings = await self._build_outputs(
func_name, type_hints, parser, drops
)
if node_name is None:
raise ValueError(f"Node {func_name} has no name")
node_model = NodeModel(
name=node_name,
function_mapping=function_mapping,
description=node_description,
input_mappings=data_mappings,
result_mappings=result_mappings,
called_nodes=[], # we need to first build the full list, will be filled in second_stage
default_parameters=parameters,
pickle_function=None,
favorite=existing_favorite,
)
logger.debug(
f"NodeModel: {node_model.name}, {node_model.function_mapping}, {node_model.input_mappings}"
)
await self.db.save(node_model)
except Exception as e:
logger.error(f"Error creating/saving NodeModel {node_name}: {e}")
import traceback
traceback.print_exc()
[docs]
async def second_stage(self, drops: str) -> None:
await update_node_children(self.db, drops)
[docs]
async def clear_table(self) -> None:
self.logger.info("Clearing NodeModel collection")
await self.db.get_collection(NodeModel).drop()
[docs]
async def make_node_table(
db: Database,
dirs: list[str] | None = None,
drops: str | None = None,
write_schema: bool = False,
clear: bool = False,
project_root: Path | None = None,
ignore_entrypoints: bool = False,
) -> None:
"""
Rebuild the node table using the given databse.
This is a thin wrapper around CreateNodeTable for backward compatibility.
"""
creator = CreateNodeTable(db, write_schema=write_schema, project_root=project_root)
await creator.build(dirs=dirs, drops=drops, clear=clear, ignore_entrypoints=ignore_entrypoints)
[docs]
def create_node_table_main() -> None:
"""
CLI-style entry point to (re)build the node table.
Uses a dedicated event loop, matching the pattern used for model table creation.
"""
TableBuilderBase.cli_main(CreateNodeTable)
if __name__ == "__main__":
create_node_table_main()