Fix ty warnings in lynxkite-graph-analytics.

lynxkite-core/src/lynxkite/core/workspace.py

@@ -230,7 +230,7 @@ class Workspace(BaseConfig):
             )
         elif "title" in kwargs:
             kwargs["data"] = WorkspaceNodeData(
-                title=kwargs["title"], op_id=kwargs["title"], params={}
+                title=kwargs["title"], op_id=kwargs["title"], params=kwargs.get("params", {})
             )
         kwargs.setdefault("type", "basic")
         kwargs.setdefault("id", f"{kwargs['data'].title} {random_string}")

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

@@ -4,7 +4,7 @@ import os
 import pandas as pd
 
 if os.environ.get("NX_CUGRAPH_AUTOCONFIG", "").strip().lower() == "true":
-    import cudf.pandas
+    import cudf.pandas  # ty: ignore[unresolved-import]
 
     cudf.pandas.install()
 

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

@@ -174,13 +174,13 @@ def disambiguate_edges(ws: workspace.Workspace):
     for edge in reversed(ws.edges):
         dst_node = nodes[edge.target]
         op = catalog.get(dst_node.data.op_id)
-        if op.get_input(edge.targetHandle).type == list[Bundle]:
+        if not op or op.get_input(edge.targetHandle).type == list[Bundle]:
             # Takes multiple bundles as an input. No need to disambiguate.
             continue
         if (edge.target, edge.targetHandle) in seen:
             i = ws.edges.index(edge)
             del ws.edges[i]
-            if hasattr(ws, "_crdt"):
+            if ws._crdt:
                 del ws._crdt["edges"][i]
         seen.add((edge.target, edge.targetHandle))
 

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

@@ -8,7 +8,8 @@ from collections import deque
 
 from . import core
 import grandcypher
-import matplotlib
+import matplotlib.cm
+import matplotlib.colors
 import networkx as nx
 import pandas as pd
 import polars as pl
@@ -205,6 +206,7 @@ def _map_color(value):
     else:
         cmap = matplotlib.cm.get_cmap("Paired")
         categories = pd.Index(value.unique())
+        assert isinstance(cmap, matplotlib.colors.ListedColormap)
         colors = cmap.colors[: len(categories)]
         return [
             "#{:02x}{:02x}{:02x}".format(int(r * 255), int(g * 255), int(b * 255))
@@ -312,7 +314,7 @@ def view_tables(bundle: core.Bundle, *, _tables_open: str = "", limit: int = 100
     view="graph_creation_view",
     outputs=["output"],
 )
-def organize(bundles: list[core.Bundle], *, relations: str = None) -> core.Bundle:
+def organize(bundles: list[core.Bundle], *, relations: str = ""):
     """Merge multiple inputs and construct graphs from the tables.
 
     To create a graph, import tables for edges and nodes, and combine them in this operation.
@@ -322,6 +324,9 @@ def organize(bundles: list[core.Bundle], *, relations: str = None) -> core.Bundl
         bundle.dfs.update(b.dfs)
         bundle.relations.extend(b.relations)
         bundle.other.update(b.other)
-    if not (relations is None or relations.strip() == ""):
-        bundle.relations = [core.RelationDefinition(**r) for r in json.loads(relations).values()]
+    if relations.strip():
+        bundle.relations = [
+            core.RelationDefinition(**r)  # ty: ignore[missing-argument]
+            for r in json.loads(relations).values()
+        ]
     return ops.Result(output=bundle, display=bundle.to_dict(limit=100))

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

@@ -216,7 +216,7 @@ def view_vectors(
     metric: UMAPMetric = UMAPMetric.euclidean,
 ):
     try:
-        from cuml.manifold.umap import UMAP
+        from cuml.manifold.umap import UMAP  # ty: ignore[unresolved-import]
     except ImportError:
         from umap import UMAP
     vec = np.stack(bundle.dfs[table_name][vector_column].to_numpy())

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

@@ -1,14 +1,14 @@
 """Automatically wraps all NetworkX functions as LynxKite operations."""
 
-import collections
-import types
 from lynxkite.core import ops
+import collections.abc
+import enum
 import functools
 import inspect
 import networkx as nx
-import re
-
 import pandas as pd
+import re
+import types
 
 ENV = "LynxKite Graph Analytics"
 
@@ -17,20 +17,22 @@ class UnsupportedParameterType(Exception):
     pass
 
 
-_UNSUPPORTED = object()
-_SKIP = object()
+class Failure(str, enum.Enum):
+    UNSUPPORTED = "unsupported"  # This parameter will be hidden.
+    SKIP = "skip"  # We have to skip the whole function.
 
 
-def doc_to_type(name: str, type_hint: str) -> type:
+def doc_to_type(name: str, type_hint: str) -> type | types.UnionType | Failure:
     type_hint = type_hint.lower()
     type_hint = re.sub("[(][^)]+[)]", "", type_hint).strip().strip(".")
     if " " in name or "http" in name:
-        return _UNSUPPORTED  # Not a parameter type.
+        return Failure.UNSUPPORTED  # Not a parameter type.
     if type_hint.endswith(", optional"):
         w = doc_to_type(name, type_hint.removesuffix(", optional").strip())
-        if w is _UNSUPPORTED:
-            return _SKIP
-        return w if w is _SKIP else w | None
+        if w is Failure.UNSUPPORTED or w is Failure.SKIP:
+            return Failure.SKIP
+        assert not isinstance(w, Failure)
+        return w | None
     if type_hint in [
         "a digraph or multidigraph",
         "a graph g",
@@ -54,15 +56,15 @@ def doc_to_type(name: str, type_hint: str) -> type:
     ]:
         return nx.DiGraph
     elif type_hint == "node":
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint == '"node (optional)"':
-        return _SKIP
+        return Failure.SKIP
     elif type_hint == '"edge"':
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint == '"edge (optional)"':
-        return _SKIP
+        return Failure.SKIP
     elif type_hint in ["class", "data type"]:
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint in ["string", "str", "node label"]:
         return str
     elif type_hint in ["string or none", "none or string", "string, or none"]:
@@ -72,27 +74,27 @@ def doc_to_type(name: str, type_hint: str) -> type:
     elif type_hint in ["bool", "boolean"]:
         return bool
     elif type_hint == "tuple":
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint == "set":
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint == "list of floats":
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint == "list of floats or float":
         return float
     elif type_hint in ["dict", "dictionary"]:
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint == "scalar or dictionary":
         return float
     elif type_hint == "none or dict":
-        return _SKIP
+        return Failure.SKIP
     elif type_hint in ["function", "callable"]:
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint in [
         "collection",
         "container of nodes",
         "list of nodes",
     ]:
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint in [
         "container",
         "generator",
@@ -104,13 +106,13 @@ def doc_to_type(name: str, type_hint: str) -> type:
         "list or tuple",
         "list",
     ]:
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint == "generator of sets":
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint == "dict or a set of 2 or 3 tuples":
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint == "set of 2 or 3 tuples":
-        return _UNSUPPORTED
+        return Failure.UNSUPPORTED
     elif type_hint == "none, string or function":
         return str | None
     elif type_hint == "string or function" and name == "weight":
@@ -135,8 +137,8 @@ def doc_to_type(name: str, type_hint: str) -> type:
     elif name == "weight":
         return str
     elif type_hint == "object":
-        return _UNSUPPORTED
-    return _SKIP
+        return Failure.UNSUPPORTED
+    return Failure.SKIP
 
 
 def types_from_doc(doc: str) -> dict[str, type]:
@@ -186,13 +188,13 @@ def wrapped(name: str, func):
     return wrapper
 
 
-def _get_params(func) -> dict | None:
+def _get_params(func) -> list[ops.Parameter | ops.ParameterGroup]:
     sig = inspect.signature(func)
     # Get types from docstring.
     types = types_from_doc(func.__doc__)
     # Always hide these.
     for k in ["backend", "backend_kwargs", "create_using"]:
-        types[k] = _SKIP
+        types[k] = Failure.SKIP
     # Add in types based on signature.
     for k, param in sig.parameters.items():
         if k in types:
@@ -203,10 +205,10 @@ def _get_params(func) -> dict | None:
             types[k] = int
     params = []
     for name, param in sig.parameters.items():
-        _type = types.get(name, _UNSUPPORTED)
-        if _type is _UNSUPPORTED:
+        _type = types.get(name, Failure.UNSUPPORTED)
+        if _type is Failure.UNSUPPORTED:
             raise UnsupportedParameterType(name)
-        if _type is _SKIP or _type in [nx.Graph, nx.DiGraph]:
+        if _type is Failure.SKIP or _type in [nx.Graph, nx.DiGraph]:
             continue
         p = ops.Parameter.basic(
             name=name,

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch/pytorch_core.py

@@ -3,7 +3,7 @@
 import copy
 import graphlib
 import io
-import numpy as np
+import typing
 import pydantic
 from lynxkite.core import ops, workspace
 import torch
@@ -36,8 +36,12 @@ def reg(name, inputs=[], outputs=None, params=[], **kwargs):
     return ops.register_passive_op(
         ENV,
         name,
-        inputs=[ops.Input(name=name, position="bottom", type="tensor") for name in inputs],
-        outputs=[ops.Output(name=name, position="top", type="tensor") for name in outputs],
+        inputs=[
+            ops.Input(name=name, position=ops.Position.BOTTOM, type="tensor") for name in inputs
+        ],
+        outputs=[
+            ops.Output(name=name, position=ops.Position.TOP, type="tensor") for name in outputs
+        ],
         params=params,
         **kwargs,
     )
@@ -75,16 +79,16 @@ class ModelMapping(pydantic.BaseModel):
     map: dict[str, ColumnSpec]
 
 
-def _torch_save(data) -> str:
+def _torch_save(data) -> bytes:
     """Saves PyTorch data (modules, tensors) as a string."""
     buffer = io.BytesIO()
     torch.save(data, buffer)
     return buffer.getvalue()
 
 
-def _torch_load(data: str) -> any:
+def _torch_load(data: bytes) -> typing.Any:
     """Loads PyTorch data (modules, tensors) from a string."""
-    buffer = io.BytesIO(data)  # noqa: F821
+    buffer = io.BytesIO(data)
     return torch.load(buffer)
 
 
@@ -98,7 +102,7 @@ class ModelConfig:
     input_output_names: list[str]
     loss: torch.nn.Module
     source_workspace_json: str
-    optimizer_parameters: dict[str, any]
+    optimizer_parameters: dict[str, typing.Any]
     optimizer: torch.optim.Optimizer | None = None
     source_workspace: str | None = None
     trained: bool = False
@@ -116,7 +120,7 @@ class ModelConfig:
         values = {k: v for k, v in zip(self.model_outputs, output)}
         return values
 
-    def inference(self, inputs: dict[str, torch.Tensor]) -> dict[str, np.ndarray]:
+    def inference(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         """Inference on a single batch."""
         self.model.eval()
         return self._forward(inputs)
@@ -422,7 +426,7 @@ class ModelBuilder:
         op = self.catalog["Optimizer"]
         cfg["optimizer_parameters"] = op.convert_params(self.nodes[self.optimizer].data.params)
         cfg["source_workspace_json"] = self.ws.model_dump_json()
-        return ModelConfig(**cfg)
+        return ModelConfig(**cfg)  # ty: ignore[missing-argument]
 
     def get_names(self, *ids: list[str]) -> dict[str, str]:
         """Returns a mapping from internal IDs to human-readable names."""
@@ -448,7 +452,7 @@ class ModelBuilder:
 
 
 def to_batch_tensors(
-    b: core.Bundle, batch_size: int, batch_index: int, m: ModelMapping | None
+    b: core.Bundle, batch_size: int, batch_index: int, m: ModelMapping
 ) -> dict[str, torch.Tensor]:
     """Extracts tensors from a bundle for a specific batch using a model mapping."""
     tensors = {}

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch/pytorch_ops.py

@@ -117,7 +117,7 @@ def neural_ode_mlp(
 
 @op("Attention", outputs=["outputs", "weights"])
 def attention(query, key, value, *, embed_dim=1024, num_heads=1, dropout=0.0):
-    return torch.nn.MultiHeadAttention(embed_dim, num_heads, dropout=dropout, need_weights=True)
+    return torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
 
 
 @op("LayerNorm", outputs=["outputs", "weights"])
@@ -250,8 +250,8 @@ reg(
 ops.register_passive_op(
     ENV,
     "Repeat",
-    inputs=[ops.Input(name="input", position="top", type="tensor")],
-    outputs=[ops.Output(name="output", position="bottom", type="tensor")],
+    inputs=[ops.Input(name="input", position=ops.Position.TOP, type="tensor")],
+    outputs=[ops.Output(name="output", position=ops.Position.BOTTOM, type="tensor")],
     params=[
         ops.Parameter.basic("times", 1, int),
         ops.Parameter.basic("same_weights", False, bool),
@@ -261,8 +261,8 @@ ops.register_passive_op(
 ops.register_passive_op(
     ENV,
     "Recurrent chain",
-    inputs=[ops.Input(name="input", position="top", type="tensor")],
-    outputs=[ops.Output(name="output", position="bottom", type="tensor")],
+    inputs=[ops.Input(name="input", position=ops.Position.TOP, type="tensor")],
+    outputs=[ops.Output(name="output", position=ops.Position.BOTTOM, type="tensor")],
     params=[],
 )
 

lynxkite-graph-analytics/tests/test_core.py

@@ -20,19 +20,19 @@ async def test_multi_input_box():
     ws.add_node(
         id="1",
         type="node_type",
-        data=workspace.WorkspaceNodeData(title="Create Bundle", params={}),
+        title="Create Bundle",
         position=workspace.Position(x=0, y=0),
     )
     ws.add_node(
         id="2",
         type="node_type",
-        data=workspace.WorkspaceNodeData(title="Create Bundle", params={}),
+        title="Create Bundle",
         position=workspace.Position(x=0, y=0),
     )
     ws.add_node(
         id="3",
         type="node_type",
-        data=workspace.WorkspaceNodeData(title="Multi input op", params={}),
+        title="Multi input op",
         position=workspace.Position(x=0, y=0),
     )
     ws.edges = [

lynxkite-graph-analytics/tests/test_lynxkite_ops.py

@@ -12,7 +12,7 @@ async def test_execute_operation_not_in_catalog():
     ws.add_node(
         id="1",
         type="node_type",
-        data=workspace.WorkspaceNodeData(title="Non existing op", params={}),
+        title="Non existing op",
         position=workspace.Position(x=0, y=0),
     )
     await execute(ws)
@@ -78,25 +78,25 @@ async def test_execute_operation_inputs_correct_cast():
     ws.add_node(
         id="1",
         type="node_type",
-        data=workspace.WorkspaceNodeData(title="Create Bundle", params={}),
+        title="Create Bundle",
         position=workspace.Position(x=0, y=0),
     )
     ws.add_node(
         id="2",
         type="node_type",
-        data=workspace.WorkspaceNodeData(title="Bundle to Graph", params={}),
+        title="Bundle to Graph",
         position=workspace.Position(x=100, y=0),
     )
     ws.add_node(
         id="3",
         type="node_type",
-        data=workspace.WorkspaceNodeData(title="Graph to Bundle", params={}),
+        title="Graph to Bundle",
         position=workspace.Position(x=200, y=0),
     )
     ws.add_node(
         id="4",
         type="node_type",
-        data=workspace.WorkspaceNodeData(title="Dataframe to Bundle", params={}),
+        title="Dataframe to Bundle",
         position=workspace.Position(x=300, y=0),
     )
     ws.edges = [
@@ -136,19 +136,19 @@ async def test_multiple_inputs():
     ws.add_node(
         id="one",
         type="cool",
-        data=workspace.WorkspaceNodeData(title="One", params={}),
+        title="One",
         position=workspace.Position(x=0, y=0),
     )
     ws.add_node(
         id="two",
         type="cool",
-        data=workspace.WorkspaceNodeData(title="Two", params={}),
+        title="Two",
         position=workspace.Position(x=100, y=0),
     )
     ws.add_node(
         id="smaller",
         type="cool",
-        data=workspace.WorkspaceNodeData(title="Smaller?", params={}),
+        title="Smaller?",
         position=workspace.Position(x=200, y=0),
     )
     ws.edges = [
@@ -188,8 +188,8 @@ async def test_optional_inputs():
         return a + (b or 0)
 
     assert maybe_add.__op__.inputs == [
-        ops.Input(name="a", type=int, position="left"),
-        ops.Input(name="b", type=int | None, position="left"),
+        ops.Input(name="a", type=int, position=ops.Position.LEFT),
+        ops.Input(name="b", type=int | None, position=ops.Position.LEFT),
     ]
     ws = workspace.Workspace(env="test", nodes=[], edges=[])
     a = ws.add_node(one)

lynxkite-graph-analytics/tests/test_pytorch_model_ops.py

@@ -1,3 +1,4 @@
+import torch_geometric.nn as pyg_nn
 from lynxkite.core import workspace
 from lynxkite_graph_analytics.pytorch import pytorch_core
 import torch
@@ -11,7 +12,8 @@ def make_ws(env, nodes: dict[str, dict], edges: list[tuple[str, str]]):
         del data["title"]
         ws.add_node(
             id=id,
-            data=workspace.WorkspaceNodeData(title=title, params=data),
+            title=title,
+            params=data,
         )
     ws.edges = [
         workspace.WorkspaceEdge(
@@ -27,10 +29,12 @@ def make_ws(env, nodes: dict[str, dict], edges: list[tuple[str, str]]):
 
 
 def summarize_layers(m: pytorch_core.ModelConfig) -> str:
+    assert isinstance(m.model, pyg_nn.Sequential)
     return "".join(str(e)[:2] for e in m.model)
 
 
 def summarize_connections(m: pytorch_core.ModelConfig) -> str:
+    assert isinstance(m.model, pyg_nn.Sequential)
     return " ".join(
         "".join(n[0] for n in c.param_names) + "->" + "".join(n[0] for n in c.return_names)
         for c in m.model._children