Improve and refactor tests (#48)

Author: michaelchu

optopsy/core.py

@@ -46,7 +46,7 @@ def _calculate_profit_loss(data):
 def _calculate_profit_loss_pct(data):
     return data.assign(
         long_profit_pct=lambda r: round((r["exit"] - r["entry"]) / r["entry"], 2)
-    ).assign(short_profit_pct=lambda r: round((r["entry"] - r["exit"]) / r["exit"], 2))
+    ).assign(short_profit_pct=lambda r: round((r["entry"] - r["exit"]) / r["entry"], 2))
 
 
 def _select_final_output_column(data, cols, side):
@@ -81,7 +81,9 @@ def _cut_options_by_otm(data, otm_pct_interval, max_otm_pct_interval):
 
 
 def _group_by_intervals(data, cols, drop_na):
-    grouped_dataset = data.groupby(cols)[["long_profit_pct", "short_profit_pct"]].describe()
+    grouped_dataset = data.groupby(cols)[
+        ["long_profit_pct", "short_profit_pct"]
+    ].describe()
     grouped_dataset.columns = [
         "_".join(col).rstrip("_") for col in grouped_dataset.columns.values
     ]
@@ -94,48 +96,29 @@ def _group_by_intervals(data, cols, drop_na):
     return grouped_dataset
 
 
-def _evaluate(entries, exits):
-    return (
-        entries.merge(
-            right=exits,
-            on=["underlying_symbol", "option_type", "expiration", "strike"],
-            suffixes=("_entry", "_exit"),
-        )
-        # by default we use the midpoint spread price to calculate entry and exit costs
-        .assign(entry=lambda r: (r["bid_entry"] + r["ask_entry"]) / 2).assign(
-            exit=lambda r: (r["bid_exit"] + r["ask_exit"]) / 2
-        )
-    )
-
-
 def _evaluate_options(data, min_bid_ask, exit_dte):
-    """
-    Take all option entries and exits and calculate the profit/loss
-    per option chain.
-
-    Args:
-        data: Dataframe containing option chains
-
-    Returns:
-        DataFrame with each option chain evaluated in all entry and exit combinations
-
-    """
-
     # remove option chains that are worthless, it's unrealistic to enter
     # trades with worthless options
     entries = _remove_min_bid_ask(data, min_bid_ask)
 
     # to reduce unnecessary computation, filter for options with the desired exit DTE
     exits = _get(data, "dte", exit_dte)
-    evaluated_options = _evaluate(entries, exits)
 
     return (
-        evaluated_options.pipe(_remove_invalid_evaluated_options)
+        entries.merge(
+            right=exits,
+            on=["underlying_symbol", "option_type", "expiration", "strike"],
+            suffixes=("_entry", "_exit"),
+        )
+        # by default we use the midpoint spread price to calculate entry and exit costs
+        .assign(entry=lambda r: (r["bid_entry"] + r["ask_entry"]) / 2)
+        .assign(exit=lambda r: (r["bid_exit"] + r["ask_exit"]) / 2)
+        .pipe(_remove_invalid_evaluated_options)
         .pipe(_calculate_profit_loss)
     )[evaluated_cols]
 
 
-def _process_entries_and_exits(data, **kwargs):
+def _evaluate_all_options(data, **kwargs):
     return (
         data.pipe(_assign_dte)
         .pipe(_trim_data, "dte", kwargs["exit_dte"], kwargs["max_entry_dte"])

optopsy/definitions.py

@@ -16,74 +16,93 @@
 # columns of dataframe after generating strategy
 single_strike_internal_cols = [
     "underlying_symbol",
+    "underlying_price_entry",
+    "option_type",
     "expiration",
     "dte_entry",
-    "dte_range",
-    "otm_pct_range",
     "strike",
+    "entry",
+    "exit",
+    "long_profit",
+    "short_profit",
+    "long_profit_pct",
+    "short_profit_pct",
+]
+
+
+straddle_internal_cols = [
+    "underlying_symbol",
     "underlying_price_entry",
-    "underlying_price_exit",
+    "expiration",
+    "dte_entry",
+    "option_type_leg1",
+    "option_type_leg2",
+    "strike",
     "entry",
     "exit",
     "long_profit",
     "short_profit",
+    "long_profit_pct",
+    "short_profit_pct",
 ]
 
+
 double_strike_internal_cols = [
     "underlying_symbol",
+    "underlying_price_entry",
     "expiration",
     "dte_entry",
-    "dte_range",
-    "otm_pct_range_leg1",
+    "option_type_leg1",
     "strike_leg1",
-    "otm_pct_range_leg2",
+    "option_type_leg2",
     "strike_leg2",
-    "underlying_price_entry",
-    "underlying_price_exit",
     "entry",
     "exit",
     "long_profit",
     "short_profit",
+    "long_profit_pct",
+    "short_profit_pct",
 ]
 
 triple_strike_internal_cols = [
     "underlying_symbol",
+    "underlying_price_entry",
     "expiration",
     "dte_entry",
-    "dte_range",
-    "leg1_otm_pct_range",
-    "leg1_strike",
-    "leg2_otm_pct_range",
-    "leg2_strike",
-    "leg3_otm_pct_range",
-    "leg3_strike",
-    "underlying_price_entry",
-    "underlying_price_exit",
+    "option_type_leg1",
+    "strike_leg1",
+    "option_type_leg2",
+    "strike_leg2",
+    "option_type_leg3",
+    "strike_leg3",
     "entry",
     "exit",
     "long_profit",
     "short_profit",
+    "long_profit_pct",
+    "short_profit_pct",
 ]
 
 quadruple_strike_internal_cols = [
     "underlying_symbol",
+    "underlying_price_entry",
     "expiration",
     "dte_entry",
     "dte_range",
-    "leg1_otm_pct_range",
-    "leg1_strike",
-    "leg2_otm_pct_range",
-    "leg2_strike",
-    "leg3_otm_pct_range",
-    "leg3_strike",
-    "leg4_otm_pct_range",
-    "leg4_strike",
-    "underlying_price_entry",
-    "underlying_price_exit",
+    "option_type_leg1",
+    "strike_leg1",
+    "option_type_leg2",
+    "strike_leg2",
+    "option_type_leg3",
+    "strike_leg3",
+    "option_type_leg4",
+    "strike_leg4",
     "entry",
     "exit",
     "long_profit",
     "short_profit",
+    "long_profit_pct",
+    "short_profit_pct",
 ]
 
 # base columns of dataframe after aggregation(minus the calculated columns)

optopsy/strategies.py

@@ -1,5 +1,5 @@
 from .core import (
-    _process_entries_and_exits,
+    _evaluate_all_options,
     _calls,
     _puts,
     _group_by_intervals,
@@ -16,16 +16,16 @@
     "otm_pct_interval": 0.05,
     "max_otm_pct": 1,
     "min_bid_ask": 0.05,
-    "side": None,
+    "side": "long",
     "drop_nan": True,
     "raw": False,
     "on": None,
 }
 
 
-def _process_strategy(data, internal_cols, external_cols, strategy_func, params):
+def _process_strategy(data, internal_cols, external_cols, _strategy_func, params):
     return (
-        _process_entries_and_exits(
+        _evaluate_all_options(
             data,
             dte_interval=params["dte_interval"],
             max_entry_dte=params["max_entry_dte"],
@@ -34,15 +34,18 @@ def _process_strategy(data, internal_cols, external_cols, strategy_func, params)
             max_otm_pct=params["max_otm_pct"],
             min_bid_ask=params["min_bid_ask"],
         )
-        .pipe(strategy_func, params["on"], internal_cols)
-        .pipe(_format_output, params, external_cols)
+        .pipe(_strategy_func, params["on"])
+        .pipe(_calculate_profit_loss_pct)
+        .pipe(_format_output, params, internal_cols, external_cols)
     )
 
 
-def _format_output(data, params, external_cols):
+def _format_output(data, params, internal_cols, external_cols):
+    if params["raw"]:
+        return data[internal_cols]
+
     return (
-        data.pipe(_calculate_profit_loss_pct)
-        .pipe(_group_by_intervals, external_cols, params["drop_nan"])
+        data.pipe(_group_by_intervals, external_cols, params["drop_nan"])
         .reset_index()
         .pipe(
             _select_final_output_column,
@@ -52,7 +55,7 @@ def _format_output(data, params, external_cols):
     )
 
 
-def _straddles(data, on, internal_cols):
+def _straddles(data, on):
     calls = _calls(data)
     puts = _puts(data)
 
@@ -63,13 +66,15 @@ def _straddles(data, on, internal_cols):
             suffixes=("_leg1", "_leg2"),
         )
         .assign(long_profit=lambda r: (r["long_profit_leg1"] + r["long_profit_leg2"]))
-        .assign(short_profit=lambda r: (r["short_profit_leg1"] + r["short_profit_leg2"]))
+        .assign(
+            short_profit=lambda r: (r["short_profit_leg1"] + r["short_profit_leg2"])
+        )
         .assign(entry=lambda r: (r["entry_leg1"] + r["entry_leg2"]))
         .assign(exit=lambda r: (r["exit_leg1"] + r["exit_leg2"]))
-    )[internal_cols]
+    )
 
 
-def _strangles(data, on, internal_cols):
+def _strangles(data, on):
     calls = _calls(data)
     puts = _puts(data)
 
@@ -85,10 +90,12 @@ def _apply_strangle_rules(d):
         )
         .pipe(_apply_strangle_rules)
         .assign(long_profit=lambda r: (r["long_profit_leg1"] + r["long_profit_leg2"]))
-        .assign(short_profit=lambda r: (r["short_profit_leg1"] + r["short_profit_leg2"]))
+        .assign(
+            short_profit=lambda r: (r["short_profit_leg1"] + r["short_profit_leg2"])
+        )
         .assign(entry=lambda r: (r["entry_leg1"] + r["entry_leg2"]))
         .assign(exit=lambda r: (r["exit_leg1"] + r["exit_leg2"]))
-    )[internal_cols]
+    )
 
 
 def singles_calls(data, **kwargs):
@@ -97,7 +104,7 @@ def singles_calls(data, **kwargs):
         _process_strategy,
         single_strike_internal_cols,
         single_strike_external_cols,
-        lambda d, o, c: d,
+        lambda d, o: d,
         params,
     )
 
@@ -108,7 +115,7 @@ def singles_puts(data, **kwargs):
         _process_strategy,
         single_strike_internal_cols,
         single_strike_external_cols,
-        lambda d, o, c: d,
+        lambda d, o: d,
         params,
     )
 
@@ -128,7 +135,7 @@ def straddles(data, **kwargs):
     params = {**default_kwargs, **kwargs}
     return _process_strategy(
         data,
-        single_strike_internal_cols,
+        straddle_internal_cols,
         single_strike_external_cols,
         _straddles,
         params,

samples/spx_strangles_example.py

@@ -49,6 +49,7 @@ def run_strategy():
 
 if __name__ == "__main__":
     import timeit
+
     start = timeit.default_timer()
 
     # All the program statements

tests/integration/__init__.py

@@ -1,5 +1,5 @@
-col1,col2,col3,col4,col5,col6,col7,col8,col9
-SPX,359.69,call,1/20/1990,1/2/1990,225,135.5,135.5,0
-SPX,359.69,call,1/20/2000,1/2/2000,320,40.9,40.9,0
-SPX,359.69,call,1/20/2010,1/2/2010,325,35.9,35.9,0
-SPX,359.69,call,1/20/2020,1/2/2020,330,30.9,30.9,0
+col1,col2,col3,col4,col5,col6,col7,col8,col9
+SPX,359.69,call,1/20/1990,1/2/1990,225,135.5,135.5,0
+SPX,359.69,call,1/20/2000,1/2/2000,320,40.9,40.9,0
+SPX,359.69,call,1/20/2010,1/2/2010,325,35.9,35.9,0
+SPX,359.69,call,1/20/2020,1/2/2020,330,30.9,30.9,0