from pandas import DataFrame, read_csv def replace_values(df) -> DataFrame: for column in df.columns: df[column].fillna(value=df[column].mean(), inplace=True) return df def process_na(df, action) -> DataFrame: if action == "drop": return df.dropna() elif action == "fill": return replace_values(df) else: print("Unknown action selected. The choices are: ") print("fill: fills the na values with the mean") print("drop: drops the na values") exit() def filter_dataframe(df) -> DataFrame: relevant_columns = [ "HORA", "DIASEMANA", "COMUNIDAD_AUTONOMA", "ISLA", "TOT_HERIDOS_LEVES", "TOT_HERIDOS_GRAVES", "TOT_VEHICULOS_IMPLICADOS", "TOT_MUERTOS", "TIPO_VIA", "LUMINOSIDAD", "FACTORES_ATMOSFERICOS", ] filtered_df = df.filter(items=relevant_columns) return filtered_df def parse_data(source, action) -> DataFrame: df = read_csv(filepath_or_buffer=source, na_values="?") processed_df = process_na(df=df, action=action) return processed_df