medax_pipeline/graphCreation/node_typing.py
2025-04-16 22:12:19 +02:00

40 lines
1.4 KiB
Python

import networkx as nx
class Resource:
def __init__(self, resource_type):
self.resource_type = resource_type
def create_resource_class(resource_type):
return type(resource_type, (Resource,), {})
def set_resource_type(graph):
for node, data in graph.nodes(data=True):
print(node, data)
print("-----------------------------")
nodes_to_replace = []
for node, data in graph.nodes(data=True):
print(isinstance(node, Resource), node, type(node))
if isinstance(node, Resource):
print("Found a resource: ", node)
resource_type = node.resource_type
if resource_type:
# Dynamically create a new class based on the resource_type
NewResourceClass = create_resource_class(resource_type)
new_node = NewResourceClass(resource_type)
nodes_to_replace.append((node, new_node, data))
else:
print(f"Warning: Node {node} is a resource but has no resource_type")
# Replace old nodes with new ones
for old_node, new_node, data in nodes_to_replace:
graph.add_node(new_node, **data)
for pred in graph.predecessors(old_node):
graph.add_edge(pred, new_node)
for succ in graph.successors(old_node):
graph.add_edge(new_node, succ)
graph.remove_node(old_node)
"""
for node, data in graph.nodes(data=True):
print(node, data) """