Spaces:
Running
Running
"""Automatically wraps all NetworkX functions as LynxKite operations.""" | |
from . import ops | |
import functools | |
import inspect | |
import networkx as nx | |
def wrapped(name: str, func): | |
def wrapper(*args, **kwargs): | |
for k, v in kwargs.items(): | |
if v == 'None': | |
kwargs[k] = None | |
res = func(*args, **kwargs) | |
if isinstance(res, nx.Graph): | |
return res | |
# Otherwise it's a node attribute. | |
graph = args[0].copy() | |
nx.set_node_attributes(graph, values=res, name=name) | |
return graph | |
return wrapper | |
def register_networkx(env: str): | |
ops.CATALOGS.setdefault(env, {}) | |
for (name, func) in nx.__dict__.items(): | |
if hasattr(func, 'graphs'): | |
sig = inspect.signature(func) | |
inputs = {k: ops.Input(name=k, type=nx.Graph) for k in func.graphs} | |
params = { | |
name: ops.Parameter.basic( | |
name, str(param.default) | |
if type(param.default) in [str, int, float] | |
else None, | |
param.annotation) | |
for name, param in sig.parameters.items() | |
if name not in ['G', 'backend', 'backend_kwargs', 'create_using']} | |
for p in params.values(): | |
if not p.type: | |
# Guess the type based on the name. | |
if len(p.name) == 1: | |
p.type = int | |
name = "NX › " + name.replace('_', ' ').title() | |
op = ops.Op( | |
func=wrapped(name, func), | |
name=name, | |
params=params, | |
inputs=inputs, | |
outputs={'output': ops.Output(name='output', type=nx.Graph)}, | |
type='basic', | |
) | |
ops.CATALOGS[env][name] = op | |
register_networkx('LynxKite') | |