implemented getting specific financial statement for multiple companies; fixed tests

This commit is contained in:
Daniil Fajnberg 2021-11-28 17:07:23 +01:00
parent 1529df252c
commit e522897991
2 changed files with 37 additions and 22 deletions

View File

@ -90,28 +90,43 @@ async def _get_financial_statement(statement: str, ticker_symbol: str, quarterly
return extract_all_data(soup) 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: session: ClientSession = None) -> ResultDict:
""" """
Returns data from the balance sheet of the specified company. 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: session: ClientSession = None) -> ResultDict:
""" """
Returns data from the income statement of the specified company. 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: session: ClientSession = None) -> ResultDict:
""" """
Returns data from the cash flow statement of the specified company. 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, 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: if len(ticker_symbols) == 1:
return { return {
constants.BS: await get_balance_sheet(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, 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, session) constants.CF: await get_cash_flow_statement(ticker_symbols[0], quarterly=quarterly, session=session)
} }
return { return {
sym: { sym: {
constants.BS: await get_balance_sheet(sym, quarterly, session), constants.BS: await get_balance_sheet(sym, quarterly=quarterly, session=session),
constants.IS: await get_income_statement(sym, quarterly, session), constants.IS: await get_income_statement(sym, quarterly=quarterly, session=session),
constants.CF: await get_cash_flow_statement(sym, quarterly, session) constants.CF: await get_cash_flow_statement(sym, quarterly=quarterly, session=session)
} for sym in ticker_symbols } for sym in ticker_symbols
} }

View File

@ -137,7 +137,7 @@ class FunctionsTestCase(IsolatedAsyncioTestCase):
self.assertEqual(expected_output, output) self.assertEqual(expected_output, output)
mock__get_financial_statement.assert_called_once_with(statement, symbol, quarterly, mock_session) mock__get_financial_statement.assert_called_once_with(statement, symbol, quarterly, mock_session)
mock__get_financial_statement.reset_mock() mock__get_financial_statement.reset_mock()
symbol1, symbol2 = 'x', 'y' symbol1, symbol2 = 'x', 'y'
expected_output = {symbol1: expected_output, symbol2: expected_output} expected_output = {symbol1: expected_output, symbol2: expected_output}
output = await function(symbol1, symbol2, quarterly=quarterly, session=mock_session) output = await function(symbol1, symbol2, quarterly=quarterly, session=mock_session)
@ -175,9 +175,9 @@ class FunctionsTestCase(IsolatedAsyncioTestCase):
symbol, quarterly, mock_session = 'foo', False, MagicMock() symbol, quarterly, mock_session = 'foo', False, MagicMock()
output = await functions.get_company_financials(symbol, quarterly=quarterly, session=mock_session) output = await functions.get_company_financials(symbol, quarterly=quarterly, session=mock_session)
self.assertDictEqual(expected_output, output) self.assertDictEqual(expected_output, output)
mock_get_bs.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, mock_session) mock_get_is.assert_called_once_with(symbol, quarterly=quarterly, session=mock_session)
mock_get_cf.assert_called_once_with(symbol, quarterly, mock_session) mock_get_cf.assert_called_once_with(symbol, quarterly=quarterly, session=mock_session)
mock_get_bs.reset_mock() mock_get_bs.reset_mock()
mock_get_is.reset_mock() mock_get_is.reset_mock()
mock_get_cf.reset_mock() mock_get_cf.reset_mock()
@ -188,16 +188,16 @@ class FunctionsTestCase(IsolatedAsyncioTestCase):
quarterly=quarterly, session=mock_session) quarterly=quarterly, session=mock_session)
self.assertDictEqual(expected_output, output) self.assertDictEqual(expected_output, output)
mock_get_bs.assert_has_calls([ mock_get_bs.assert_has_calls([
call(test_symbol1, quarterly, mock_session), call(test_symbol1, quarterly=quarterly, session=mock_session),
call(test_symbol2, quarterly, mock_session) call(test_symbol2, quarterly=quarterly, session=mock_session)
]) ])
mock_get_is.assert_has_calls([ mock_get_is.assert_has_calls([
call(test_symbol1, quarterly, mock_session), call(test_symbol1, quarterly=quarterly, session=mock_session),
call(test_symbol2, quarterly, mock_session) call(test_symbol2, quarterly=quarterly, session=mock_session)
]) ])
mock_get_cf.assert_has_calls([ mock_get_cf.assert_has_calls([
call(test_symbol1, quarterly, mock_session), call(test_symbol1, quarterly=quarterly, session=mock_session),
call(test_symbol2, quarterly, mock_session) call(test_symbol2, quarterly=quarterly, session=mock_session)
]) ])
@patch.object(functions, 'ClientSession') @patch.object(functions, 'ClientSession')