server/networkx_ops.py

"""Automatically wraps all NetworkX functions as LynxKite operations."""
from . import ops
import functools
import inspect
import networkx as nx


def wrapped(func):
  @functools.wraps(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, 'attr', name)
    return graph
  return wrapper


for (name, func) in nx.__dict__.items():
  if hasattr(func, 'graphs'):
    sig = inspect.signature(func)
    inputs = {k: 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']}
    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(wrapped(func), name, params=params, inputs=inputs, outputs={'output': 'yes'}, type='basic')
    ops.ALL_OPS[name] = op