|
| 1 | +"""Tests for the dataframe backend dispatch. |
| 2 | +
|
| 3 | +Covers ``YfConfig.dataframe.backend`` validation and parity of every |
| 4 | +wrapped boundary in the codebase — using mocked-data unit tests so the |
| 5 | +suite stays network-independent. |
| 6 | +""" |
| 7 | +from __future__ import annotations |
| 8 | + |
| 9 | +import unittest |
| 10 | +from unittest.mock import MagicMock |
| 11 | + |
| 12 | +import pandas as pd |
| 13 | + |
| 14 | +from yfinance._backend import df_to_backend, series_to_backend |
| 15 | +from yfinance.config import YfConfig |
| 16 | +from yfinance.lookup import Lookup |
| 17 | + |
| 18 | + |
| 19 | +_LOOKUP_RESPONSE = { |
| 20 | + "finance": { |
| 21 | + "result": [ |
| 22 | + { |
| 23 | + "documents": [ |
| 24 | + {"symbol": "AAPL", "longName": "Apple Inc.", "quoteType": "EQUITY"}, |
| 25 | + {"symbol": "MSFT", "longName": "Microsoft", "quoteType": "EQUITY"}, |
| 26 | + ] |
| 27 | + } |
| 28 | + ] |
| 29 | + } |
| 30 | +} |
| 31 | + |
| 32 | + |
| 33 | +def _polars_df_type(): |
| 34 | + import polars as pl |
| 35 | + return pl.DataFrame |
| 36 | + |
| 37 | + |
| 38 | +# --------------------------------------------------------------------------- |
| 39 | +# YfConfig.dataframe.backend validation |
| 40 | +# --------------------------------------------------------------------------- |
| 41 | +class TestDataframeBackendConfig(unittest.TestCase): |
| 42 | + |
| 43 | + def tearDown(self): |
| 44 | + YfConfig.dataframe.backend = "pandas" |
| 45 | + |
| 46 | + def test_default_backend_is_pandas(self): |
| 47 | + self.assertEqual(YfConfig.dataframe.backend, "pandas") |
| 48 | + |
| 49 | + def test_polars_backend_accepted(self): |
| 50 | + YfConfig.dataframe.backend = "polars" |
| 51 | + self.assertEqual(YfConfig.dataframe.backend, "polars") |
| 52 | + |
| 53 | + def test_pandas_backend_accepted(self): |
| 54 | + YfConfig.dataframe.backend = "polars" |
| 55 | + YfConfig.dataframe.backend = "pandas" |
| 56 | + self.assertEqual(YfConfig.dataframe.backend, "pandas") |
| 57 | + |
| 58 | + def test_unknown_backend_raises(self): |
| 59 | + with self.assertRaises(ValueError): |
| 60 | + YfConfig.dataframe.backend = "modin" |
| 61 | + |
| 62 | + def test_unknown_backend_does_not_corrupt_state(self): |
| 63 | + try: |
| 64 | + YfConfig.dataframe.backend = "modin" |
| 65 | + except ValueError: |
| 66 | + pass |
| 67 | + self.assertEqual(YfConfig.dataframe.backend, "pandas") |
| 68 | + |
| 69 | + |
| 70 | +# --------------------------------------------------------------------------- |
| 71 | +# Helpers in _backend.py |
| 72 | +# --------------------------------------------------------------------------- |
| 73 | +class TestDfToBackend(unittest.TestCase): |
| 74 | + |
| 75 | + def tearDown(self): |
| 76 | + YfConfig.dataframe.backend = "pandas" |
| 77 | + |
| 78 | + def test_pandas_is_passthrough(self): |
| 79 | + df = pd.DataFrame({"x": [1, 2]}) |
| 80 | + out = df_to_backend(df) |
| 81 | + self.assertIs(out, df) |
| 82 | + |
| 83 | + def test_polars_returns_polars_frame(self): |
| 84 | + YfConfig.dataframe.backend = "polars" |
| 85 | + df = pd.DataFrame({"x": [1, 2]}) |
| 86 | + out = df_to_backend(df) |
| 87 | + self.assertIsInstance(out, _polars_df_type()) |
| 88 | + |
| 89 | + def test_polars_promotes_named_index_to_column(self): |
| 90 | + YfConfig.dataframe.backend = "polars" |
| 91 | + df = pd.DataFrame({"v": [1, 2]}, index=pd.Index(["a", "b"], name="key")) |
| 92 | + out = df_to_backend(df) |
| 93 | + self.assertIn("key", out.columns) |
| 94 | + self.assertEqual(out["key"].to_list(), ["a", "b"]) |
| 95 | + |
| 96 | + def test_polars_override_index_name(self): |
| 97 | + YfConfig.dataframe.backend = "polars" |
| 98 | + df = pd.DataFrame({"v": [1, 2]}, index=pd.Index(["a", "b"])) |
| 99 | + out = df_to_backend(df, index_as_column="Date") |
| 100 | + self.assertIn("Date", out.columns) |
| 101 | + |
| 102 | + def test_polars_rangeindex_unnamed_is_dropped(self): |
| 103 | + YfConfig.dataframe.backend = "polars" |
| 104 | + df = pd.DataFrame({"v": [1, 2]}) # default RangeIndex, no name |
| 105 | + out = df_to_backend(df) |
| 106 | + self.assertNotIn("index", out.columns) |
| 107 | + self.assertEqual(out["v"].to_list(), [1, 2]) |
| 108 | + |
| 109 | + def test_polars_empty_frame(self): |
| 110 | + YfConfig.dataframe.backend = "polars" |
| 111 | + out = df_to_backend(pd.DataFrame()) |
| 112 | + self.assertIsInstance(out, _polars_df_type()) |
| 113 | + self.assertEqual(out.shape, (0, 0)) |
| 114 | + |
| 115 | + |
| 116 | +class TestSeriesToBackend(unittest.TestCase): |
| 117 | + |
| 118 | + def tearDown(self): |
| 119 | + YfConfig.dataframe.backend = "pandas" |
| 120 | + |
| 121 | + def test_pandas_is_passthrough(self): |
| 122 | + s = pd.Series([1, 2], name="x") |
| 123 | + out = series_to_backend(s) |
| 124 | + self.assertIs(out, s) |
| 125 | + |
| 126 | + def test_polars_no_index_two_columns(self): |
| 127 | + YfConfig.dataframe.backend = "polars" |
| 128 | + s = pd.Series([1.0, 2.0], index=pd.Index(["a", "b"], name="Date"), name="Dividends") |
| 129 | + out = series_to_backend(s, index_as_column="Date", value_name="Dividends") |
| 130 | + self.assertIsInstance(out, _polars_df_type()) |
| 131 | + self.assertEqual(out.columns, ["Date", "Dividends"]) |
| 132 | + self.assertEqual(out["Dividends"].to_list(), [1.0, 2.0]) |
| 133 | + |
| 134 | + |
| 135 | +# --------------------------------------------------------------------------- |
| 136 | +# lookup.py |
| 137 | +# --------------------------------------------------------------------------- |
| 138 | +class TestLookupBackendParity(unittest.TestCase): |
| 139 | + |
| 140 | + def tearDown(self): |
| 141 | + YfConfig.dataframe.backend = "pandas" |
| 142 | + |
| 143 | + def test_pandas_output_uses_symbol_index(self): |
| 144 | + YfConfig.dataframe.backend = "pandas" |
| 145 | + df = Lookup._parse_response(_LOOKUP_RESPONSE) |
| 146 | + self.assertIsInstance(df, pd.DataFrame) |
| 147 | + self.assertEqual(df.index.name, "symbol") |
| 148 | + self.assertEqual(list(df.index), ["AAPL", "MSFT"]) |
| 149 | + |
| 150 | + def test_polars_output_keeps_symbol_column(self): |
| 151 | + YfConfig.dataframe.backend = "polars" |
| 152 | + df = Lookup._parse_response(_LOOKUP_RESPONSE) |
| 153 | + self.assertIsInstance(df, _polars_df_type()) |
| 154 | + self.assertIn("symbol", df.columns) |
| 155 | + self.assertEqual(df["symbol"].to_list(), ["AAPL", "MSFT"]) |
| 156 | + |
| 157 | + def test_empty_result_returns_empty_frame(self): |
| 158 | + for backend in ("pandas", "polars"): |
| 159 | + YfConfig.dataframe.backend = backend |
| 160 | + df = Lookup._parse_response({"finance": {"result": []}}) |
| 161 | + self.assertEqual(len(df), 0, f"{backend}: expected empty frame") |
| 162 | + |
| 163 | + |
| 164 | +# --------------------------------------------------------------------------- |
| 165 | +# base.py wrappers — exercise the get_* methods directly with fake scrapers. |
| 166 | +# --------------------------------------------------------------------------- |
| 167 | +def _fake_ticker(scraper_attrs: dict) -> MagicMock: |
| 168 | + """Build a TickerBase-like mock pre-populated with scraper attributes |
| 169 | + so we can call get_* methods without network.""" |
| 170 | + from yfinance.base import TickerBase |
| 171 | + t = TickerBase.__new__(TickerBase) |
| 172 | + for k, v in scraper_attrs.items(): |
| 173 | + setattr(t, k, v) |
| 174 | + return t |
| 175 | + |
| 176 | + |
| 177 | +class TestBaseGetterWrappers(unittest.TestCase): |
| 178 | + |
| 179 | + def tearDown(self): |
| 180 | + YfConfig.dataframe.backend = "pandas" |
| 181 | + |
| 182 | + def _ticker_with_quote(self, **kwargs): |
| 183 | + quote = MagicMock() |
| 184 | + for k, v in kwargs.items(): |
| 185 | + setattr(quote, k, v) |
| 186 | + return _fake_ticker({"_quote": quote}) |
| 187 | + |
| 188 | + def _ticker_with_holders(self, **kwargs): |
| 189 | + holders = MagicMock() |
| 190 | + for k, v in kwargs.items(): |
| 191 | + setattr(holders, k, v) |
| 192 | + return _fake_ticker({"_holders": holders}) |
| 193 | + |
| 194 | + def _ticker_with_analysis(self, **kwargs): |
| 195 | + analysis = MagicMock() |
| 196 | + for k, v in kwargs.items(): |
| 197 | + setattr(analysis, k, v) |
| 198 | + return _fake_ticker({"_analysis": analysis}) |
| 199 | + |
| 200 | + def _df(self): |
| 201 | + return pd.DataFrame({"a": [1, 2], "b": [3, 4]}) |
| 202 | + |
| 203 | + def test_get_recommendations(self): |
| 204 | + t = self._ticker_with_quote(recommendations=self._df()) |
| 205 | + YfConfig.dataframe.backend = "pandas" |
| 206 | + self.assertIsInstance(t.get_recommendations(), pd.DataFrame) |
| 207 | + YfConfig.dataframe.backend = "polars" |
| 208 | + self.assertIsInstance(t.get_recommendations(), _polars_df_type()) |
| 209 | + |
| 210 | + def test_get_upgrades_downgrades(self): |
| 211 | + t = self._ticker_with_quote(upgrades_downgrades=self._df()) |
| 212 | + YfConfig.dataframe.backend = "polars" |
| 213 | + self.assertIsInstance(t.get_upgrades_downgrades(), _polars_df_type()) |
| 214 | + |
| 215 | + def test_get_sustainability(self): |
| 216 | + t = self._ticker_with_quote(sustainability=self._df()) |
| 217 | + YfConfig.dataframe.backend = "polars" |
| 218 | + self.assertIsInstance(t.get_sustainability(), _polars_df_type()) |
| 219 | + |
| 220 | + def test_get_valuation_measures(self): |
| 221 | + t = self._ticker_with_quote(valuation_measures=self._df()) |
| 222 | + YfConfig.dataframe.backend = "polars" |
| 223 | + self.assertIsInstance(t.get_valuation_measures(), _polars_df_type()) |
| 224 | + |
| 225 | + def test_get_major_holders(self): |
| 226 | + t = self._ticker_with_holders(major=self._df()) |
| 227 | + YfConfig.dataframe.backend = "polars" |
| 228 | + self.assertIsInstance(t.get_major_holders(), _polars_df_type()) |
| 229 | + |
| 230 | + def test_get_institutional_holders(self): |
| 231 | + t = self._ticker_with_holders(institutional=self._df()) |
| 232 | + YfConfig.dataframe.backend = "polars" |
| 233 | + self.assertIsInstance(t.get_institutional_holders(), _polars_df_type()) |
| 234 | + |
| 235 | + def test_get_mutualfund_holders(self): |
| 236 | + t = self._ticker_with_holders(mutualfund=self._df()) |
| 237 | + YfConfig.dataframe.backend = "polars" |
| 238 | + self.assertIsInstance(t.get_mutualfund_holders(), _polars_df_type()) |
| 239 | + |
| 240 | + def test_get_insider_purchases(self): |
| 241 | + t = self._ticker_with_holders(insider_purchases=self._df()) |
| 242 | + YfConfig.dataframe.backend = "polars" |
| 243 | + self.assertIsInstance(t.get_insider_purchases(), _polars_df_type()) |
| 244 | + |
| 245 | + def test_get_insider_transactions(self): |
| 246 | + t = self._ticker_with_holders(insider_transactions=self._df()) |
| 247 | + YfConfig.dataframe.backend = "polars" |
| 248 | + self.assertIsInstance(t.get_insider_transactions(), _polars_df_type()) |
| 249 | + |
| 250 | + def test_get_insider_roster_holders(self): |
| 251 | + t = self._ticker_with_holders(insider_roster=self._df()) |
| 252 | + YfConfig.dataframe.backend = "polars" |
| 253 | + self.assertIsInstance(t.get_insider_roster_holders(), _polars_df_type()) |
| 254 | + |
| 255 | + def test_get_earnings_estimate(self): |
| 256 | + t = self._ticker_with_analysis(earnings_estimate=self._df()) |
| 257 | + YfConfig.dataframe.backend = "polars" |
| 258 | + self.assertIsInstance(t.get_earnings_estimate(), _polars_df_type()) |
| 259 | + |
| 260 | + def test_get_revenue_estimate(self): |
| 261 | + t = self._ticker_with_analysis(revenue_estimate=self._df()) |
| 262 | + YfConfig.dataframe.backend = "polars" |
| 263 | + self.assertIsInstance(t.get_revenue_estimate(), _polars_df_type()) |
| 264 | + |
| 265 | + def test_get_earnings_history(self): |
| 266 | + t = self._ticker_with_analysis(earnings_history=self._df()) |
| 267 | + YfConfig.dataframe.backend = "polars" |
| 268 | + self.assertIsInstance(t.get_earnings_history(), _polars_df_type()) |
| 269 | + |
| 270 | + def test_get_eps_trend(self): |
| 271 | + t = self._ticker_with_analysis(eps_trend=self._df()) |
| 272 | + YfConfig.dataframe.backend = "polars" |
| 273 | + self.assertIsInstance(t.get_eps_trend(), _polars_df_type()) |
| 274 | + |
| 275 | + def test_get_eps_revisions(self): |
| 276 | + t = self._ticker_with_analysis(eps_revisions=self._df()) |
| 277 | + YfConfig.dataframe.backend = "polars" |
| 278 | + self.assertIsInstance(t.get_eps_revisions(), _polars_df_type()) |
| 279 | + |
| 280 | + def test_get_growth_estimates(self): |
| 281 | + t = self._ticker_with_analysis(growth_estimates=self._df()) |
| 282 | + YfConfig.dataframe.backend = "polars" |
| 283 | + self.assertIsInstance(t.get_growth_estimates(), _polars_df_type()) |
| 284 | + |
| 285 | + def test_as_dict_short_circuit_returns_dict(self): |
| 286 | + """`as_dict=True` must always return a plain dict regardless of backend.""" |
| 287 | + t = self._ticker_with_quote(recommendations=self._df()) |
| 288 | + YfConfig.dataframe.backend = "polars" |
| 289 | + self.assertIsInstance(t.get_recommendations(as_dict=True), dict) |
| 290 | + |
| 291 | + |
| 292 | +# --------------------------------------------------------------------------- |
| 293 | +# Cache invariant: changing backend mid-session must propagate to next access. |
| 294 | +# --------------------------------------------------------------------------- |
| 295 | +class TestBackendSwitchInvariant(unittest.TestCase): |
| 296 | + |
| 297 | + def tearDown(self): |
| 298 | + YfConfig.dataframe.backend = "pandas" |
| 299 | + |
| 300 | + def test_lookup_repeated_parse_honors_current_backend(self): |
| 301 | + YfConfig.dataframe.backend = "pandas" |
| 302 | + pd_df = Lookup._parse_response(_LOOKUP_RESPONSE) |
| 303 | + self.assertIsInstance(pd_df, pd.DataFrame) |
| 304 | + YfConfig.dataframe.backend = "polars" |
| 305 | + pl_df = Lookup._parse_response(_LOOKUP_RESPONSE) |
| 306 | + self.assertIsInstance(pl_df, _polars_df_type()) |
| 307 | + YfConfig.dataframe.backend = "pandas" |
| 308 | + pd_df_again = Lookup._parse_response(_LOOKUP_RESPONSE) |
| 309 | + self.assertIsInstance(pd_df_again, pd.DataFrame) |
| 310 | + |
| 311 | + |
| 312 | +# --------------------------------------------------------------------------- |
| 313 | +# calendars.py — _cleanup_df + _to_backend |
| 314 | +# --------------------------------------------------------------------------- |
| 315 | +class TestCalendarsBackendParity(unittest.TestCase): |
| 316 | + |
| 317 | + def tearDown(self): |
| 318 | + YfConfig.dataframe.backend = "pandas" |
| 319 | + |
| 320 | + def _make_calendars(self, calendar_type="sp_earnings"): |
| 321 | + from yfinance.calendars import Calendars |
| 322 | + # Synthetic raw frame matching what _create_df builds. |
| 323 | + if calendar_type == "sp_earnings": |
| 324 | + df = pd.DataFrame({ |
| 325 | + "Symbol": ["AAPL", "MSFT"], |
| 326 | + "Company Name": ["Apple", "Microsoft"], |
| 327 | + "Market Cap (Intraday)": [1.0, 2.0], |
| 328 | + "Event Name": ["Q1", "Q2"], |
| 329 | + "Event Start Date": ["2025-01-01", "2025-02-01"], |
| 330 | + "Timing": ["BMO", "AMC"], |
| 331 | + "EPS Estimate": [1.0, 2.0], |
| 332 | + "Reported EPS": [1.1, 2.1], |
| 333 | + "Surprise (%)": [0.1, 0.05], |
| 334 | + }) |
| 335 | + c = Calendars() |
| 336 | + c.calendars[calendar_type] = df |
| 337 | + return c |
| 338 | + |
| 339 | + def test_to_backend_pandas(self): |
| 340 | + c = self._make_calendars() |
| 341 | + df = c._to_backend("sp_earnings") |
| 342 | + self.assertIsInstance(df, pd.DataFrame) |
| 343 | + self.assertEqual(df.index.name, "Symbol") |
| 344 | + |
| 345 | + def test_to_backend_polars_keeps_index_as_column(self): |
| 346 | + YfConfig.dataframe.backend = "polars" |
| 347 | + c = self._make_calendars() |
| 348 | + df = c._to_backend("sp_earnings") |
| 349 | + self.assertIsInstance(df, _polars_df_type()) |
| 350 | + self.assertIn("Symbol", df.columns) |
| 351 | + |
| 352 | + |
| 353 | +# --------------------------------------------------------------------------- |
| 354 | +# domain/sector.py |
| 355 | +# --------------------------------------------------------------------------- |
| 356 | +class TestSectorBackendParity(unittest.TestCase): |
| 357 | + |
| 358 | + def tearDown(self): |
| 359 | + YfConfig.dataframe.backend = "pandas" |
| 360 | + |
| 361 | + def test_industries_property_switches_with_backend(self): |
| 362 | + from yfinance.domain.sector import Sector |
| 363 | + s = Sector.__new__(Sector) |
| 364 | + s._industries = pd.DataFrame( |
| 365 | + {"name": ["A", "B"], "symbol": ["X", "Y"], "market weight": [0.5, 0.5]}, |
| 366 | + index=pd.Index(["a", "b"], name="key"), |
| 367 | + ) |
| 368 | + s._ensure_fetched = lambda *_a, **_kw: None # type: ignore[assignment] |
| 369 | + |
| 370 | + YfConfig.dataframe.backend = "pandas" |
| 371 | + self.assertIsInstance(s.industries, pd.DataFrame) |
| 372 | + YfConfig.dataframe.backend = "polars" |
| 373 | + self.assertIsInstance(s.industries, _polars_df_type()) |
| 374 | + self.assertIn("key", s.industries.columns) |
| 375 | + |
| 376 | + |
| 377 | +if __name__ == "__main__": |
| 378 | + unittest.main() |
0 commit comments