Implement BoolValue.select

Author: FancyNeuron

src/spellbind/bool_values.py

@@ -2,13 +2,25 @@
 
 import operator
 from abc import ABC
-from typing import TypeVar, Generic
+from typing import TypeVar, Generic, overload, TYPE_CHECKING, TypeAlias
+from spellbind.values import Value, OneToOneValue, Constant, SimpleVariable, TwoToOneValue, SelectValue
 
-from spellbind.values import Value, OneToOneValue, Constant, SimpleVariable, TwoToOneValue
+if TYPE_CHECKING:
+    from spellbind.float_values import FloatValue  # pragma: no cover
+    from spellbind.int_values import IntValue  # pragma: no cover
+    from spellbind.str_values import StrValue  # pragma: no cover
+
+IntValueLike: TypeAlias = 'IntValue | int'
+FloatValueLike: TypeAlias = 'FloatValue | float'
+StrValueLike: TypeAlias = 'StrValue | str'
+BoolValueLike: TypeAlias = 'BoolValue | bool'
 
 _S = TypeVar('_S')
 
-BoolLike = Value[bool] | bool
+BoolLike = bool | Value[bool]
+IntLike = int | Value[int]
+FloatLike = float | Value[float]
+StrLike = str | Value[str]
 
 
 class BoolValue(Value[bool], ABC):
@@ -33,6 +45,37 @@ def __xor__(self, other: BoolLike) -> BoolValue:
     def __rxor__(self, other: bool) -> BoolValue:
         return XorBoolValues(other, self)
 
+    @overload
+    def select(self, if_true: IntValueLike, if_false: IntValueLike) -> IntValue: ...
+
+    @overload
+    def select(self, if_true: FloatValueLike, if_false: FloatValueLike) -> FloatValue: ...
+
+    @overload
+    def select(self, if_true: StrValueLike, if_false: StrValueLike) -> StrValue: ...
+
+    @overload
+    def select(self, if_true: BoolValue, if_false: BoolValue) -> BoolValue: ...
+
+    @overload
+    def select(self, if_true: Value[_S] | _S, if_false: Value[_S] | _S) -> Value[_S]: ...
+
+    def select(self, if_true, if_false):
+        from spellbind.float_values import FloatValue, SelectFloatValue
+        from spellbind.int_values import IntValue, SelectIntValue
+        from spellbind.str_values import StrValue, SelectStrValue
+
+        if isinstance(if_true, (FloatValue, float)) and isinstance(if_false, (FloatValue, float)):
+            return SelectFloatValue(self, if_true, if_false)
+        elif isinstance(if_true, (StrValue, str)) and isinstance(if_false, (StrValue, str)):
+            return SelectStrValue(self, if_true, if_false)
+        elif isinstance(if_true, (BoolValue, bool)) and isinstance(if_false, (BoolValue, bool)):
+            return SelectBoolValue(self, if_true, if_false)
+        elif isinstance(if_true, (IntValue, int)) and isinstance(if_false, (IntValue, int)):
+            return SelectIntValue(self, if_true, if_false)
+        else:
+            return SelectValue(self, if_true, if_false)
+
 
 class OneToBoolValue(OneToOneValue[_S, bool], BoolValue, Generic[_S]):
     pass
@@ -66,5 +109,10 @@ class BoolVariable(SimpleVariable[bool], BoolValue):
     pass
 
 
+class SelectBoolValue(SelectValue[bool], BoolValue):
+    def __init__(self, condition: BoolLike, if_true: BoolLike, if_false: BoolLike):
+        super().__init__(condition, if_true, if_false)
+
+
 TRUE = BoolConstant(True)
 FALSE = BoolConstant(False)

src/spellbind/float_values.py

@@ -7,9 +7,10 @@
 from typing_extensions import Self
 from typing_extensions import TYPE_CHECKING
 
-from spellbind.bool_values import BoolValue
+from spellbind.bool_values import BoolValue, BoolLike
 from spellbind.functions import clamp_float, multiply_all_floats
-from spellbind.values import Value, SimpleVariable, OneToOneValue, DerivedValueBase, Constant, TwoToOneValue
+from spellbind.values import Value, SimpleVariable, OneToOneValue, DerivedValueBase, Constant, TwoToOneValue, \
+    SelectValue
 
 if TYPE_CHECKING:
     from spellbind.int_values import IntValue, IntLike  # pragma: no cover
@@ -235,3 +236,8 @@ def __init__(self, left: FloatLike, right: FloatLike, op: Callable[[float, float
 class ClampFloatValue(ThreeFloatToOneValue[float], FloatValue):
     def __init__(self, value: FloatLike, min_value: FloatLike, max_value: FloatLike):
         super().__init__(clamp_float, value, min_value, max_value)
+
+
+class SelectFloatValue(SelectValue[float], FloatValue):
+    def __init__(self, condition: BoolLike, if_true: float | Value[float], if_false: float | Value[float]):
+        super().__init__(condition, if_true, if_false)

src/spellbind/int_values.py

@@ -7,12 +7,12 @@
 
 from typing_extensions import Self, TypeVar
 
-from spellbind.bool_values import BoolValue
+from spellbind.bool_values import BoolValue, BoolLike
 from spellbind.float_values import FloatValue, MultiplyFloatValues, DivideValues, SubtractFloatValues, \
     AddFloatValues, CompareNumbersValues
 from spellbind.functions import clamp_int, multiply_all_ints
 from spellbind.values import Value, ManyToOneValue, SimpleVariable, TwoToOneValue, OneToOneValue, Constant, \
-    ThreeToOneValue
+    ThreeToOneValue, SelectValue
 
 IntLike = int | Value[int]
 FloatLike = IntLike | float | FloatValue
@@ -217,3 +217,8 @@ def __init__(self, value: Value[float]):
 class ClampIntValue(ThreeToOneValue[int, int, int, int], IntValue):
     def __init__(self, value: IntLike, min_value: IntLike, max_value: IntLike) -> None:
         super().__init__(clamp_int, value, min_value, max_value)
+
+
+class SelectIntValue(SelectValue[int], IntValue):
+    def __init__(self, condition: BoolLike, if_true: IntLike, if_false: IntLike):
+        super().__init__(condition, if_true, if_false)

src/spellbind/str_values.py

@@ -3,7 +3,8 @@
 from abc import ABC
 from typing import Any, Generic, TypeVar
 
-from spellbind.values import Value, OneToOneValue, ManyToOneValue, SimpleVariable, Constant
+from spellbind.bool_values import BoolLike, StrLike
+from spellbind.values import Value, OneToOneValue, ManyToOneValue, SimpleVariable, Constant, SelectValue
 
 StringLike = str | Value[str]
 
@@ -41,3 +42,8 @@ def __init__(self, value: Value[Any]):
 class ConcatenateStrValues(ManyToOneValue[str, str], StrValue):
     def __init__(self, *values: StringLike):
         super().__init__(''.join, *values)
+
+
+class SelectStrValue(SelectValue[str], StrValue):
+    def __init__(self, condition: BoolLike, if_true: StrLike, if_false: StrLike):
+        super().__init__(condition, if_true, if_false)

src/spellbind/values.py

@@ -9,7 +9,7 @@
 if TYPE_CHECKING:
     from spellbind.str_values import StrValue  # pragma: no cover
     from spellbind.int_values import IntValue  # pragma: no cover
-    from spellbind.bool_values import BoolValue  # pragma: no cover
+    from spellbind.bool_values import BoolValue, BoolLike  # pragma: no cover
 
 
 EMPTY_FROZEN_SET: frozenset = frozenset()
@@ -275,3 +275,8 @@ def __init__(self, transformer: Callable[[_S, _T, _U], _V],
 
     def _calculate_value(self) -> _V:
         return self._transformer(self._first_getter(), self._second_getter(), self._third_getter())
+
+
+class SelectValue(ThreeToOneValue[bool, _S, _S, _S], Generic[_S]):
+    def __init__(self, condition: BoolLike, if_true: Value[_S] | _S, if_false: Value[_S] | _S):
+        super().__init__(lambda b, t, f: t if b else f, condition, if_true, if_false)