def call_load(self, other_args: List[str]):
"""Process load"""
parser = argparse.ArgumentParser(
add_help=False,
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
prog="load",
description="Load custom data set into a dataframe",
)
parser.add_argument(
"-f",
"--file",
choices=self.DATA_FILES,
help="File to load in.",
default="test.csv",
type=str,
)
if other_args and "-" not in other_args[0][0]:
other_args.insert(0, "-f")
ns_parser = parse_known_args_and_warn(parser, other_args)
if ns_parser:
file = Path("custom_imports") / ns_parser.file
self.df = custom_model.load(file)
self.df.columns = self.df.columns.map(lambda x: x.lower())
for col in self.df.columns:
if col in ["date", "time", "timestamp"]: # Could be others?
self.df[col] = pd.to_datetime(self.df[col])
self.df = self.df.set_index(col)
self.file = ns_parser.file
self.update_runtime_choices()
console.print("")
def test_test_empty_load(recorder):
result_df = custom_model.load("")
assert result_df.empty
recorder.capture(result_df)
def test_test_load(recorder, file):
result_df = custom_model.load(os.path.join("custom_imports", file))
assert not result_df.empty
recorder.capture(result_df)