Take NX parameter types from the docstrings.

lynxkite-core/src/lynxkite/core/ops.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 import enum
 import functools
 import inspect
+import types
 import pydantic
 import typing
 from dataclasses import dataclass
@@ -123,6 +124,26 @@ def basic_outputs(*names):
     return {name: Output(name=name, type=None) for name in names}
 
 
+def _param_to_type(name, value, type):
+    value = value or ""
+    print(f'Converting "{name}" {value} to {type}')
+    if type is int:
+        assert value != "", f"{name} is unset."
+        return int(value)
+    if type is float:
+        assert value != "", f"{name} is unset."
+        return float(value)
+    if isinstance(type, enum.EnumMeta):
+        return type[value]
+    if isinstance(type, types.UnionType):
+        match type.__args__:
+            case (None, type):
+                return None if value == "" else _param_to_type(name, value, type)
+            case (type, None):
+                return None if value == "" else _param_to_type(name, value, type)
+    return value
+
+
 class Op(BaseConfig):
     func: typing.Callable = pydantic.Field(exclude=True)
     name: str
@@ -136,12 +157,8 @@ class Op(BaseConfig):
         # Convert parameters.
         for p in params:
             if p in self.params:
-                if self.params[p].type is int:
-                    params[p] = int(params[p])
-                elif self.params[p].type is float:
-                    params[p] = float(params[p])
-                elif isinstance(self.params[p].type, enum.EnumMeta):
-                    params[p] = self.params[p].type[params[p]]
+                params[p] = _param_to_type(p, params[p], self.params[p].type)
+                print(self.name, p, params[p])
         res = self.func(*inputs, **params)
         if not isinstance(res, Result):
             # Automatically wrap the result in a Result object, if it isn't already.

lynxkite-graph-analytics/src/lynxkite_graph_analytics/networkx_ops.py

@@ -4,10 +4,149 @@ from lynxkite.core import ops
 import functools
 import inspect
 import networkx as nx
+import re
 
 ENV = "LynxKite Graph Analytics"
 
 
+class UnsupportedType(Exception):
+    pass
+
+
+nx.ladder_graph
+
+
+def doc_to_type(name: str, t: str) -> type:
+    t = t.lower()
+    t = re.sub("[(][^)]+[)]", "", t).strip().strip(".")
+    if " " in name or "http" in name:
+        return None  # Not a parameter type.
+    if t.endswith(", optional"):
+        w = doc_to_type(name, t.removesuffix(", optional").strip())
+        if w is None:
+            return None
+        return w | None
+    if t in [
+        "a digraph or multidigraph",
+        "a graph g",
+        "graph",
+        "graphs",
+        "networkx graph instance",
+        "networkx graph",
+        "networkx undirected graph",
+        "nx.graph",
+        "undirected graph",
+        "undirected networkx graph",
+    ] or t.startswith("networkx graph"):
+        return nx.Graph
+    elif t in [
+        "digraph-like",
+        "digraph",
+        "directed graph",
+        "networkx digraph",
+        "networkx directed graph",
+        "nx.digraph",
+    ]:
+        return nx.DiGraph
+    elif t == "node":
+        raise UnsupportedType(t)
+    elif t == '"node (optional)"':
+        return None
+    elif t == '"edge"':
+        raise UnsupportedType(t)
+    elif t == '"edge (optional)"':
+        return None
+    elif t in ["class", "data type"]:
+        raise UnsupportedType(t)
+    elif t in ["string", "str", "node label"]:
+        return str
+    elif t in ["string or none", "none or string", "string, or none"]:
+        return str | None
+    elif t in ["int", "integer"]:
+        return int
+    elif t in ["bool", "boolean"]:
+        return bool
+    elif t == "tuple":
+        raise UnsupportedType(t)
+    elif t == "set":
+        raise UnsupportedType(t)
+    elif t == "list of floats":
+        raise UnsupportedType(t)
+    elif t == "list of floats or float":
+        return float
+    elif t in ["dict", "dictionary"]:
+        raise UnsupportedType(t)
+    elif t == "scalar or dictionary":
+        return float
+    elif t == "none or dict":
+        return None
+    elif t in ["function", "callable"]:
+        raise UnsupportedType(t)
+    elif t in [
+        "collection",
+        "container of nodes",
+        "list of nodes",
+    ]:
+        raise UnsupportedType(t)
+    elif t in [
+        "container",
+        "generator",
+        "iterable",
+        "iterator",
+        "list or iterable container",
+        "list or iterable",
+        "list or set",
+        "list or tuple",
+        "list",
+    ]:
+        raise UnsupportedType(t)
+    elif t == "generator of sets":
+        raise UnsupportedType(t)
+    elif t == "dict or a set of 2 or 3 tuples":
+        raise UnsupportedType(t)
+    elif t == "set of 2 or 3 tuples":
+        raise UnsupportedType(t)
+    elif t == "none, string or function":
+        return str | None
+    elif t == "string or function" and name == "weight":
+        return str
+    elif t == "integer, float, or none":
+        return float | None
+    elif t in [
+        "float",
+        "int or float",
+        "integer or float",
+        "integer, float",
+        "number",
+        "numeric",
+        "real",
+        "scalar",
+    ]:
+        return float
+    elif t in ["integer or none", "int or none"]:
+        return int | None
+    elif name == "seed":
+        return int | None
+    elif name == "weight":
+        return str
+    elif t == "object":
+        raise UnsupportedType(t)
+    return None
+
+
+def types_from_doc(doc: str) -> dict[str, type]:
+    types = {}
+    for line in doc.splitlines():
+        if ":" in line:
+            a, b = line.split(":", 1)
+            for a in a.split(","):
+                a = a.strip()
+                t = doc_to_type(a, b)
+                if t is not None:
+                    types[a] = t
+    return types
+
+
 def wrapped(name: str, func):
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
@@ -27,36 +166,45 @@ def wrapped(name: str, func):
 
 def register_networkx(env: str):
     cat = ops.CATALOGS.setdefault(env, {})
+    counter = 0
     for name, func in nx.__dict__.items():
         if hasattr(func, "graphs"):
             sig = inspect.signature(func)
+            try:
+                types = types_from_doc(func.__doc__)
+            except UnsupportedType:
+                continue
+            for k, param in sig.parameters.items():
+                if k in types:
+                    continue
+                if param.annotation is not param.empty:
+                    types[k] = param.annotation
+                if k in ["i", "j", "n"]:
+                    types[k] = int
             inputs = {k: ops.Input(name=k, type=nx.Graph) for k in func.graphs}
             params = {
                 name: ops.Parameter.basic(
-                    name,
-                    str(param.default)
+                    name=name,
+                    default=str(param.default)
                     if type(param.default) in [str, int, float]
                     else None,
-                    param.annotation,
+                    type=types[name],
                 )
                 for name, param in sig.parameters.items()
-                if name not in ["G", "backend", "backend_kwargs", "create_using"]
+                if name in types and types[name] not in [nx.Graph, nx.DiGraph]
             }
-            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()
+            nicename = name.replace("_", " ").title()
             op = ops.Op(
                 func=wrapped(name, func),
-                name=name,
+                name=nicename,
                 params=params,
                 inputs=inputs,
                 outputs={"output": ops.Output(name="output", type=nx.Graph)},
                 type="basic",
             )
-            cat[name] = op
+            cat[nicename] = op
+            counter += 1
+    print(f"Registered {counter} NetworkX operations.")
 
 
 register_networkx(ENV)