66import numpy as np
77import pandas as pd
88
9- offline = os .getenv ('FORESTPLOT_OFFLINE' )
9+ offline = os .getenv ("FORESTPLOT_OFFLINE" )
10+
1011
1112def insert_groups (
1213 dataframe : pd .core .frame .DataFrame , groupvar : str , varlabel : str
@@ -125,9 +126,9 @@ def insert_empty_row(dataframe: pd.core.frame.DataFrame) -> pd.core.frame.DataFr
125126
126127def load_data (
127128 name : str ,
128- data_path : Path = Path ("./examples/data/" ),
129- ** param_dict : Optional [Any ]
130- ) -> pd .core .frame .DataFrame :
129+ data_path : Union [ Path , str ] = Path ("./examples/data/" ),
130+ ** param_dict : Optional [Any ],
131+ ) -> pd .core .frame .DataFrame :
131132 """
132133 Load example dataset for quickstart.
133134
@@ -151,14 +152,16 @@ def load_data(
151152 available_data = ["mortality" , "sleep" , "sleep-untruncated" ]
152153 name = name .lower ().strip ()
153154 if name in available_data :
154- data = Path (data_path ) / f"{ name } .csv"
155- if not data .is_file ():
156- if offline :
157- raise AssertionError ( f" { data } not found. Working offline (FORESTPLOT_OFFLINE= { offline } )." )
158- data = (
159- f"https://raw.githubusercontent.com/lsys/forestplot/main/examples/data/ { name } .csv "
155+ data_path = Path (data_path ) / f"{ name } .csv"
156+ if data_path .is_file ():
157+ df = pd . read_csv ( data_path , ** param_dict )
158+ elif offline :
159+ raise AssertionError (
160+ f"{ data_path } not found. Working offline (FORESTPLOT_OFFLINE= { offline } ). "
160161 )
161- df = pd .read_csv (data , ** param_dict )
162+ else :
163+ url = f"https://github.com/LSYS/forestplot/tree/main/examples/data/{ name } .csv"
164+ df = pd .read_csv (url , ** param_dict )
162165 if name == "sleep" :
163166 df ["n" ] = df ["n" ].astype ("str" )
164167 return df
0 commit comments