Merge pull request #41 from lynxkite/darabos-ty

.github/workflows/test.yaml

@@ -6,19 +6,20 @@ on:
 jobs:
   test:
     runs-on: ubuntu-latest
-    env:
-      UV_SYSTEM_PYTHON: 1
     steps:
       - uses: actions/checkout@v4
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v5
-        with:
-            enable-cache: true
+        uses: astral-sh/setup-uv@v6
 
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "3.12"
+      - name: Set up Python
+        run: uv python install
+
+      # - name: Debug ty issue
+      #   run: |
+      #     uv pip install ty
+      #     uv pip list
+      #     uv run python -m ty check
 
       - name: Install dependencies
         run: |
@@ -33,27 +34,16 @@ jobs:
       - name: Run pre-commits
         run: |
           uv pip install pre-commit
-          pre-commit run --all-files
-
-      - name: Run core tests
-        run: |
-          cd lynxkite-core
-          pytest
-
-      - name: Run app tests
-        run: |
-          cd lynxkite-app
-          pytest
+          uv run pre-commit run --all-files
 
-      - name: Run graph analytics tests
+      - name: Run Python unittests
         run: |
-          cd lynxkite-graph-analytics
-          pytest
+          uv run pytest --asyncio-mode=auto
 
-      - name: Try building the documentation
+      - name: Build the documentation
         run: |
           uv pip install mkdocs-material mkdocstrings[python]
-          mkdocs build
+          uv run mkdocs build
 
       - uses: actions/setup-node@v4
         with:
@@ -67,9 +57,7 @@ jobs:
 
       - name: Run Playwright tests
         run: |
-          cd lynxkite-app/web
-          npm run build
-          npm run test
+          uv run bash -c 'cd lynxkite-app/web; npm run build; npm run test'
 
       - uses: actions/upload-artifact@v4
         name: Upload playwright report

.pre-commit-config.yaml

@@ -15,3 +15,13 @@ repos:
   rev: v1.9.4
   hooks:
     - id: biome-check
+# https://github.com/astral-sh/ty/issues/269
+- repo: local
+  hooks:
+    - id: ty-check
+      name: ty-check
+      language: python
+      entry: uv run ty check
+      pass_filenames: false
+      args: [--python=.venv/]
+      additional_dependencies: [ty]

examples/fake_data.py

@@ -1,5 +1,5 @@
 from lynxkite.core.ops import op
-from faker import Faker
+from faker import Faker  # ty: ignore[unresolved-import]
 import pandas as pd
 
 faker = Faker()

examples/word2vec.py

@@ -6,7 +6,7 @@ ENV = "LynxKite Graph Analytics"
 
 @op(ENV, "Word2vec for the top 1000 words", slow=True)
 def word2vec_1000():
-    import staticvectors
+    import staticvectors  # ty: ignore[unresolved-import]
 
     model = staticvectors.StaticVectors("neuml/word2vec-quantized")
     df = pd.read_csv(

lynxkite-app/src/lynxkite_app/crdt.py

@@ -7,8 +7,8 @@ import pathlib
 import fastapi
 import os.path
 import pycrdt.websocket
-import pycrdt.store
-import uvicorn
+import pycrdt.store.file
+import uvicorn.protocols.utils
 import builtins
 from lynxkite.core import workspace, ops
 
@@ -59,7 +59,8 @@ class WorkspaceWebsocketServer(pycrdt.websocket.WebsocketServer):
         room = pycrdt.websocket.YRoom(
             ystore=ystore, ydoc=ydoc, exception_handler=ws_exception_handler
         )
-        room.ws = ws
+        # We hang the YDoc pointer on the room, so it only gets garbage collected when the room does.
+        room.ws = ws  # ty: ignore[unresolved-attribute]
 
         def on_change(changes):
             task = asyncio.create_task(workspace_changed(name, changes, ws))
@@ -106,7 +107,8 @@ class CodeWebsocketServer(WorkspaceWebsocketServer):
         room = pycrdt.websocket.YRoom(
             ystore=ystore, ydoc=ydoc, exception_handler=ws_exception_handler
         )
-        room.text = text
+        # We hang the YDoc pointer on the room, so it only gets garbage collected when the room does.
+        room.text = text  # ty: ignore[unresolved-attribute]
 
         def on_change(changes):
             asyncio.create_task(code_changed(name, changes, text))
@@ -163,6 +165,7 @@ def crdt_update(
         ValueError: If the Python object provided is not a dict or list.
     """
     if isinstance(python_obj, dict):
+        assert isinstance(crdt_obj, pycrdt.Map), "CRDT object must be a Map for a dict input"
         for key, value in python_obj.items():
             if key in non_collaborative_fields:
                 crdt_obj[key] = value
@@ -179,6 +182,7 @@ def crdt_update(
             else:
                 crdt_obj[key] = value
     elif isinstance(python_obj, list):
+        assert isinstance(crdt_obj, pycrdt.Array), "CRDT object must be an Array for a list input"
         for i, value in enumerate(python_obj):
             if isinstance(value, dict):
                 if i >= len(crdt_obj):
@@ -220,7 +224,7 @@ last_known_versions = {}
 delayed_executions = {}
 
 
-async def workspace_changed(name: str, changes: pycrdt.MapEvent, ws_crdt: pycrdt.Map):
+async def workspace_changed(name: str, changes: list[pycrdt.MapEvent], ws_crdt: pycrdt.Map):
     """Callback to react to changes in the workspace.
 
     Args:
@@ -295,6 +299,14 @@ async def code_changed(name: str, changes: pycrdt.TextEvent, text: pycrdt.Text):
         f.write(contents)
 
 
+ws_websocket_server: WorkspaceWebsocketServer
+code_websocket_server: CodeWebsocketServer
+
+
+def get_room(name):
+    return ws_websocket_server.get_room(name)
+
+
 @contextlib.asynccontextmanager
 async def lifespan(app):
     global ws_websocket_server

lynxkite-app/src/lynxkite_app/main.py

@@ -8,7 +8,7 @@ import pathlib
 import pkgutil
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.gzip import GZipMiddleware
-import starlette
+import starlette.exceptions
 from lynxkite.core import ops
 from lynxkite.core import workspace
 from . import crdt
@@ -136,7 +136,7 @@ async def upload(req: fastapi.Request):
 @app.post("/api/execute_workspace")
 async def execute_workspace(name: str):
     """Trigger and await the execution of a workspace."""
-    room = await crdt.ws_websocket_server.get_room(name)
+    room = await crdt.get_room(name)
     ws_pyd = workspace.Workspace.model_validate(room.ws.to_py())
     await crdt.execute(name, room.ws, ws_pyd)
 

lynxkite-core/src/lynxkite/core/executors/one_by_one.py

@@ -48,10 +48,10 @@ def register(env: str, cache: bool = True):
     """
     if cache:
         CACHES[env] = {}
-        cache = CACHES[env]
+        _cache = CACHES[env]
     else:
-        cache = None
-    ops.EXECUTORS[env] = lambda ws: _execute(ws, ops.CATALOGS[env], cache=cache)
+        _cache = None
+    ops.EXECUTORS[env] = lambda ws: _execute(ws, ops.CATALOGS[env], cache=_cache)
 
 
 def _get_stages(ws, catalog: ops.Catalog):
@@ -102,7 +102,9 @@ async def _await_if_needed(obj):
     return obj
 
 
-async def _execute(ws: workspace.Workspace, catalog: ops.Catalog, cache=None):
+async def _execute(
+    ws: workspace.Workspace, catalog: ops.Catalog, cache: typing.Optional[dict] = None
+):
     nodes = {n.id: n for n in ws.nodes}
     contexts = {n.id: Context(node=n) for n in ws.nodes}
     edges = {n.id: [] for n in ws.nodes}

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

@@ -6,7 +6,7 @@ import asyncio
 import enum
 import functools
 import json
-import importlib
+import importlib.util
 import inspect
 import pathlib
 import subprocess
@@ -498,6 +498,7 @@ def install_requirements(req: pathlib.Path):
 
 def run_user_script(script_path: pathlib.Path):
     spec = importlib.util.spec_from_file_location(script_path.stem, str(script_path))
+    assert spec
     module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(module)
 

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

@@ -1,7 +1,7 @@
 """For working with LynxKite workspaces."""
 
 import json
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 import dataclasses
 import enum
 import os
@@ -10,6 +10,9 @@ import pydantic
 import tempfile
 from . import ops
 
+if TYPE_CHECKING:
+    from lynxkite.core import ops
+
 
 class BaseConfig(pydantic.BaseModel):
     model_config = pydantic.ConfigDict(
@@ -36,8 +39,7 @@ class WorkspaceNodeData(BaseConfig):
     input_metadata: Optional[object] = None
     error: Optional[str] = None
     status: NodeStatus = NodeStatus.done
-    # Also contains a "meta" field when going out.
-    # This is ignored when coming back from the frontend.
+    meta: Optional["ops.Op"] = None
 
     @pydantic.model_validator(mode="before")
     @classmethod
@@ -47,6 +49,13 @@ class WorkspaceNodeData(BaseConfig):
             data["op_id"] = data["title"]
         return data
 
+    @pydantic.model_validator(mode="before")
+    @classmethod
+    def ignore_meta(cls, data: dict) -> dict:
+        """Metadata is never loaded. We will use fresh metadata."""
+        data["meta"] = None
+        return data
+
 
 class WorkspaceNode(BaseConfig):
     # Most of these fields are shared with ReactFlow.
@@ -56,13 +65,13 @@ class WorkspaceNode(BaseConfig):
     position: Position
     width: Optional[float] = None
     height: Optional[float] = None
-    _crdt: pycrdt.Map
+    _crdt: Optional[pycrdt.Map] = None
 
     def publish_started(self):
         """Notifies the frontend that work has started on this node."""
         self.data.error = None
         self.data.status = NodeStatus.active
-        if hasattr(self, "_crdt") and "data" in self._crdt:
+        if self._crdt and "data" in self._crdt:
             with self._crdt.doc.transaction():
                 self._crdt["data"]["error"] = None
                 self._crdt["data"]["status"] = NodeStatus.active
@@ -73,7 +82,7 @@ class WorkspaceNode(BaseConfig):
         self.data.input_metadata = result.input_metadata
         self.data.error = result.error
         self.data.status = NodeStatus.done
-        if hasattr(self, "_crdt") and "data" in self._crdt:
+        if self._crdt and "data" in self._crdt:
             with self._crdt.doc.transaction():
                 try:
                     self._crdt["data"]["status"] = NodeStatus.done
@@ -109,7 +118,7 @@ class Workspace(BaseConfig):
     env: str = ""
     nodes: list[WorkspaceNode] = dataclasses.field(default_factory=list)
     edges: list[WorkspaceEdge] = dataclasses.field(default_factory=list)
-    _crdt: pycrdt.Map
+    _crdt: Optional[pycrdt.Map] = None
 
     def normalize(self):
         if self.env not in ops.CATALOGS:
@@ -194,19 +203,19 @@ class Workspace(BaseConfig):
                 if getattr(data, "meta", None) != op:
                     data.meta = op
                     # If the node is connected to a CRDT, update that too.
-                    if hasattr(node, "_crdt"):
+                    if node._crdt:
                         node._crdt["data"]["meta"] = op.model_dump()
                 if node.type != op.type:
                     node.type = op.type
-                    if hasattr(node, "_crdt"):
+                    if node._crdt:
                         node._crdt["type"] = op.type
                 if data.error == "Unknown operation.":
                     data.error = None
-                    if hasattr(node, "_crdt"):
+                    if node._crdt:
                         node._crdt["data"]["error"] = None
             else:
                 data.error = "Unknown operation."
-                if hasattr(node, "_crdt"):
+                if node._crdt:
                     node._crdt["data"]["meta"] = {}
                     node._crdt["data"]["error"] = "Unknown operation."
 
@@ -226,12 +235,16 @@ class Workspace(BaseConfig):
             kwargs["data"] = WorkspaceNodeData(
                 title=func.__op__.name, op_id=func.__op__.id, params={}
             )
+        elif "title" in kwargs:
+            kwargs["data"] = WorkspaceNodeData(
+                title=kwargs["title"], op_id=kwargs["title"], params=kwargs.get("params", {})
+            )
         kwargs.setdefault("type", "basic")
         kwargs.setdefault("id", f"{kwargs['data'].title} {random_string}")
         kwargs.setdefault("position", Position(x=0, y=0))
         kwargs.setdefault("width", 100)
         kwargs.setdefault("height", 100)
-        node = WorkspaceNode(**kwargs)
+        node = WorkspaceNode(**kwargs)  # ty: ignore[missing-argument]
         self.nodes.append(node)
         return node
 

lynxkite-core/tests/test_one_by_one.py

@@ -13,8 +13,8 @@ async def test_optional_inputs():
         return [a + b for a, b in zip(a, b)] if b else a
 
     assert maybe_add.__op__.inputs == [
-        ops.Input(name="a", type=list[int], position="bottom"),
-        ops.Input(name="b", type=list[int] | None, position="bottom"),
+        ops.Input(name="a", type=list[int], position=ops.Position.BOTTOM),
+        ops.Input(name="b", type=list[int] | None, position=ops.Position.BOTTOM),
     ]
     one_by_one.register("test")
     ws = workspace.Workspace(env="test", nodes=[], edges=[])

lynxkite-core/tests/test_ops.py

@@ -11,10 +11,10 @@ def test_op_decorator_no_params_no_types_default_positions():
     assert add.__op__.name == "add"
     assert add.__op__.params == []
     assert add.__op__.inputs == [
-        ops.Input(name="a", type=inspect._empty, position="left"),
-        ops.Input(name="b", type=inspect._empty, position="left"),
+        ops.Input(name="a", type=inspect._empty, position=ops.Position.LEFT),
+        ops.Input(name="b", type=inspect._empty, position=ops.Position.LEFT),
     ]
-    assert add.__op__.outputs == [ops.Output(name="result", type=None, position="right")]
+    assert add.__op__.outputs == [ops.Output(name="result", type=None, position=ops.Position.RIGHT)]
     assert add.__op__.type == "basic"
     assert ops.CATALOGS["test"]["add"] == add.__op__
 
@@ -29,10 +29,12 @@ def test_op_decorator_custom_positions():
     assert add.__op__.name == "add"
     assert add.__op__.params == []
     assert add.__op__.inputs == [
-        ops.Input(name="a", type=inspect._empty, position="right"),
-        ops.Input(name="b", type=inspect._empty, position="top"),
+        ops.Input(name="a", type=inspect._empty, position=ops.Position.RIGHT),
+        ops.Input(name="b", type=inspect._empty, position=ops.Position.TOP),
+    ]
+    assert add.__op__.outputs == [
+        ops.Output(name="result", type=None, position=ops.Position.BOTTOM)
     ]
-    assert add.__op__.outputs == [ops.Output(name="result", type=None, position="bottom")]
     assert add.__op__.type == "basic"
     assert ops.CATALOGS["test"]["add"] == add.__op__
 
@@ -45,10 +47,12 @@ def test_op_decorator_with_params_and_types_():
     assert multiply.__op__.name == "multiply"
     assert multiply.__op__.params == [ops.Parameter(name="param", default="param", type=str)]
     assert multiply.__op__.inputs == [
-        ops.Input(name="a", type=int, position="left"),
-        ops.Input(name="b", type=float, position="left"),
+        ops.Input(name="a", type=int, position=ops.Position.LEFT),
+        ops.Input(name="b", type=float, position=ops.Position.LEFT),
+    ]
+    assert multiply.__op__.outputs == [
+        ops.Output(name="result", type=None, position=ops.Position.RIGHT)
     ]
-    assert multiply.__op__.outputs == [ops.Output(name="result", type=None, position="right")]
     assert multiply.__op__.type == "basic"
     assert ops.CATALOGS["test"]["multiply"] == multiply.__op__
 
@@ -66,12 +70,14 @@ def test_op_decorator_with_complex_types():
     assert complex_op.__op__.name == "color_op"
     assert complex_op.__op__.params == []
     assert complex_op.__op__.inputs == [
-        ops.Input(name="color", type=Color, position="left"),
-        ops.Input(name="color_list", type=list[Color], position="left"),
-        ops.Input(name="color_dict", type=dict[str, Color], position="left"),
+        ops.Input(name="color", type=Color, position=ops.Position.LEFT),
+        ops.Input(name="color_list", type=list[Color], position=ops.Position.LEFT),
+        ops.Input(name="color_dict", type=dict[str, Color], position=ops.Position.LEFT),
     ]
     assert complex_op.__op__.type == "basic"
-    assert complex_op.__op__.outputs == [ops.Output(name="result", type=None, position="right")]
+    assert complex_op.__op__.outputs == [
+        ops.Output(name="result", type=None, position=ops.Position.RIGHT)
+    ]
     assert ops.CATALOGS["test"]["color_op"] == complex_op.__op__
 
 

lynxkite-core/tests/test_simple.py

@@ -12,8 +12,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),
     ]
     simple.register("test")
     ws = workspace.Workspace(env="test", nodes=[], edges=[])

lynxkite-core/tests/test_workspace.py

@@ -9,12 +9,12 @@ def test_save_load():
     ws.add_node(
         id="1",
         type="node_type",
-        data=workspace.WorkspaceNodeData(title="Node 1", params={}),
+        title="Node 1",
     )
     ws.add_node(
         id="2",
         type="node_type",
-        data=workspace.WorkspaceNodeData(title="Node 2", params={}),
+        title="Node 2",
     )
     ws.edges.append(
         workspace.WorkspaceEdge(
@@ -70,17 +70,18 @@ def test_update_metadata():
     ws.add_node(
         id="1",
         type="basic",
-        data=workspace.WorkspaceNodeData(title="Test Operation", params={}),
+        title="Test Operation",
     )
     ws.add_node(
         id="2",
         type="basic",
-        data=workspace.WorkspaceNodeData(title="Unknown Operation", params={}),
+        title="Unknown Operation",
     )
     ws.update_metadata()
+    assert ws.nodes[0].data.meta
     assert ws.nodes[0].data.meta.name == "Test Operation"
     assert ws.nodes[0].data.error is None
-    assert not hasattr(ws.nodes[1].data, "meta")
+    assert ws.nodes[1].data.meta is None
     assert ws.nodes[1].data.error == "Unknown operation."
 
 

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

@@ -227,13 +227,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

lynxkite-pillow-example/src/lynxkite_pillow_example/__init__.py

@@ -20,48 +20,48 @@ def open_image(*, filename: str):
 
 
 @op("Save image")
-def save_image(image: Image, *, filename: str):
+def save_image(image: Image.Image, *, filename: str):
     with fsspec.open(filename, "wb") as f:
         image.save(f)
 
 
 @op("Crop")
-def crop(image: Image, *, top: int, left: int, bottom: int, right: int):
+def crop(image: Image.Image, *, top: int, left: int, bottom: int, right: int):
     return image.crop((left, top, right, bottom))
 
 
 @op("Flip horizontally")
-def flip_horizontally(image: Image):
+def flip_horizontally(image: Image.Image):
     return image.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
 
 
 @op("Flip vertically")
-def flip_vertically(image: Image):
+def flip_vertically(image: Image.Image):
     return image.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
 
 
 @op("Blur")
-def blur(image: Image, *, radius: float = 5):
+def blur(image: Image.Image, *, radius: float = 5):
     return image.filter(ImageFilter.GaussianBlur(radius))
 
 
 @op("Detail")
-def detail(image: Image):
+def detail(image: Image.Image):
     return image.filter(ImageFilter.DETAIL)
 
 
 @op("Edge enhance")
-def edge_enhance(image: Image):
+def edge_enhance(image: Image.Image):
     return image.filter(ImageFilter.EDGE_ENHANCE)
 
 
 @op("To grayscale")
-def to_grayscale(image: Image):
+def to_grayscale(image: Image.Image):
     return image.convert("L")
 
 
 @op("View image", view="image")
-def view_image(image: Image):
+def view_image(image: Image.Image):
     buffered = io.BytesIO()
     image.save(buffered, format="webp")
     b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")