DX enhancements

pytest.ini

@@ -0,0 +1,6 @@
+[pytest]
+minversion = 5.4
+testpaths =
+    tests
+filterwarnings =
+    ignore::DeprecationWarning

setup.py

@@ -25,6 +25,9 @@
         'ipywidgets>=7.0.0',
         'numba>=0.51.2'
     ],
+    extras_require={
+        "empyrical":["empyrical"]
+    },
     python_requires='>=3.6, <3.9',
     license='Apache 2.0',
     classifiers=[

tests/test_signals.py

@@ -1,4 +1,7 @@
+from vectorbt.signals.enums import StopType
 import vectorbt as vbt
+from enum import Enum, IntEnum
+from typing import Union
 import numpy as np
 import pandas as pd
 from numba import njit
@@ -2291,3 +2294,12 @@ def test_IOHLCSTEX(self):
             )
         )
 
+def test_enum_in_numba():
+    @njit
+    def try_enum(my_enum:StopType) -> int:
+        return my_enum.value + 1
+
+    assert try_enum(StopType.StopLoss) == 1
+    assert try_enum(StopType.TrailStop) == 2
+    assert try_enum(StopType.TakeProfit) == 3
+

vectorbt/portfolio/enums.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from collections import namedtuple
+from typing import NamedTuple, Any
 import json
 
 __all__ = [
@@ -39,20 +40,20 @@
 
 # ############# Portfolio ############# #
 
-SimulationContext = namedtuple('SimulationContext', [
-    'target_shape',
-    'close',
-    'group_lens',
-    'init_cash',
-    'cash_sharing',
-    'call_seq',
-    'active_mask',
-    'order_records',
-    'log_records',
-    'last_cash',
-    'last_shares',
-    'last_val_price'
-])
+class SimulationContext(NamedTuple):
+    target_shape: Any
+    close: Any
+    group_lens: Any
+    init_cash: Any
+    cash_sharing: Any
+    call_seq: Any
+    active_mask: Any
+    order_records: Any
+    log_records: Any
+    last_cash: Any
+    last_shares: Any
+    last_val_price: Any
+
 
 __pdoc__['SimulationContext'] = "A named tuple representing context of the simulation."
 __pdoc__['SimulationContext.target_shape'] = """Target shape.
@@ -69,13 +70,13 @@
 """
 __pdoc__['SimulationContext.init_cash'] = """Initial capital per column, or per group if cash sharing is enabled.
 
-If `cash_sharing` is True, has shape `(target_shape[0], group_lens.shape[0])`. 
+If `cash_sharing` is True, has shape `(target_shape[0], group_lens.shape[0])`.
 Otherwise, has shape `target_shape`.
 """
 __pdoc__['SimulationContext.cash_sharing'] = """Whether cash sharing is enabled."""
 __pdoc__['SimulationContext.call_seq'] = """Default sequence of calls per segment.
 
-Controls the sequence in which `order_func_nb` is executed within a segment. 
+Controls the sequence in which `order_func_nb` is executed within a segment.
 
 Has shape `target_shape` and each value must exist in the range `[0, group_len)`.
 
@@ -130,7 +131,7 @@
 """
 __pdoc__['GroupContext.to_col'] = """Index of the first column in the next group.
 
-Has range `[1, target_shape[1] + 1)`. 
+Has range `[1, target_shape[1] + 1)`.
 
 If columns are not grouped, equals `from_col + 1`.
 """
@@ -166,7 +167,7 @@
 __pdoc__['SegmentContext.to_col'] = "See `GroupContext.to_col`."
 __pdoc__['SegmentContext.call_seq_now'] = """Current sequence of calls.
 
-Has shape `(group_len,)`. 
+Has shape `(group_len,)`.
 """
 
 OrderContext = namedtuple('OrderContext', [

vectorbt/signals/enums.py

@@ -1,6 +1,10 @@
+#%%
 """Named tuples and enumerated types."""
 
 from collections import namedtuple
+from enum import IntEnum,Enum
+from typing import Union, Type
+
 import json
 
 __all__ = [
@@ -9,18 +13,25 @@
 
 __pdoc__ = {}
 
-# We use namedtuple for enums and classes to be able to use them in Numba
 
-StopType = namedtuple('StopType', [
-    'StopLoss',
-    'TrailStop',
-    'TakeProfit'
-])(*range(3))
-"""_"""
+class StopType(IntEnum):
+    """_"""
+    StopLoss = 0
+    TrailStop = 1
+    TakeProfit = 2
+
+def _enum_to_json(enum:Union[Type[Enum],Type[IntEnum]]) -> str:
+    """Render an (Int)Enum class to json
+
+    Returns:
+        str: A json string
+    """
+    return json.dumps({enum_option.name:enum_option.value for enum_option in StopType}, indent=2,default=str)
+
 
 __pdoc__['StopType'] = f"""Stop type.
 
 ```plaintext
-{json.dumps(dict(zip(StopType._fields, StopType)), indent=2, default=str)}
+{_enum_to_json(StopType)}
 ```
 """