Adds long-only and short-only config options.

Author: edtechre

src/pybroker/common.py

@@ -173,6 +173,20 @@ class FeeInfo(NamedTuple):
     order_type: Literal["buy", "sell"]
 
 
+class PositionMode(Enum):
+    """Position mode for backtesting.
+
+    Attributes:
+        DEFAULT: Long and short positions.
+        LONG_ONLY: Long-only positions.
+        SHORT_ONLY: Short-only positions.
+    """
+
+    DEFAULT = "default"
+    LONG_ONLY = "long_only"
+    SHORT_ONLY = "short_only"
+
+
 class BarData:
     r"""Contains data for a series of bars. Each field is a
     :class:`numpy.ndarray` that contains bar values in the series. The values

src/pybroker/config.py

@@ -6,7 +6,7 @@
 (see LICENSE for details).
 """
 
-from pybroker.common import BarData, FeeInfo, FeeMode, PriceType
+from pybroker.common import BarData, FeeInfo, FeeMode, PositionMode, PriceType
 from dataclasses import dataclass, field
 from decimal import Decimal
 from typing import Callable, Optional, Union
@@ -35,6 +35,12 @@ class StrategyConfig:
             Set to ``True`` for crypto trading. Defaults to ``False``.
         round_fill_price: Whether to round fill prices to the nearest cent.
             Defaults to ``True``.
+        position_mode: Position mode for :class:`pybroker.strategy.Strategy`.
+            Supports one of:
+
+            - ``DEFAULT``: Long and short positions.
+            - ``LONG_ONLY``: Long-only positions.
+            - ``SHORT_ONLY``: Short-only positions.
         max_long_positions: Maximum number of long positions that can be held
             at any time in :class:`pybroker.portfolio.Portfolio`. Unlimited
             when ``None``. Defaults to ``None``.
@@ -81,6 +87,7 @@ class StrategyConfig:
     subtract_fees: bool = field(default=False)
     enable_fractional_shares: bool = field(default=False)
     round_fill_price: bool = field(default=True)
+    position_mode: PositionMode = field(default=PositionMode.DEFAULT)
     max_long_positions: Optional[int] = field(default=None)
     max_short_positions: Optional[int] = field(default=None)
     buy_delay: int = field(default=1)

src/pybroker/portfolio.py

@@ -16,6 +16,7 @@
     DataCol,
     FeeInfo,
     FeeMode,
+    PositionMode,
     PriceType,
     StopType,
     to_decimal,
@@ -342,6 +343,7 @@ class Portfolio:
         subtract_fees: Whether to subtract fees from the cash balance after an
             order is filled.
         enable_fractional_shares: Whether to enable trading fractional shares.
+        position_mode: Position mode for :class:`.Portfolio`.
         max_long_positions: Maximum number of long :class:`.Position`\ s that
             can be held at a time. If ``None``, then unlimited.
         max_short_positions: Maximum number of short :class:`.Position`\ s that
@@ -384,6 +386,7 @@ def __init__(
         fee_amount: Optional[float] = None,
         subtract_fees: bool = False,
         enable_fractional_shares: bool = False,
+        position_mode: PositionMode = PositionMode.DEFAULT,
         max_long_positions: Optional[int] = None,
         max_short_positions: Optional[int] = None,
         record_stops: Optional[bool] = False,
@@ -396,6 +399,7 @@ def __init__(
         )
         self._subtract_fees = subtract_fees
         self._enable_fractional_shares = enable_fractional_shares
+        self._position_mode = position_mode
         self.equity: Decimal = self.cash
         self.market_value: Decimal = self.cash
         self.fees = Decimal()
@@ -639,7 +643,7 @@ def buy(
         if shares == 0:
             return None
         covered = self._cover(date, symbol, shares, fill_price)
-        bought_shares = self._buy(
+        bought_shares = self._long(
             date, symbol, covered.rem_shares, fill_price, limit_price, stops
         )
         if not covered.filled_shares and not bought_shares:
@@ -724,7 +728,7 @@ def _exit_short(
             mfe=mfe,
         )
 
-    def _buy(
+    def _long(
         self,
         date: np.datetime64,
         symbol: str,
@@ -733,6 +737,8 @@ def _buy(
         limit_price: Optional[Decimal],
         stops: Optional[Iterable[Stop]],
     ) -> Decimal:
+        if self._position_mode == PositionMode.SHORT_ONLY:
+            return Decimal()
         clamped_shares = self._clamp_shares(fill_price, shares)
         if clamped_shares < shares:
             self._logger.debug_buy_shares_exceed_cash(
@@ -925,6 +931,8 @@ def _short(
             and len(self.short_positions) == self._max_short_positions
         ):
             return Decimal()
+        if self._position_mode == PositionMode.LONG_ONLY:
+            return Decimal()
         if symbol not in self.short_positions:
             self.symbols.add(symbol)
             pos = Position(symbol=symbol, shares=shares, type="short")

src/pybroker/strategy.py

@@ -1251,6 +1251,7 @@ def walkforward(
                     self._config.fee_amount,
                     self._config.subtract_fees,
                     self._fractional_shares_enabled(),
+                    self._config.position_mode,
                     self._config.max_long_positions,
                     self._config.max_short_positions,
                     self._config.return_stops,

tests/test_portfolio.py

@@ -11,7 +11,7 @@
 import pytest
 from collections import deque
 from decimal import Decimal
-from pybroker.common import FeeMode, PriceType, StopType
+from pybroker.common import FeeMode, PositionMode, PriceType, StopType
 from pybroker.portfolio import Portfolio, Stop
 from pybroker.scope import ColumnScope, PriceScope
 
@@ -3536,3 +3536,25 @@ def test_mae_mfe_when_long_position():
     assert len(portfolio.trades) == 1
     assert portfolio.trades[0].mae == low_price - fill_price
     assert portfolio.trades[0].mfe == high_price - fill_price
+
+
+def test_long_only_mode():
+    cash = 100_000
+    portfolio = Portfolio(cash, position_mode=PositionMode.LONG_ONLY)
+    portfolio.buy(DATE_1, SYMBOL_1, 100, FILL_PRICE_1)
+    portfolio.sell(DATE_2, SYMBOL_1, 200, FILL_PRICE_1)
+    assert not portfolio.long_positions
+    assert not portfolio.short_positions
+    assert len(portfolio.trades) == 1
+    assert portfolio.trades[0].shares == 100
+
+
+def test_short_only_mode():
+    cash = 100_000
+    portfolio = Portfolio(cash, position_mode=PositionMode.SHORT_ONLY)
+    portfolio.sell(DATE_1, SYMBOL_1, 100, FILL_PRICE_1)
+    portfolio.buy(DATE_2, SYMBOL_1, 200, FILL_PRICE_1)
+    assert not portfolio.long_positions
+    assert not portfolio.short_positions
+    assert len(portfolio.trades) == 1
+    assert portfolio.trades[0].shares == 100