diff --git a/backtester/option.py b/backtester/option.py index e0c597c..d382aa0 100644 --- a/backtester/option.py +++ b/backtester/option.py @@ -1,14 +1,20 @@ # Option Enum types from enum import Enum -Type = Enum("Type", {"CALL": "call", "PUT": "put"}) -Direction = Enum("Direction", {"BUY": "ask", "SELL": "bid"}) + +class Type(Enum): + CALL = 'call' + PUT = 'put' + + def __invert__(self): + flip = Type.PUT if self == Type.CALL else Type.CALL + return flip -def flip(enum_type): - """Returns the other value of `Type` or `Direction`""" - assert isinstance(enum_type, Type) or isinstance(enum_type, Direction) +class Direction(Enum): + BUY = 'ask' + SELL = 'bid' - for e in type(enum_type): - if e != enum_type: - return e + def __invert__(self): + flip = Direction.SELL if self == Direction.BUY else Direction.SELL + return flip diff --git a/backtester/strategy/strategy.py b/backtester/strategy/strategy.py index 34c5d2a..a5c1223 100644 --- a/backtester/strategy/strategy.py +++ b/backtester/strategy/strategy.py @@ -3,7 +3,6 @@ from collections import namedtuple import pandas as pd from backtester.datahandler import Schema -from backtester.option import flip from .strategy_leg import StrategyLeg from .signal import Signal, get_order @@ -95,7 +94,7 @@ class Strategy: price = leg.direction.value else: flt = leg.exit_filter - price = flip(leg.direction).value + price = (~leg.direction).value df = flt(data) fields = {