239 lines
6.5 KiB
Python
239 lines
6.5 KiB
Python
from collections.abc import Callable
|
|
|
|
ROOT_COLOR = "#ff9999"
|
|
GHOST_COLOR = "#cccccc"
|
|
|
|
def from_nested_dict(trees: list[dict]) -> tuple[list, list]:
|
|
"""
|
|
Convert a list of nested dictionaries to vis.js nodes and edges format.
|
|
|
|
Args:
|
|
trees: List of nested dictionaries where keys are node names and values are
|
|
dictionaries of children (e.g., [{"root1": {"child1": {}}}, {"root2": {}}])
|
|
|
|
Returns:
|
|
tuple: (nodes, edges) where:
|
|
- nodes: list of dicts with auto-incremented numeric IDs
|
|
- edges: list of dicts with 'from' and 'to' keys
|
|
|
|
Example:
|
|
>>> trees = [{"root1": {"child1": {}}}, {"root2": {"child2": {}}}]
|
|
>>> nodes, edges = from_nested_dict(trees)
|
|
>>> # nodes = [{"id": 1, "label": "root1"}, {"id": 2, "label": "child1"}, ...]
|
|
"""
|
|
nodes = []
|
|
edges = []
|
|
node_id_counter = 1
|
|
node_id_map = {} # maps node_label -> node_id
|
|
|
|
def traverse(subtree: dict, parent_id: int | None = None):
|
|
nonlocal node_id_counter
|
|
|
|
for node_label, children in subtree.items():
|
|
# Create node with auto-incremented ID
|
|
current_id = node_id_counter
|
|
node_id_map[node_label] = current_id
|
|
nodes.append({
|
|
"id": current_id,
|
|
"label": node_label
|
|
})
|
|
node_id_counter += 1
|
|
|
|
# Create edge from parent to current node
|
|
if parent_id is not None:
|
|
edges.append({
|
|
"from": parent_id,
|
|
"to": current_id
|
|
})
|
|
|
|
# Recursively process children
|
|
if children:
|
|
traverse(children, parent_id=current_id)
|
|
|
|
# Process each tree in the list
|
|
for tree in trees:
|
|
traverse(tree)
|
|
|
|
return nodes, edges
|
|
|
|
|
|
def from_tree_with_metadata(
|
|
trees: list[dict],
|
|
id_getter: Callable = None,
|
|
label_getter: Callable = None,
|
|
children_getter: Callable = None
|
|
) -> tuple[list, list]:
|
|
"""
|
|
Convert a list of trees with metadata to vis.js nodes and edges format.
|
|
|
|
Args:
|
|
trees: List of dictionaries with 'id', 'label', and 'children' keys
|
|
(e.g., [{"id": "root1", "label": "Root 1", "children": [...]}, ...])
|
|
id_getter: Optional callback to extract node ID from dict
|
|
Default: lambda n: n.get("id")
|
|
label_getter: Optional callback to extract node label from dict
|
|
Default: lambda n: n.get("label", "")
|
|
children_getter: Optional callback to extract children list from dict
|
|
Default: lambda n: n.get("children", [])
|
|
|
|
Returns:
|
|
tuple: (nodes, edges) where:
|
|
- nodes: list of dicts with IDs from tree or auto-incremented
|
|
- edges: list of dicts with 'from' and 'to' keys
|
|
|
|
Example:
|
|
>>> trees = [
|
|
... {
|
|
... "id": "root1",
|
|
... "label": "Root Node 1",
|
|
... "children": [
|
|
... {"id": "child1", "label": "Child 1", "children": []}
|
|
... ]
|
|
... },
|
|
... {"id": "root2", "label": "Root Node 2"}
|
|
... ]
|
|
>>> nodes, edges = from_tree_with_metadata(trees)
|
|
"""
|
|
# Default getters
|
|
if id_getter is None:
|
|
id_getter = lambda n: n.get("id")
|
|
if label_getter is None:
|
|
label_getter = lambda n: n.get("label", "")
|
|
if children_getter is None:
|
|
children_getter = lambda n: n.get("children", [])
|
|
|
|
nodes = []
|
|
edges = []
|
|
node_id_counter = 1
|
|
|
|
def traverse(node_dict: dict, parent_id: int | str | None = None):
|
|
nonlocal node_id_counter
|
|
|
|
# Extract ID (use provided or auto-increment)
|
|
node_id = id_getter(node_dict)
|
|
if node_id is None:
|
|
node_id = node_id_counter
|
|
node_id_counter += 1
|
|
|
|
# Extract label
|
|
node_label = label_getter(node_dict)
|
|
|
|
# Create node
|
|
nodes.append({
|
|
"id": node_id,
|
|
"label": node_label
|
|
})
|
|
|
|
# Create edge from parent to current node
|
|
if parent_id is not None:
|
|
edges.append({
|
|
"from": parent_id,
|
|
"to": node_id
|
|
})
|
|
|
|
# Recursively process children
|
|
children = children_getter(node_dict)
|
|
if children:
|
|
for child in children:
|
|
traverse(child, parent_id=node_id)
|
|
|
|
# Process each tree in the list
|
|
for tree in trees:
|
|
traverse(tree)
|
|
|
|
return nodes, edges
|
|
|
|
|
|
def from_parent_child_list(
|
|
items: list,
|
|
id_getter: Callable = None,
|
|
label_getter: Callable = None,
|
|
parent_getter: Callable = None,
|
|
ghost_color: str = GHOST_COLOR,
|
|
root_color: str | None = ROOT_COLOR
|
|
) -> tuple[list, list]:
|
|
"""
|
|
Convert a list of items with parent references to vis.js nodes and edges format.
|
|
|
|
Args:
|
|
items: List of items (dicts or objects) with parent references
|
|
id_getter: callback to extract node ID
|
|
label_getter: callback to extract node label
|
|
parent_getter: callback to extract parent ID
|
|
ghost_color: color for ghost nodes (referenced parents)
|
|
root_color: color for root nodes (nodes without parent)
|
|
|
|
Returns:
|
|
tuple: (nodes, edges)
|
|
"""
|
|
|
|
# Default getters
|
|
if id_getter is None:
|
|
id_getter = lambda item: item.get("id")
|
|
|
|
if label_getter is None:
|
|
label_getter = lambda item: item.get("label", "")
|
|
|
|
if parent_getter is None:
|
|
parent_getter = lambda item: item.get("parent")
|
|
|
|
nodes = []
|
|
edges = []
|
|
|
|
# Track all existing node IDs
|
|
existing_ids = set()
|
|
|
|
# First pass: create nodes for all items
|
|
for item in items:
|
|
node_id = id_getter(item)
|
|
node_label = label_getter(item)
|
|
|
|
existing_ids.add(node_id)
|
|
nodes.append({
|
|
"id": node_id,
|
|
"label": node_label,
|
|
# root color assigned later
|
|
})
|
|
|
|
# Track ghost nodes
|
|
ghost_nodes = set()
|
|
|
|
# Track which nodes have parents
|
|
nodes_with_parent = set()
|
|
|
|
# Second pass: create edges and detect ghost nodes
|
|
for item in items:
|
|
node_id = id_getter(item)
|
|
parent_id = parent_getter(item)
|
|
|
|
# Skip roots
|
|
if parent_id is None or parent_id == "":
|
|
continue
|
|
|
|
# Child has a parent
|
|
nodes_with_parent.add(node_id)
|
|
|
|
# Create edge parent → child
|
|
edges.append({
|
|
"from": parent_id,
|
|
"to": node_id
|
|
})
|
|
|
|
# Create ghost node if parent not found
|
|
if parent_id not in existing_ids and parent_id not in ghost_nodes:
|
|
ghost_nodes.add(parent_id)
|
|
nodes.append({
|
|
"id": parent_id,
|
|
"label": str(parent_id),
|
|
"color": ghost_color
|
|
})
|
|
|
|
# Final pass: assign color to root nodes
|
|
if root_color is not None:
|
|
for node in nodes:
|
|
if node["id"] not in nodes_with_parent and node["id"] not in ghost_nodes:
|
|
# Root node
|
|
node["color"] = root_color
|
|
|
|
return nodes, edges
|