Skip to content

Commit 5353ffa

Browse files
committed
Add opt-in polars output via pandas pipeline + boundary adapter
Closes #1868 (alternative to #2782 / #2808). Internal yfinance pipelines stay 100% pandas. Polars support is a thin output adapter applied at every public-API boundary: yf.config.dataframe.backend = 'polars' yf.Ticker('MSFT').history(period='1mo') # -> polars.DataFrame Direct answer to the maintenance concern raised on #2808: a parallel polars implementation doubles surface area; an output adapter does not. There is *one* code path for transformations. It isn't. yfinance frames are small enough that pandas vs polars makes no practical difference -- exactly the point #2808's review flagged. The value is *convenience*: users who downstream-process with polars get polars frames straight from yfinance instead of threading `pl.from_pandas(...)` after every call. - setup.py: optional dep group `yfinance[polars]` (polars only; no narwhals because with no transformation layer there is nothing to make backend-agnostic). - config.py: `YfConfig.dataframe.backend` defaults to 'pandas', raises ValueError on unknown values, leaves prior state intact on validation failure. - _backend.py: helpers (`current_backend`, `df_to_backend`, `series_to_backend`, `empty_df`) and public type aliases `DataFrameLike` / `SeriesLike`. All pure pass-through when backend == 'pandas'. Wrapped at the public boundary: - lookup.py, calendars.py, domain/{sector,industry,domain}.py - base.py: get_dividends/splits/capital_gains/actions/shares/ shares_full/earnings_dates/earnings, financial statements, recommendations, upgrades_downgrades, holders x6, valuation_measures, sustainability, analysis getters x6. - ticker.py: option_chain (calls / puts) and all @Property return annotations. - scrapers/funds.py: fund_operations, top_holdings, equity_holdings, bond_holdings. - scrapers/history.py: 4 return points of `history()`. - multi.py: `download()` final return. Type annotations: public-facing return types changed from `pd.DataFrame` / `pd.Series` to `DataFrameLike` / `SeriesLike` (union of pandas + polars). Not a breaking change: pandas remains a valid subtype, and `polars` is only imported at type-check time (TYPE_CHECKING guard). - Default backend is 'pandas'. Existing users see no behavioural difference. - `df_to_backend(df)` is identity (`assertIs`) when backend == 'pandas'. Verified by unit test. - Caches keep frames in pandas; conversion happens at each read, so changing backend mid-session is honoured immediately by every subsequent call. `tests/test_dataframe_backend.py` -- 37 tests, no network: - Config validation (default / accepted / unknown / no-corruption). - `df_to_backend` semantics (passthrough, polars conversion, named-index promotion, override, RangeIndex drop, empty frame). - `series_to_backend` semantics. - Lookup parity on both backends + empty result. - All 17 wrapped base.py getters return the right backend type, plus `as_dict=True` always returns a plain dict. - Backend switch invariant: pandas -> polars -> pandas in sequence returns the matching type each time. - Calendars `_to_backend` on both backends. - Sector.industries property switches with backend. All 58 unit tests (37 backend + 21 utils) pass. Ruff clean for the changed files. - `download()` long-form polars shape: kept as MultiIndex pandas converted via reset_index. Native polars long-form is a separate opinionated decision. - Price-repair engine stays pandas-internal (numerical kernel with scipy.ndimage, not idiomatic DataFrame ops). The pandas frame is converted at the same `history()` boundary.
1 parent 46ce563 commit 5353ffa

14 files changed

Lines changed: 676 additions & 164 deletions

File tree

setup.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@
6868
extras_require={
6969
'nospam': ['requests_cache>=1.0', 'requests_ratelimiter>=0.3.1'],
7070
'repair': ['scipy>=1.6.3'],
71+
'polars': ['polars>=1.0'],
7172
},
7273
# Include protobuf files for websocket support
7374
package_data={

tests/test_dataframe_backend.py

Lines changed: 378 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,378 @@
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

Comments
 (0)