get it right pyright

This commit is contained in:
Khaleel Al-Adhami 2024-11-13 19:03:42 -08:00
parent 3cdd2097b6
commit a9db61b371
4 changed files with 57 additions and 37 deletions

View File

@ -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

View File

@ -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):

View File

@ -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.

View File

@ -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()
first_arg_type = self[number_var]._var_type
# generic number var
number_var = Var("").to(NumberVar, int)
arg_name = get_unique_variable_name()
first_arg_type = self[number_var]._var_type
# get first argument type
first_arg = Var(
arg_name = get_unique_variable_name()
# 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: