get it right pyright

reflex/vars/base.py

@@ -1491,9 +1491,6 @@ class TypeComputer(Protocol):
 
         Args:
             *args: The arguments to compute the type of.
-
-        Returns:
-            The type of the operation.
         """
         ...
 
@@ -1554,6 +1551,9 @@ def var_operation(
 
     Returns:
         The decorated function.
+
+    Raises:
+        TypeError: If the function has keyword-only arguments or arguments without Var type hints.
     """
     from .function import ArgsFunctionOperation, ReflexCallable
 

reflex/vars/function.py

@@ -416,9 +416,6 @@ def format_args_function_operation(
 
     Args:
         self: The function operation.
-        args: The function arguments.
-        return_expr: The return expression.
-        explicit_return: Whether to use explicit return syntax.
 
     Returns:
         The formatted args function operation.
@@ -454,6 +451,9 @@ def pre_check_args(
 
     Returns:
         True if the function can be called with the given arguments.
+
+    Raises:
+        VarTypeError: If the arguments are invalid.
     """
     for i, (validator, arg) in enumerate(zip(self._validators, args)):
         if not validator(arg):

reflex/vars/number.py

@@ -367,7 +367,7 @@ class NumberVar(Var[NUMBER_T], python_types=(int, float)):
         Returns:
             The boolean NOT operation.
         """
-        return boolean_not_operation(self.bool())
+        return boolean_not_operation(self.bool()).guess_type()
 
     def __pos__(self) -> NumberVar:
         """Positive the number.
@@ -518,7 +518,7 @@ class NumberVar(Var[NUMBER_T], python_types=(int, float)):
             The boolean value of the number.
         """
         if is_optional(self._var_type):
-            return boolify((self != None) & (self != 0))  # noqa: E711
+            return boolify((self != None) & (self != 0)).guess_type()  # noqa: E711
         return self != 0
 
     def _is_strict_float(self) -> bool:
@@ -777,7 +777,7 @@ class BooleanVar(NumberVar[bool], python_types=bool):
         Returns:
             The boolean NOT operation.
         """
-        return boolean_not_operation(self)
+        return boolean_not_operation(self).guess_type()
 
     def __int__(self):
         """Convert the boolean to an int.
@@ -785,7 +785,7 @@ class BooleanVar(NumberVar[bool], python_types=bool):
         Returns:
             The boolean to int operation.
         """
-        return boolean_to_number_operation(self)
+        return boolean_to_number_operation(self).guess_type()
 
     def __pos__(self):
         """Convert the boolean to an int.
@@ -793,7 +793,7 @@ class BooleanVar(NumberVar[bool], python_types=bool):
         Returns:
             The boolean to int operation.
         """
-        return boolean_to_number_operation(self)
+        return boolean_to_number_operation(self).guess_type()
 
     def bool(self) -> BooleanVar:
         """Boolean conversion.

reflex/vars/sequence.py

@@ -12,6 +12,7 @@ import typing
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     Dict,
     List,
     Literal,
@@ -20,10 +21,11 @@ from typing import (
     Tuple,
     Type,
     Union,
+    cast,
     overload,
 )
 
-from typing_extensions import TypeVar
+from typing_extensions import TypeAliasType, TypeVar
 
 from reflex import constants
 from reflex.constants.base import REFLEX_VAR_OPENING_TAG
@@ -59,6 +61,7 @@ from .number import (
 )
 
 if TYPE_CHECKING:
+    from .function import FunctionVar
     from .object import ObjectVar
 
 STRING_TYPE = TypeVar("STRING_TYPE", default=str)
@@ -751,6 +754,7 @@ ARRAY_VAR_TYPE = TypeVar("ARRAY_VAR_TYPE", bound=Union[List, Tuple, Set])
 OTHER_TUPLE = TypeVar("OTHER_TUPLE")
 
 INNER_ARRAY_VAR = TypeVar("INNER_ARRAY_VAR")
+ANOTHER_ARRAY_VAR = TypeVar("ANOTHER_ARRAY_VAR")
 
 KEY_TYPE = TypeVar("KEY_TYPE")
 VALUE_TYPE = TypeVar("VALUE_TYPE")
@@ -939,7 +943,7 @@ class ArrayVar(Var[ARRAY_VAR_TYPE], python_types=(list, tuple, set)):
             isinstance(i, NumberVar) and i._is_strict_float()
         ):
             raise_unsupported_operand_types("[]", (type(self), type(i)))
-        return array_item_operation(self, i)
+        return array_item_operation(self, i).guess_type()
 
     def length(self) -> NumberVar:
         """Get the length of the array.
@@ -1143,7 +1147,11 @@ class ArrayVar(Var[ARRAY_VAR_TYPE], python_types=(list, tuple, set)):
 
         return array_ge_operation(self, other).guess_type()
 
-    def foreach(self, fn: Any):
+    def foreach(
+        self: ARRAY_VAR_OF_LIST_ELEMENT[INNER_ARRAY_VAR],
+        fn: Callable[[Var[INNER_ARRAY_VAR]], ANOTHER_ARRAY_VAR]
+        | Callable[[], ANOTHER_ARRAY_VAR],
+    ) -> ArrayVar[List[ANOTHER_ARRAY_VAR]]:
         """Apply a function to each element of the array.
 
         Args:
@@ -1167,37 +1175,49 @@ class ArrayVar(Var[ARRAY_VAR_TYPE], python_types=(list, tuple, set)):
             )
 
         if num_args == 0:
-            return_value = fn()
-            function_var = ArgsFunctionOperation.create(tuple(), return_value)
-        else:
-            # generic number var
-            number_var = Var("").to(NumberVar, int)
+            return_value = fn()  # type: ignore
+            simple_function_var: FunctionVar[ReflexCallable[[], ANOTHER_ARRAY_VAR]] = (
+                ArgsFunctionOperation.create(tuple(), return_value)
+            )
+            return map_array_operation(self, simple_function_var).guess_type()
+
+        # generic number var
+        number_var = Var("").to(NumberVar, int)
 
-            first_arg_type = self[number_var]._var_type
+        first_arg_type = self[number_var]._var_type
 
-            arg_name = get_unique_variable_name()
+        arg_name = get_unique_variable_name()
 
-            # get first argument type
-            first_arg = Var(
+        # get first argument type
+        first_arg = cast(
+            Var[Any],
+            Var(
                 _js_expr=arg_name,
                 _var_type=first_arg_type,
-            ).guess_type()
+            ).guess_type(),
+        )
 
-            function_var = ArgsFunctionOperation.create(
-                (arg_name,),
-                Var.create(fn(first_arg)),
-            )
+        function_var: FunctionVar[
+            ReflexCallable[[INNER_ARRAY_VAR], ANOTHER_ARRAY_VAR]
+        ] = ArgsFunctionOperation.create(
+            (arg_name,),
+            Var.create(fn(first_arg)),  # type: ignore
+        )
 
-        return map_array_operation(self, function_var)
+        return map_array_operation(self, function_var).guess_type()
 
 
 LIST_ELEMENT = TypeVar("LIST_ELEMENT")
 
-ARRAY_VAR_OF_LIST_ELEMENT = Union[
-    ArrayVar[List[LIST_ELEMENT]],
-    ArrayVar[Set[LIST_ELEMENT]],
-    ArrayVar[Tuple[LIST_ELEMENT, ...]],
-]
+ARRAY_VAR_OF_LIST_ELEMENT = TypeAliasType(
+    "ARRAY_VAR_OF_LIST_ELEMENT",
+    Union[
+        ArrayVar[List[LIST_ELEMENT]],
+        ArrayVar[Tuple[LIST_ELEMENT, ...]],
+        ArrayVar[Set[LIST_ELEMENT]],
+    ],
+    type_params=(LIST_ELEMENT,),
+)
 
 
 @dataclasses.dataclass(
@@ -1663,9 +1683,9 @@ if TYPE_CHECKING:
 
 @var_operation
 def map_array_operation(
-    array: Var[ARRAY_VAR_TYPE],
-    function: Var[ReflexCallable],
-):
+    array: Var[ARRAY_VAR_OF_LIST_ELEMENT[INNER_ARRAY_VAR]],
+    function: Var[ReflexCallable[[INNER_ARRAY_VAR], ANOTHER_ARRAY_VAR]],
+) -> CustomVarOperationReturn[List[ANOTHER_ARRAY_VAR]]:
     """Map a function over an array.
 
     Args: