diff --git a/src/mwfin/functions.py b/src/mwfin/functions.py index e89875e..2d03225 100644 --- a/src/mwfin/functions.py +++ b/src/mwfin/functions.py @@ -90,28 +90,43 @@ async def _get_financial_statement(statement: str, ticker_symbol: str, quarterly return extract_all_data(soup) -async def get_balance_sheet(ticker_symbol: str, quarterly: bool = False, +async def get_balance_sheet(*ticker_symbols: str, quarterly: bool = False, session: ClientSession = None) -> ResultDict: """ Returns data from the balance sheet of the specified company. """ - return await _get_financial_statement(constants.BS, ticker_symbol, quarterly, session) + if len(ticker_symbols) == 1: + return await _get_financial_statement(constants.BS, ticker_symbols[0], quarterly, session) + return { + sym: await _get_financial_statement(constants.BS, sym, quarterly, session) + for sym in ticker_symbols + } -async def get_income_statement(ticker_symbol: str, quarterly: bool = False, +async def get_income_statement(*ticker_symbols: str, quarterly: bool = False, session: ClientSession = None) -> ResultDict: """ Returns data from the income statement of the specified company. """ - return await _get_financial_statement(constants.IS, ticker_symbol, quarterly, session) + if len(ticker_symbols) == 1: + return await _get_financial_statement(constants.IS, ticker_symbols[0], quarterly, session) + return { + sym: await _get_financial_statement(constants.IS, sym, quarterly, session) + for sym in ticker_symbols + } -async def get_cash_flow_statement(ticker_symbol: str, quarterly: bool = False, +async def get_cash_flow_statement(*ticker_symbols: str, quarterly: bool = False, session: ClientSession = None) -> ResultDict: """ Returns data from the cash flow statement of the specified company. """ - return await _get_financial_statement(constants.CF, ticker_symbol, quarterly, session) + if len(ticker_symbols) == 1: + return await _get_financial_statement(constants.CF, ticker_symbols[0], quarterly, session) + return { + sym: await _get_financial_statement(constants.CF, sym, quarterly, session) + for sym in ticker_symbols + } async def get_company_financials(*ticker_symbols: str, quarterly: bool = False, @@ -121,14 +136,14 @@ async def get_company_financials(*ticker_symbols: str, quarterly: bool = False, """ if len(ticker_symbols) == 1: return { - constants.BS: await get_balance_sheet(ticker_symbols[0], quarterly, session), - constants.IS: await get_income_statement(ticker_symbols[0], quarterly, session), - constants.CF: await get_cash_flow_statement(ticker_symbols[0], quarterly, session) + constants.BS: await get_balance_sheet(ticker_symbols[0], quarterly=quarterly, session=session), + constants.IS: await get_income_statement(ticker_symbols[0], quarterly=quarterly, session=session), + constants.CF: await get_cash_flow_statement(ticker_symbols[0], quarterly=quarterly, session=session) } return { sym: { - constants.BS: await get_balance_sheet(sym, quarterly, session), - constants.IS: await get_income_statement(sym, quarterly, session), - constants.CF: await get_cash_flow_statement(sym, quarterly, session) + constants.BS: await get_balance_sheet(sym, quarterly=quarterly, session=session), + constants.IS: await get_income_statement(sym, quarterly=quarterly, session=session), + constants.CF: await get_cash_flow_statement(sym, quarterly=quarterly, session=session) } for sym in ticker_symbols } diff --git a/tests/test_functions.py b/tests/test_functions.py index b46d0db..47b91c4 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -137,7 +137,7 @@ class FunctionsTestCase(IsolatedAsyncioTestCase): self.assertEqual(expected_output, output) mock__get_financial_statement.assert_called_once_with(statement, symbol, quarterly, mock_session) mock__get_financial_statement.reset_mock() - + symbol1, symbol2 = 'x', 'y' expected_output = {symbol1: expected_output, symbol2: expected_output} output = await function(symbol1, symbol2, quarterly=quarterly, session=mock_session) @@ -175,9 +175,9 @@ class FunctionsTestCase(IsolatedAsyncioTestCase): symbol, quarterly, mock_session = 'foo', False, MagicMock() output = await functions.get_company_financials(symbol, quarterly=quarterly, session=mock_session) self.assertDictEqual(expected_output, output) - mock_get_bs.assert_called_once_with(symbol, quarterly, mock_session) - mock_get_is.assert_called_once_with(symbol, quarterly, mock_session) - mock_get_cf.assert_called_once_with(symbol, quarterly, mock_session) + mock_get_bs.assert_called_once_with(symbol, quarterly=quarterly, session=mock_session) + mock_get_is.assert_called_once_with(symbol, quarterly=quarterly, session=mock_session) + mock_get_cf.assert_called_once_with(symbol, quarterly=quarterly, session=mock_session) mock_get_bs.reset_mock() mock_get_is.reset_mock() mock_get_cf.reset_mock() @@ -188,16 +188,16 @@ class FunctionsTestCase(IsolatedAsyncioTestCase): quarterly=quarterly, session=mock_session) self.assertDictEqual(expected_output, output) mock_get_bs.assert_has_calls([ - call(test_symbol1, quarterly, mock_session), - call(test_symbol2, quarterly, mock_session) + call(test_symbol1, quarterly=quarterly, session=mock_session), + call(test_symbol2, quarterly=quarterly, session=mock_session) ]) mock_get_is.assert_has_calls([ - call(test_symbol1, quarterly, mock_session), - call(test_symbol2, quarterly, mock_session) + call(test_symbol1, quarterly=quarterly, session=mock_session), + call(test_symbol2, quarterly=quarterly, session=mock_session) ]) mock_get_cf.assert_has_calls([ - call(test_symbol1, quarterly, mock_session), - call(test_symbol2, quarterly, mock_session) + call(test_symbol1, quarterly=quarterly, session=mock_session), + call(test_symbol2, quarterly=quarterly, session=mock_session) ]) @patch.object(functions, 'ClientSession')