Auxiliary class to collect input data and options for plotting functions, and to check if the inputs are consistent
| 1341 | |
| 1342 | |
| 1343 | class PlottingInput(object): |
| 1344 | """ |
| 1345 | Auxiliary class to collect input data and options for plotting |
| 1346 | functions, and to check if the inputs are consistent |
| 1347 | """ |
| 1348 | supported_datatypes = ( |
| 1349 | pd.Series, |
| 1350 | pd.DataFrame, |
| 1351 | xr.DataArray, |
| 1352 | xr.Dataset, |
| 1353 | ) |
| 1354 | |
| 1355 | def __init__(self, datasets, fields, **argd): |
| 1356 | # Add all arguments as class attributes |
| 1357 | self.__dict__.update({'datasets':datasets.copy(), |
| 1358 | 'fields':fields, |
| 1359 | **argd}) |
| 1360 | |
| 1361 | # Check consistency of all attributes |
| 1362 | self._check_consistency() |
| 1363 | |
| 1364 | def _check_consistency(self): |
| 1365 | """ |
| 1366 | Check consistency of all input data |
| 1367 | """ |
| 1368 | |
| 1369 | # ---------------------- |
| 1370 | # Check dataset argument |
| 1371 | # ---------------------- |
| 1372 | # If a single dataset is provided, convert to a dictionary |
| 1373 | # under a generic key 'Dataset' |
| 1374 | if isinstance(self.datasets, self.supported_datatypes): |
| 1375 | self.datasets = {'Dataset': self.datasets} |
| 1376 | for dfname,df in self.datasets.items(): |
| 1377 | # convert dataset types here |
| 1378 | if isinstance(df, (xr.Dataset,xr.DataArray)): |
| 1379 | # handle xarray datatypes |
| 1380 | if isinstance(df, xr.Dataset): |
| 1381 | df = df[self.fields] |
| 1382 | self.datasets[dfname] = df.to_dataframe() |
| 1383 | columns = self.datasets[dfname].columns |
| 1384 | if len(columns) == 1: |
| 1385 | # convert to pd.Series |
| 1386 | self.datasets[dfname] = self.datasets[dfname][columns[0]] |
| 1387 | else: |
| 1388 | assert(isinstance(df, self.supported_datatypes)), \ |
| 1389 | "Dataset {:s} of type {:s} not supported".format(dfname,str(type(df))) |
| 1390 | |
| 1391 | # ---------------------- |
| 1392 | # Check fields argument |
| 1393 | # ---------------------- |
| 1394 | # If no fields are specified, check that |
| 1395 | # - all datasets are series |
| 1396 | # - the name of every series is either None or matches other series names |
| 1397 | if self.fields is None: |
| 1398 | assert(all([isinstance(self.datasets[dfname],pd.Series) for dfname in self.datasets])), \ |
| 1399 | "'fields' argument must be specified unless all datasets are pandas Series" |
| 1400 | series_names = set() |
no outgoing calls
no test coverage detected