From 56709a210b0b722c05b72dc515634d19afb1586b Mon Sep 17 00:00:00 2001
From: Khaleel Al-Adhami <khaleel.aladhami@gmail.com>
Date: Thu, 3 Oct 2024 13:01:19 -0700
Subject: [PATCH] add of_type to _evaluate (#4051)

* add of_type to _evaluate

* get it right pyright
---
 reflex/state.py                              | 20 ++++++++++++++++----
 tests/integration/test_dynamic_components.py |  4 +++-
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/reflex/state.py b/reflex/state.py
index 64ea960e1..1746835be 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -699,11 +699,14 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         )
 
     @classmethod
-    def _evaluate(cls, f: Callable[[Self], Any]) -> Var:
+    def _evaluate(
+        cls, f: Callable[[Self], Any], of_type: Union[type, None] = None
+    ) -> Var:
         """Evaluate a function to a ComputedVar. Experimental.
 
         Args:
             f: The function to evaluate.
+            of_type: The type of the ComputedVar. Defaults to Component.
 
         Returns:
             The ComputedVar.
@@ -711,14 +714,23 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         console.warn(
             "The _evaluate method is experimental and may be removed in future versions."
         )
-        from reflex.components.base.fragment import fragment
         from reflex.components.component import Component
 
+        of_type = of_type or Component
+
         unique_var_name = get_unique_variable_name()
 
-        @computed_var(_js_expr=unique_var_name, return_type=Component)
+        @computed_var(_js_expr=unique_var_name, return_type=of_type)
         def computed_var_func(state: Self):
-            return fragment(f(state))
+            result = f(state)
+
+            if not isinstance(result, of_type):
+                console.warn(
+                    f"Inline ComputedVar {f} expected type {of_type}, got {type(result)}. "
+                    "You can specify expected type with `of_type` argument."
+                )
+
+            return result
 
         setattr(cls, unique_var_name, computed_var_func)
         cls.computed_vars[unique_var_name] = computed_var_func
diff --git a/tests/integration/test_dynamic_components.py b/tests/integration/test_dynamic_components.py
index 5a4d99f9e..aeebd10e9 100644
--- a/tests/integration/test_dynamic_components.py
+++ b/tests/integration/test_dynamic_components.py
@@ -65,7 +65,9 @@ def DynamicComponents():
             DynamicComponentsState.client_token_component,
             DynamicComponentsState.button,
             rx.text(
-                DynamicComponentsState._evaluate(lambda state: factorial(state.value)),
+                DynamicComponentsState._evaluate(
+                    lambda state: factorial(state.value), of_type=int
+                ),
                 id="factorial",
             ),
         )