refactor: switch SignalTypes to StrEnum

This commit is contained in:
Matthias
2026-02-10 07:16:01 +01:00
parent 812f357d7e
commit 24319585b4
3 changed files with 20 additions and 26 deletions

View File

@@ -1,7 +1,7 @@
from enum import Enum, StrEnum
from enum import StrEnum
class SignalType(Enum):
class SignalType(StrEnum):
"""
Enum to distinguish between enter and exit signals
"""
@@ -11,11 +11,8 @@ class SignalType(Enum):
ENTER_SHORT = "enter_short"
EXIT_SHORT = "exit_short"
def __str__(self):
return f"{self.name.lower()}"
class SignalTagType(Enum):
class SignalTagType(StrEnum):
"""
Enum for signal columns
"""
@@ -23,9 +20,6 @@ class SignalTagType(Enum):
ENTER_TAG = "enter_tag"
EXIT_TAG = "exit_tag"
def __str__(self):
return f"{self.name.lower()}"
class SignalDirection(StrEnum):
LONG = "long"

View File

@@ -210,12 +210,12 @@ def remove_entry_exit_signals(dataframe: pd.DataFrame):
:param dataframe: The DataFrame to remove signals from
"""
dataframe[SignalType.ENTER_LONG.value] = 0
dataframe[SignalType.EXIT_LONG.value] = 0
dataframe[SignalType.ENTER_SHORT.value] = 0
dataframe[SignalType.EXIT_SHORT.value] = 0
dataframe[SignalTagType.ENTER_TAG.value] = None
dataframe[SignalTagType.EXIT_TAG.value] = None
dataframe[SignalType.ENTER_LONG] = 0
dataframe[SignalType.EXIT_LONG] = 0
dataframe[SignalType.ENTER_SHORT] = 0
dataframe[SignalType.EXIT_SHORT] = 0
dataframe[SignalTagType.ENTER_TAG] = None
dataframe[SignalTagType.EXIT_TAG] = None
return dataframe

View File

@@ -1329,13 +1329,13 @@ class IStrategy(ABC, HyperStrategyMixin):
return False, False, None
if is_short:
enter = latest.get(SignalType.ENTER_SHORT.value, 0) == 1
exit_ = latest.get(SignalType.EXIT_SHORT.value, 0) == 1
enter = latest.get(SignalType.ENTER_SHORT, 0) == 1
exit_ = latest.get(SignalType.EXIT_SHORT, 0) == 1
else:
enter = latest.get(SignalType.ENTER_LONG.value, 0) == 1
exit_ = latest.get(SignalType.EXIT_LONG.value, 0) == 1
exit_tag = latest.get(SignalTagType.EXIT_TAG.value, None)
enter = latest.get(SignalType.ENTER_LONG, 0) == 1
exit_ = latest.get(SignalType.EXIT_LONG, 0) == 1
exit_tag = latest.get(SignalTagType.EXIT_TAG, None)
# Tags can be None, which does not resolve to False.
exit_tag = exit_tag if isinstance(exit_tag, str) and exit_tag != "nan" else None
@@ -1362,16 +1362,16 @@ class IStrategy(ABC, HyperStrategyMixin):
if latest is None or latest_date is None:
return None, None
enter_long = latest.get(SignalType.ENTER_LONG.value, 0) == 1
exit_long = latest.get(SignalType.EXIT_LONG.value, 0) == 1
enter_short = latest.get(SignalType.ENTER_SHORT.value, 0) == 1
exit_short = latest.get(SignalType.EXIT_SHORT.value, 0) == 1
enter_long = latest.get(SignalType.ENTER_LONG, 0) == 1
exit_long = latest.get(SignalType.EXIT_LONG, 0) == 1
enter_short = latest.get(SignalType.ENTER_SHORT, 0) == 1
exit_short = latest.get(SignalType.EXIT_SHORT, 0) == 1
enter_signal: SignalDirection | None = None
enter_tag: str | None = None
if enter_long == 1 and not any([exit_long, enter_short]):
enter_signal = SignalDirection.LONG
enter_tag = latest.get(SignalTagType.ENTER_TAG.value, None)
enter_tag = latest.get(SignalTagType.ENTER_TAG, None)
if (
self.config.get("trading_mode", TradingMode.SPOT) != TradingMode.SPOT
and self.can_short
@@ -1379,7 +1379,7 @@ class IStrategy(ABC, HyperStrategyMixin):
and not any([exit_short, enter_long])
):
enter_signal = SignalDirection.SHORT
enter_tag = latest.get(SignalTagType.ENTER_TAG.value, None)
enter_tag = latest.get(SignalTagType.ENTER_TAG, None)
enter_tag = enter_tag if isinstance(enter_tag, str) and enter_tag != "nan" else None