Improves economy/fred command (#4634)

* Fixed

* Started the switch

---------

Co-authored-by: James Maslek <jmaslek11@gmail.com>
This commit is contained in:
Colin Delahunty
2023-04-03 05:07:36 -07:00
committed by GitHub
parent af8cfac18f
commit f4c48dc8d4
3 changed files with 44 additions and 15 deletions

View File

@@ -1389,8 +1389,13 @@ class EconomyController(BaseController):
query = " ".join(ns_parser.query)
df_search = fred_model.get_series_notes(search_query=query)
if not df_search.empty:
fred_view.notes(search_query=query, limit=ns_parser.limit)
if isinstance(df_search, pd.DataFrame) and not df_search.empty:
fred_view.notes(
search_query=query,
limit=ns_parser.limit,
export=ns_parser.export,
sheet_name=ns_parser.sheet_name,
)
self.fred_query = df_search["id"].head(ns_parser.limit)
self.update_runtime_choices()
@@ -1431,7 +1436,7 @@ class EconomyController(BaseController):
get_data=True,
)
if not df.empty:
if isinstance(df, pd.DataFrame) and not df.empty:
for series_id, data in detail.items():
self.FRED_TITLES[
series_id
@@ -1448,7 +1453,7 @@ class EconomyController(BaseController):
if get_current_user().preferences.ENABLE_EXIT_AUTO_HELP:
self.print_help()
else:
elif not ns_parser.export and not ns_parser.raw:
console.print("[red]No data found for the given Series ID[/red]")
elif not parameters and ns_parser.raw:

View File

@@ -40,7 +40,12 @@ def format_units(num: int) -> str:
@log_start_end(log=logger)
@check_api_key(["API_FRED_KEY"])
def notes(search_query: str, limit: int = 10):
def notes(
search_query: str,
limit: int = 10,
export: str = "",
sheet_name: Optional[str] = None,
):
"""Display series notes. [Source: FRED]
Parameters
@@ -49,6 +54,10 @@ def notes(search_query: str, limit: int = 10):
Text query to search on fred series notes database
limit : int
Maximum number of series notes to display
export : str
Export data to csv,json,xlsx or png,jpg,pdf,svg file
sheet_name : Optional[str]
The name of the sheet
"""
df_search = fred_model.get_series_notes(search_query)
@@ -61,6 +70,14 @@ def notes(search_query: str, limit: int = 10):
show_index=False,
headers=["Series ID", "Title", "Description"],
limit=limit,
export=bool(export),
)
export_data(
export,
os.path.dirname(os.path.abspath(__file__)),
"fred",
df_search,
sheet_name,
)
@@ -93,6 +110,8 @@ def display_fred_series(
Output only raw data
export : str
Export data to csv,json,xlsx or png,jpg,pdf,svg file
sheet_name : Optional[str]
The name of the sheet
external_axes : bool, optional
Whether to return the figure object or not, by default False
"""
@@ -132,6 +151,17 @@ def display_fred_series(
)
data.index = [x.strftime("%Y-%m-%d") for x in data.index]
if export:
export_data(
export,
os.path.dirname(os.path.abspath(__file__)),
"fred",
data,
sheet_name,
fig,
)
return None, None
if raw:
data = data.sort_index(ascending=False)
print_rich_table(
@@ -142,15 +172,7 @@ def display_fred_series(
export=bool(export),
limit=limit,
)
export_data(
export,
os.path.dirname(os.path.abspath(__file__)),
"fred",
data,
sheet_name,
fig,
)
return None, None
if get_data:
fig.show(external=external_axes)

View File

@@ -1040,7 +1040,9 @@ def test_call_fred_query(mocker):
controller = economy_controller.EconomyController(queue=None)
controller.choices = {}
controller.call_fred(["--query", "MOCK_QUERY", "--limit", "1"])
mock.assert_called_once_with(**dict(search_query="MOCK_QUERY", limit=1))
mock.assert_called_once_with(
**dict(search_query="MOCK_QUERY", limit=1, export="", sheet_name=None)
)
@pytest.mark.vcr(record_mode="none")