Skip to content

Fix LSP type checking for ops wrapped with @api #30

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
fferflo merged 1 commit into from
Mar 31, 2026
Merged

Fix LSP type checking for ops wrapped with @api #30

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 17 additions & 6 deletions einx/_src/frontend/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,21 @@
import functools
import inspect
from .types import Tensor
import traceback
import types
from collections.abc import Callable
from functools import partial
from typing import Any, TypeVar, overload

import numpy as np

import einx._src.tracer as tracer
from einx._src.frontend.errors import CallOperationError
from einx._src.util.lru_cache import lru_cache

from .backend import registry
import numpy as np
from functools import partial
import types
import traceback
from einx._src.frontend.errors import CallOperationError
from .types import Tensor

_F = TypeVar("_F", bound=Callable[..., Any])


class TensorArg:
Expand Down Expand Up @@ -217,6 +224,10 @@ def inner(*args, graph=False, **kwargs):
return inner


@overload
def api(func: _F) -> _F: ...
@overload
def api(func: None = None, backend: Any = None, signature: Any = None) -> Callable[[_F], _F]: ...
def api(func=None, backend=None, signature=None):
if func is None:
return partial(api, backend=backend, signature=signature)
Expand Down