Files
OpenBB/openbb_terminal/common/technical_analysis/ta_helpers.py
Colin Delahunty dcf3bfe1f7 Refactor common/ta (#3226)
* Handle no adj close

* Added fixes for bad data handling

* Finished switching model layout

* Fixed all the bugs

* Big fixes

* Handle a lot more bugs

* A couple more enhancements

* Fixed issue

Co-authored-by: Henrique Joaquim <h.joaquim@campus.fct.unl.pt>
Co-authored-by: James Maslek <jmaslek11@gmail.com>
2022-11-02 11:30:28 -04:00

53 lines
1.4 KiB
Python

from typing import Optional
import logging
import pandas as pd
from openbb_terminal.rich_config import console
logger = logging.getLogger(__name__)
def check_columns(
data: pd.DataFrame, high: bool = True, low: bool = True, close: bool = True
) -> Optional[str]:
"""Returns the close columns, or None if the dataframe does not have required columns
Parameters
----------
data: pd.DataFrame
The dataframe to check
high: bool
Whether to check for high column
low: bool
Whether to check for low column
close: bool
Whether to check for close column
Returns
----------
Optional[str]
The name of the close column, none if df is invalid
"""
close_col = "Close"
for title in ["Adj Close", "Close", "adj_close", "Adj_Close"]:
if title in data.columns:
close_col = title
break
close_col = "Adj Close" if "Adj Close" in data.columns else "Close"
# pylint: disable=too-many-boolean-expressions
if (
("High" not in data.columns and high)
or ("Low" not in data.columns and low)
or (close_col not in data.columns and close)
):
logger.error("Invalid columns. data has columns %s", data.columns)
console.print(
"[red] Please make sure that the columns 'High', 'Low', and 'Close'"
" are in the dataframe.[/red]"
)
return None
return close_col