Skip to content

Commit

Permalink
Fix missing type hints for @content_component (#1146)
Browse files Browse the repository at this point in the history
Closes #912
  • Loading branch information
richard-to authored Dec 14, 2024
1 parent a53a70d commit a6efb85
Showing 1 changed file with 12 additions and 4 deletions.
16 changes: 12 additions & 4 deletions mesop/component_helpers/helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,9 @@
Any,
Callable,
Generator,
Generic,
KeysView,
ParamSpec,
Type,
TypeVar,
cast,
Expand Down Expand Up @@ -66,8 +68,12 @@ def slot():
runtime().context().save_current_node_as_slot()


class _UserCompositeComponent:
def __init__(self, fn: Callable[..., Any]):
T = TypeVar("T")
P = ParamSpec("P")


class _UserCompositeComponent(Generic[T]):
def __init__(self, fn: Callable[[], T]):
self.prev_current_node = runtime().context().current_node()
fn()
node_slot = runtime().context().node_slot()
Expand All @@ -93,9 +99,11 @@ def __exit__(self, exc_type, exc_val, exc_tb): # type: ignore
runtime().context().set_current_node(self.prev_current_node)


def content_component(fn: Callable[..., Any]):
def content_component(
fn: Callable[P, T],
) -> Callable[P, _UserCompositeComponent[T]]:
@wraps(fn)
def wrapper(*args: Any, **kwargs: Any):
def wrapper(*args: P.args, **kwargs: P.kwargs) -> _UserCompositeComponent[T]:
return _UserCompositeComponent(lambda: fn(*args, **kwargs))

return wrapper
Expand Down

0 comments on commit a6efb85

Please sign in to comment.