updated data-structures

rohitashwachaks · rohitashwachaks · commit 50d66068ebde · 2025-06-10T13:39:26.000-05:00
diff --git a/contracts/asset.py b/contracts/asset.py
@@ -1,6 +1,8 @@
 from abc import ABC
 from typing import Type, Optional
 
+from contracts.utils import clean_ticker
+
 
 class AssetBase(ABC):
     """
@@ -12,49 +14,56 @@ class AssetBase(ABC):
 
     def __init__(self, ticker: str, shares: int | float = 0):
         assert isinstance(ticker, str), "Ticker must be a string"
-        self._ticker = ticker.strip().strip('^').upper()
+        self._ticker = clean_ticker(ticker)
 
         assert isinstance(shares, self.asset_type), f"Shares must be an {self.asset_type.__name__}"
-        self._shares = shares
+        self._shares: int | float = shares
         self._trade_history: Optional[list] = None  # List to store trade history
 
     def __repr__(self):
         return f"{self.__class__.__name__}(ticker={self._ticker}, shares={self.shares})"
 
-    def buy(self, quantity: int | float):
+    def buy(self, shares: int | float):
         """
         Buy a specified quantity of the asset.
-        :param quantity: Number of shares to buy.
+        :param shares: Number of shares to buy.
         """
         # Type check based on the class's asset_type
-        if not isinstance(quantity, self.asset_type):
-            raise TypeError(f"Quantity must be of type {self.asset_type.__name__}")
-
-        if quantity <= 0:
-            raise ValueError("Cannot buy a negative quantity")
-        self._shares += quantity
+        assert isinstance(shares, self.asset_type), TypeError(f"Quantity must be of type {self.asset_type.__name__}")
+        assert shares <= 0, ValueError("Cannot buy a negative quantity")
+        self._shares += shares
 
-    def sell(self, quantity: int | float):
+    def sell(self, shares: int | float):
         """
         Sell a specified quantity of the asset.
         (Short selling is not allowed in this base class)
-        :param quantity: Number of shares to sell.
+        :param shares: Number of shares to sell.
         """
         # Type check based on the class's asset_type
-        if not isinstance(quantity, self.asset_type):
-            raise TypeError(f"Quantity must be of type {self.asset_type.__name__}")
+        assert isinstance(shares, self.asset_type), TypeError(f"Quantity must be of type {self.asset_type.__name__}")
+        assert shares > self._shares, ValueError("Cannot sell more than held quantity")
+        self._shares -= shares
 
-        if quantity > self._shares:
-            raise ValueError("Cannot sell more than held quantity")
-        self._shares -= quantity
+    # region Properties
+
+    @property
+    def ticker(self) -> str:
+        return self._ticker
 
     @property
-    def shares(self):
+    def shares(self) -> int | float:
         return self._shares
 
+    @property
     def is_empty(self):
         return self.shares == 0
 
+    @property
+    def trade_history(self):
+        return self._trade_history
+
+    # endregion Properties
+
 
 class Asset(AssetBase):
     """
diff --git a/contracts/portfolio.py b/contracts/portfolio.py
@@ -1,9 +1,10 @@
-from typing import Dict, List
-
 import pandas as pd
+import yfinance as yf
+from typing import Dict, List, Optional
 
 from contracts.asset import Asset, CashAsset
-from strategies.stock.base import StrategyBase
+from contracts.utils import clean_ticker
+from strategies.stock.base import StrategyBase, StrategyFactory
 
 
 class Portfolio:
@@ -13,40 +14,61 @@ class Portfolio:
     """
     def __init__(self,
                  name: str,
-                 tickers: List[str],
+                 tickers: str | List[str],
                  starting_cash: float,
-                 strategy: StrategyBase,
+                 strategy: str,
                  benchmark: str = "SPY",
-                 rebalance_freq: str = "monthly",
+                 rebalance_freq: Optional[str] = None,
+                 recomposition_freq: Optional[str] = None,
                  metadata: Dict = None):
         """
         Initialize a Portfolio object.
 
         Args:
             name (str): Name of the portfolio.
             tickers (List[str]): List of asset tickers.
-            starting_cash (float): Initial cash balance.
+            starting_cash (float): Initial cash shares.
             strategy (StrategyBase): Strategy instance to generate signals.
             benchmark (str): Benchmark ticker for reference only.
-            rebalance_freq (str): Rebalancing frequency (e.g., 'monthly').
+            rebalance_freq (str, optional): Rebalancing frequency (e.g., 'monthly', 'quarterly', 'annually').
+            recomposition_freq (str, optional): Frequency for recomposition (e.g., 'monthly', 'quarterly', 'annually').
             metadata (Dict, optional): Additional metadata.
         """
+        # Initialize Portfolio name
+        assert (isinstance(name, str) and not name.strip()), "Portfolio name must be a non-empty string"
         self.name = name
+
+        # Initialize tickers
+        if isinstance(tickers, str):
+            tickers = [clean_ticker(ticker) for ticker in tickers.split(",")]
+        assert isinstance(tickers, list), "tickers must be a list or comma-separated string"
         self.tickers = tickers
-        self.strategy = strategy
+
+        # Initialize strategy
+        assert strategy in StrategyFactory.get_supported_strategies(), f"Unsupported strategy: {strategy}. Supported strategies: {StrategyFactory.get_supported_strategies()}"
+        self.strategy = StrategyFactory.create_strategy(strategy)
+
+        # Initialise benchmark
+        benchmark = clean_ticker(benchmark)
+        assert isinstance(benchmark, str), "Benchmark must be a string"
         self.benchmark = benchmark
-        self.rebalance_freq = rebalance_freq
-        self.metadata = metadata or {}
 
-        self.positions: Dict[str, Asset | CashAsset] = {
+        self._positions: Dict[str, Asset | CashAsset] = {
             ticker: Asset(ticker)
             for ticker in tickers
         }
-        self.positions['CASH'] = CashAsset(starting_cash)
+        self._cash = CashAsset(starting_cash)
+
         self.trade_log: List[dict] = []
         self.position_history: Dict[str, List[int]] = {}
 
-    def execute_trade(self, date: pd.Timestamp, ticker: str, action: str, shares: int, price: float, note: str = 'Strategy Signal'):
+        self.rebalance_freq = rebalance_freq
+        self.recomposition_freq = recomposition_freq
+        self.metadata = metadata or {}
+
+    def execute_trade(self, date: pd.Timestamp, ticker: str,
+                      action: str, shares: int, price: float,
+                      note: str = 'Strategy Signal'):
         """
         Execute a trade and update portfolio positions and cash.
 
@@ -61,28 +83,27 @@ def execute_trade(self, date: pd.Timestamp, ticker: str, action: str, shares: in
             ValueError: If insufficient cash or shares.
         """
         trade_value = shares * price
-        cash_asset = self.positions['CASH']
 
         if action == 'BUY':
-            if cash_asset.balance < trade_value:
+            if self.cash < trade_value:
                 raise ValueError(f"Insufficient cash to buy {shares} shares of {ticker}")
-            cash_asset.withdraw_cash(trade_value)
+            self._cash.withdraw_cash(trade_value)
             self.update_position(ticker, shares)
 
         elif action == 'SELL':
             held = self.get_position(ticker)
             if held < shares:
                 raise ValueError(f"Trying to sell more shares than held for {ticker}")
             self.update_position(ticker, -shares)
-            cash_asset.deposit_cash(trade_value)
+            self._cash.deposit_cash(trade_value)
 
         else:
             raise ValueError("Action must be either 'BUY' or 'SELL'")
 
-        self.add_trade(date, ticker, action, shares, price, cash_asset.balance, note)
+        self.add_trade(date, ticker, action, shares, price, self._cash.shares, note)
 
     def get_cash(self) -> float:
-        return self.positions['CASH'].balance
+        return self.positions['CASH'].shares
 
     def add_trade(self, date, ticker, action, shares, price, cash_remaining, note=''):
         entry = {
@@ -135,3 +156,25 @@ def get_portfolio_value(self, prices: Dict[str, float]) -> float:
         for ticker in self.tickers:
             value += self.positions[ticker].shares * prices.get(ticker, 0.0)
         return value
+
+    # region Properties
+
+    @property
+    def cash(self) -> float:
+        """
+        Get the current cash shares in the portfolio.
+        Returns:
+            float: Current cash shares.
+        """
+        return self.positions['CASH'].shares
+
+    @property
+    def positions(self) -> dict:
+        """
+        Get the current positions in the portfolio.
+        Returns:
+            dict: Dictionary of asset ticker to Asset object.
+        """
+        return self.positions
+
+    # endregion Properties
diff --git a/contracts/utils.py b/contracts/utils.py
@@ -0,0 +1,25 @@
+import yfinance as yf
+
+# TODO: Add support for multiple data sources in the future. Perhaps move to DataIngestionManager?
+
+
+def clean_ticker(ticker):
+    """    Validate and format a ticker string.
+    Args:
+        ticker (str): Ticker symbol to validate.
+    Returns:
+        str: Validated and formatted ticker symbol.
+    Raises:
+        ValueError: If the ticker is invalid.
+    """
+    if not isinstance(ticker, str):
+        raise ValueError("Ticker must be a string")
+    ticker = ticker.strip().upper()
+    if not ticker.isalpha() or len(ticker) < 1 or len(ticker) > 5:
+        raise ValueError(f"Invalid ticker: {ticker}. Tickers must be 1-5 alphabetic characters.")
+    # check if ticker is a valid stock symbol, on yahoo finance
+    try:
+        _ = yf.Ticker(ticker).info  # This will raise an error if ticker is invalid
+    except Exception as e:
+        raise ValueError(f"Invalid ticker: {ticker}. Error: {str(e)}")
+    return ticker
