More detailed validation of the downloaded data in download_data.py. For this reason, a new validation in backtest_signals.py is not required. The cumbersome multi-index is also removed early on.

niklasniedermeier · niklasniedermeier · commit f4a622c39742 · 2025-02-22T17:52:05.000+01:00
diff --git a/src/backtest_bay/backtest/backtest_signals.py b/src/backtest_bay/backtest/backtest_signals.py
@@ -28,6 +28,8 @@ def backtest_signals(data, signals, initial_cash, tac, trade_pct, price_col="Clo
             - 'cash': Cash.
             - 'assets': Portfolio value (cash + holdings).
     """
+    # Note that the inputs 'data' and 'signals' are already validated in
+    # 'download_data.py' and 'generate_signals.py'
     _validate_backtest_signals_input(data, initial_cash, tac, trade_pct, price_col)
 
     prices = data[price_col].squeeze()
@@ -159,36 +161,10 @@ def _update_portfolio(cash, shares, price):
 
 def _validate_backtest_signals_input(data, initial_cash, tac, trade_pct, price_col):
     """Validates input for backtesting signals."""
-    # Since the variable 'signals' is generated using 'generate_signals', we already
-    # validated the input in the corresponding test script of 'generate_signals'.
-    _validate_data(data, price_col)
     _validate_initial_cash(initial_cash)
     _validate_tac(tac)
     _validate_trade_pct(trade_pct)
-
-
-def _validate_data(data, price_col):
-    """Validate the input data for backtesting.
-
-    Args:
-        data (pd.DataFrame): DataFrame containing stock data.
-        price_col (str): Column name for the stock price.
-
-    Raises:
-        TypeError: If data is not a pandas DataFrame.
-        ValueError: If the price column is missing or contains non-numeric values.
-    """
-    if not isinstance(data, pd.DataFrame):
-        error_msg = f"data must be a pandas DataFrame, got {type(data).__name__}."
-        raise TypeError(error_msg)
-
-    if price_col not in data.columns:
-        error_msg = f"data must contain a '{price_col}' column."
-        raise ValueError(error_msg)
-
-    if not pd.api.types.is_numeric_dtype(data[price_col].squeeze()):
-        error_msg = f"The '{price_col}' column must contain numeric values."
-        raise ValueError(error_msg)
+    _validate_price_col(data, price_col)
 
 
 def _validate_initial_cash(initial_cash):
@@ -245,3 +221,18 @@ def _validate_trade_pct(trade_pct):
     if not (0 < trade_pct <= 1):
         error_msg = "trade_pct must be between 0 and 1. Zero is not possible."
         raise ValueError(error_msg)
+
+
+def _validate_price_col(data, price_col):
+    """Validate the input price_col for backtesting.
+
+    Args:
+        data (pd.DataFrame): DataFrame containing stock data.
+        price_col (str): Column name for the stock price.
+
+    Raises:
+        ValueError: If the price column is missing.
+    """
+    if price_col not in data.columns:
+        error_msg = f"data must contain a '{price_col}' column."
+        raise ValueError(error_msg)
diff --git a/src/backtest_bay/backtest/generate_signals.py b/src/backtest_bay/backtest/generate_signals.py
@@ -29,7 +29,7 @@ def generate_signals(data, method, **kwargs):
         pd.Series: Trading signals (2: buy, 1: sell, 0: do nothing).
     """
     _validate_input_method(method)
-    closing_prices = data["Close"].squeeze()
+    closing_prices = data["Close"]
 
     if method == "bollinger":
         signal = _bollinger_signals(prices=closing_prices, **kwargs)
diff --git a/src/backtest_bay/backtest/task_backtest.py b/src/backtest_bay/backtest/task_backtest.py
@@ -55,5 +55,4 @@ def _merge_stock_data_with_portfolio(data, portfolio):
     Returns:
         pd.DataFrame: Merged DataFrame.
     """
-    data.columns = data.columns.droplevel(1)
     return data.merge(portfolio, how="left", left_index=True, right_index=True)
diff --git a/src/backtest_bay/data/download_data.py b/src/backtest_bay/data/download_data.py
@@ -1,16 +1,34 @@
 from datetime import datetime
 
+import pandas as pd
 import yfinance as yf
 
 
 def download_data(symbol, interval, start_date, end_date):
+    """Download historical stock data and validate it.
+
+    This function downloads historical stock data for a given symbol, date range,
+    and interval using the yfinance library. The input parameters and the downloaded
+    data are validated to ensure they meet the required criteria.
+
+    Args:
+        symbol (str): Stock symbol to download data for (e.g., 'AAPL' for Apple).
+        interval (str): Data interval (e.g., '1d' for daily, '1h' for hourly).
+        start_date (str): Start date for the data in 'YYYY-MM-DD' format.
+        end_date (str): End date for the data in 'YYYY-MM-DD' format.
+
+    Returns:
+        pd.DataFrame: DataFrame containing the downloaded stock data.
+    """
     _validate_input(symbol, interval, start_date, end_date)
     data = yf.download(symbol, start=start_date, end=end_date, interval=interval)
-    _validate_output(data, symbol, start_date, end_date, interval)
+    _validate_data(data, symbol, start_date, end_date, interval)
+    data.columns = _remove_multiindex_from_cols(data.columns)
     return data
 
 
 def _validate_input(symbol, interval, start_date, end_date):
+    """Validate symbol, interval, and date inputs."""
     _validate_symbol(symbol)
     _validate_interval(interval)
     _validate_date_format(start_date)
@@ -19,13 +37,15 @@ def _validate_input(symbol, interval, start_date, end_date):
 
 
 def _validate_symbol(symbol):
+    """Check if symbol is a non-empty string."""
     is_symbol_string = isinstance(symbol, str)
     if not is_symbol_string:
         error_msg = "Symbol must be a non-empty string."
         raise TypeError(error_msg)
 
 
 def _validate_interval(interval):
+    """Validate if interval is within the allowed set."""
     valid_intervals = {
         "1m",
         "2m",
@@ -47,6 +67,7 @@ def _validate_interval(interval):
 
 
 def _validate_date_format(date_str):
+    """Check if date string is in 'YYYY-MM-DD' format."""
     if not isinstance(date_str, str):
         error_msg = "Date must be a string in 'YYYY-MM-DD' format."
         raise TypeError(error_msg)
@@ -58,15 +79,111 @@ def _validate_date_format(date_str):
 
 
 def _validate_date_range(start_date, end_date):
+    """Ensure start date is before end date."""
     if start_date > end_date:
         error_msg = "Start date must be before end date."
         raise ValueError(error_msg)
 
 
-def _validate_output(data, symbol, start_date, end_date, interval):
+def _validate_data(data, symbol, start_date, end_date, interval):
+    """Validate the downloaded data.
+
+    Args:
+        data (pd.DataFrame): DataFrame containing stock data.
+        symbol (str): Stock symbol for the data.
+        start_date (str): Start date for the data.
+        end_date (str): End date for the data.
+        interval (str): Interval for the data.
+
+    Raises:
+        TypeError: If data is not a pandas DataFrame or index is not a DatetimeIndex.
+        ValueError: If required columns are missing or contain non-numeric values.
+        ValueError: If the DataFrame is empty.
+    """
+    _validate_data_type_dataframe(data)
+    _validate_data_empty(data, symbol, start_date, end_date, interval)
+    _validate_data_index_datetime(data.index)
+    _validate_data_multiindex(data.columns)
+    _validate_data_numeric(data)
+
+
+def _validate_data_type_dataframe(data):
+    """Check if the input is a pandas DataFrame."""
+    if not isinstance(data, pd.DataFrame):
+        error_msg = f"data must be a pandas DataFrame, got {type(data).__name__}."
+        raise TypeError(error_msg)
+
+
+def _validate_data_empty(data, symbol, start_date, end_date, interval):
+    """Check if the DataFrame is empty."""
     if data.empty:
         error_msg = (
             f"No data found for {symbol} between {start_date} and {end_date} "
             f"with interval '{interval}'."
         )
         raise ValueError(error_msg)
+
+
+def _validate_data_index_datetime(index):
+    """Check if the index is of type DatetimeIndex."""
+    if not isinstance(index, pd.DatetimeIndex):
+        error_msg = (
+            f"data index must be a pandas DatetimeIndex, got {type(index).__name__}."
+        )
+        raise TypeError(error_msg)
+
+
+def _validate_data_multiindex(columns):
+    """Check if the columns have the required MultiIndex.
+
+    Args:
+        columns (pd.MultiIndex): MultiIndex of DataFrame columns.
+
+    Raises:
+        ValueError: If the MultiIndex is not present, does not have exactly two levels,
+                    or if required columns are missing from level 0.
+    """
+    required_cols = {"Close", "Open", "High", "Low"}
+
+    if not isinstance(columns, pd.MultiIndex):
+        error_msg = "DataFrame must have a MultiIndex for columns."
+        raise TypeError(error_msg)
+
+    yfinance_index_levels = 2
+    if columns.nlevels != yfinance_index_levels:
+        error_msg = (
+            f"MultiIndex must have exactly 2 levels, got {columns.nlevels} levels."
+        )
+        raise ValueError(error_msg)
+
+    level_0_values = set(columns.get_level_values(0))
+    missing_cols = required_cols - level_0_values
+    if missing_cols:
+        error_msg = (
+            "Level 0 of MultiIndex must contain the following columns: "
+            f"{', '.join(missing_cols)}."
+        )
+        raise ValueError(error_msg)
+
+
+def _validate_data_numeric(data):
+    """Check if the required columns contain numeric values."""
+    required_cols = ["Close", "Open", "High", "Low"]
+    non_numeric_cols = [
+        col
+        for col in required_cols
+        if not pd.api.types.is_numeric_dtype(data[col].squeeze())
+    ]
+
+    if non_numeric_cols:
+        error_msg = (
+            "The following columns must contain numeric values: "
+            f"{', '.join(non_numeric_cols)}."
+        )
+        raise ValueError(error_msg)
+
+
+def _remove_multiindex_from_cols(cols):
+    """Remove MultiIndex from columns but retain level 0 as column names."""
+    cols = cols.get_level_values(0)
+    return cols
diff --git a/tests/backtest/test_backtest_signals.py b/tests/backtest/test_backtest_signals.py
@@ -8,8 +8,8 @@
     _is_buy_trade_affordable,
     _is_sell_trade_affordable,
     _update_portfolio,
-    _validate_data,
     _validate_initial_cash,
+    _validate_price_col,
     _validate_tac,
     _validate_trade_pct,
     backtest_signals,
@@ -151,67 +151,6 @@ def test_update_portfolio_correct_calculation():
     assert holdings == expected_holdings
 
 
-# Tests for _validate_data
-@pytest.mark.parametrize(
-    ("data", "price_col"),
-    [
-        (pd.DataFrame({"Close": [100, 101, 102]}), "Close"),
-        (pd.DataFrame({"Open": [99, 100, 101], "Close": [100, 101, 102]}), "Close"),
-        (pd.DataFrame({"Close": [100.5, 101.2, 102.1, 103.8]}), "Close"),
-    ],
-)
-def test_validate_data_valid_input(data, price_col):
-    """Test valid data input for _validate_data."""
-    _validate_data(data, price_col)
-
-
-@pytest.mark.parametrize(
-    ("data", "price_col", "expected_error"),
-    [
-        ([100, 101], "Close", "data must be a pandas DataFrame, got list."),
-        ({"Close": [100, 101]}, "Close", "data must be a pandas DataFrame, got dict."),
-    ],
-)
-def test_validate_data_invalid_type(data, price_col, expected_error):
-    """Test invalid data types for _validate_data."""
-    with pytest.raises(TypeError, match=expected_error):
-        _validate_data(data, price_col)
-
-
-@pytest.mark.parametrize(
-    ("data", "price_col", "expected_error"),
-    [
-        (pd.DataFrame({"Open": [100]}), "Close", "data must contain a 'Close' column."),
-        (pd.DataFrame({"Close": [100]}), "Open", "data must contain a 'Open' column."),
-    ],
-)
-def test_validate_data_missing_price_column(data, price_col, expected_error):
-    """Test missing price column for _validate_data."""
-    with pytest.raises(ValueError, match=expected_error):
-        _validate_data(data, price_col)
-
-
-@pytest.mark.parametrize(
-    ("data", "price_col", "expected_error"),
-    [
-        (
-            pd.DataFrame({"Close": ["100", 100]}),
-            "Close",
-            "The 'Close' column must contain numeric values.",
-        ),
-        (
-            pd.DataFrame({"Close": ["a", "b", "c", "d"]}),
-            "Close",
-            "The 'Close' column must contain numeric values.",
-        ),
-    ],
-)
-def test_validate_data_non_numeric_price_column(data, price_col, expected_error):
-    """Test non-numeric price column for _validate_data."""
-    with pytest.raises(ValueError, match=expected_error):
-        _validate_data(data, price_col)
-
-
 # Tests for _validate_initial_cash
 @pytest.mark.parametrize("initial_cash", [1000, 1000.50, 0.01])
 def test_validate_initial_cash_valid_input(initial_cash):
@@ -272,3 +211,17 @@ def test_validate_trade_pct_invalid_input(trade_pct, expected_error):
     """Test invalid trade_pct values."""
     with pytest.raises((TypeError, ValueError), match=expected_error):
         _validate_trade_pct(trade_pct)
+
+
+# Tests for _validate_price_col
+@pytest.mark.parametrize(
+    ("data", "price_col"),
+    [
+        (pd.DataFrame({"Close": [100, 101, 102]}), "Close"),
+        (pd.DataFrame({"Open": [99, 100, 101], "Close": [100, 101, 102]}), "Close"),
+        (pd.DataFrame({"Close": [100.5, 101.2, 102.1, 103.8]}), "Close"),
+    ],
+)
+def test_validate_data_valid_input(data, price_col):
+    """Test valid data input for _validate_data."""
+    _validate_price_col(data, price_col)
diff --git a/tests/data/test_download_data.py b/tests/data/test_download_data.py
