diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 00000000..ced19e23 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,6 @@ +[pytest] +minversion = 5.4 +testpaths = + tests +filterwarnings = + ignore::DeprecationWarning \ No newline at end of file diff --git a/setup.py b/setup.py index d046c49e..d34c9f7e 100644 --- a/setup.py +++ b/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=[ diff --git a/tests/test_signals.py b/tests/test_signals.py index 0b24e260..d0e719c2 100644 --- a/tests/test_signals.py +++ b/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 + diff --git a/vectorbt/portfolio/enums.py b/vectorbt/portfolio/enums.py index 2cee094c..87058160 100644 --- a/vectorbt/portfolio/enums.py +++ b/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', [ diff --git a/vectorbt/signals/enums.py b/vectorbt/signals/enums.py index d2937681..7ceafd1b 100644 --- a/vectorbt/signals/enums.py +++ b/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)} ``` """