Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_plot_creation_01_array(self):
"""Test the plot creation with a plotter that takes one array"""
psy.register_plotter('test_plotter',
import_plotter=True, module='test_plotter',
plotter_name='TestPlotter')
ds = psy.open_dataset(bt.get_file('test-t2m-u-v.nc'))
sp = psy.plot.test_plotter(ds, name='t2m')
self.assertEqual(len(sp), 1)
self.assertEqual(sp[0].name, 't2m')
self.assertEqual(sp[0].shape, ds.t2m.shape)
self.assertEqual(sp[0].values.tolist(), ds.t2m.values.tolist())
psy.close()
psy.unregister_plotter('test_plotter')
def test_figs(self):
"""Test the :attr:`psyplot.project.Project.figs` attribute"""
psy.register_plotter('test_plotter', import_plotter=True,
module='test_plotter', plotter_name='TestPlotter')
ds = psy.open_dataset(bt.get_file('test-t2m-u-v.nc'))
sp = psy.plot.test_plotter(ds, name='t2m', time=[1, 2])
self.assertEqual(sp[0].psy.ax.figure.number, 1)
self.assertEqual(sp[1].psy.ax.figure.number, 2)
figs = sp.figs
self.assertIn(sp[0].psy.ax.figure, figs)
self.assertIs(figs[sp[0].psy.ax.figure][0], sp[0])
self.assertIn(sp[1].psy.ax.figure, figs)
self.assertIs(figs[sp[1].psy.ax.figure][0], sp[1])
def test_plot_creation_11_post_fmt(self):
"""Test the :attr:`psyplot.plotter.Plotter.post` formatoption"""
psy.register_plotter('test_plotter',
import_plotter=True, module='test_plotter',
plotter_name='TestPlotter')
ds = psy.open_dataset(bt.get_file('test-t2m-u-v.nc'))
# test whether it is plotted automatically
sp = psy.plot.test_plotter(ds, name='t2m',
post='self.ax.set_title("test")')
self.assertEqual(sp.plotters[0].ax.get_title(), 'test')
# test whether the disabling works
sp = psy.plot.test_plotter(ds, name='t2m', enable_post=False,
post='self.ax.set_title("test")')
self.assertEqual(sp.plotters[0].ax.get_title(), '')
def test_plot_creation_02_array_default_dims(self):
# add a default value for the y dimension
psy.register_plotter('test_plotter',
import_plotter=True, module='test_plotter',
plotter_name='TestPlotter',
default_dims={'y': 0})
ds = psy.open_dataset(bt.get_file('test-t2m-u-v.nc'))
sp = psy.plot.test_plotter(ds, name='t2m')
self.assertEqual(len(sp), 1)
self.assertEqual(sp[0].name, 't2m')
self.assertEqual(sp[0].shape, ds.t2m.isel(lat=0).shape)
self.assertEqual(sp[0].values.tolist(),
ds.t2m.isel(lat=0).values.tolist())
psy.close()
psy.unregister_plotter('test_plotter')
def test_plot_creation_03_2arrays(self):
# try multiple names and dimension
psy.register_plotter('test_plotter',
import_plotter=True, module='test_plotter',
plotter_name='TestPlotter',
default_dims={'y': 0})
ds = psy.open_dataset(bt.get_file('test-t2m-u-v.nc'))
sp = psy.plot.test_plotter(ds, name=['t2m', 'u'], x=slice(3, 5))
self.assertEqual(len(sp), 2)
self.assertEqual(sp[0].name, 't2m')
self.assertEqual(sp[1].name, 'u')
self.assertEqual(sp[0].shape,
ds.t2m.isel(lat=0, lon=slice(3, 5)).shape)
self.assertEqual(sp[1].shape,
ds.u.isel(lat=0, lon=slice(3, 5)).shape)
self.assertEqual(sp[0].values.tolist(),
ds.t2m.isel(lat=0, lon=slice(3, 5)).values.tolist())
self.assertEqual(sp[1].values.tolist(),
ds.u.isel(lat=0, lon=slice(3, 5)).values.tolist())
psy.close()
psy.unregister_plotter('test_plotter')
def test_export_01_replacement(self):
"""Test exporting a project"""
import matplotlib.pyplot as plt
from matplotlib.testing.compare import compare_images
import pandas as pd
import numpy as np
from tempfile import NamedTemporaryFile
self._register_export_plotter()
with psy.open_dataset(bt.get_file('test-t2m-u-v.nc')) as ds:
time = ds.time
time.values # make sure the data is loaded
ds = xr.Dataset(
{"v0": xr.Variable(('x', 'y'), np.arange(3 * 5).reshape(3, 5)),
"v1": xr.Variable(('time', 'y'), np.arange(5 * 5).reshape(5, 5))},
{"x": xr.Variable(('x', ), [4, 5, 6]),
"y": xr.Variable(('y', ), [6, 7, 8, 9, 10]),
'time': time})
# create reference plots
reffiles = []
fig, ax = plt.subplots()
ds.v0[1].plot(ax=ax)
reffiles.append(
NamedTemporaryFile(prefix='psyplot_', suffix='.png').name)
self._created_files.update(reffiles)
# test list of list with array and an array out of 2 variables
# psyplot.project.Project([
# arr0: psyplot.data.InteractiveList([
# arr0: 3-dim DataArray of t2m, with
# (time, lev, lat)=(5, 4, 96), lon=1.875,
# arr1: 4-dim DataArray of u, v, with
# (variable, time, lev, lat)=(2, 5, 4, 96), lon=1.875]),
# arr1: psyplot.data.InteractiveList([
# arr0: 3-dim DataArray of t2m, with
# (time, lev, lat)=(5, 4, 96), lon=1.875,
# arr1: 4-dim DataArray of u, v, with
# (variable, time, lev, lat)=(2, 5, 4, 96), lon=1.875])])
psy.register_plotter('test_plotter',
import_plotter=True, module='test_plotter',
plotter_name='TestPlotter', prefer_list=True)
ds = psy.open_dataset(bt.get_file('test-t2m-u-v.nc'))
sp = psy.plot.test_plotter(ds, name=[['t2m', ['u', 'v']]], x=[1, 2])
self.assertEqual(len(sp), 2)
self.assertEqual(len(sp[0]), 2)
self.assertEqual(len(sp[1]), 2)
self.assertEqual(sp[0][0].name, 't2m')
self.assertIn('variable', sp[0][1].dims)
self.assertEqual(sp[0][1].coords['variable'].values.tolist(),
['u', 'v'])
self.assertEqual(sp[1][0].name, 't2m')
self.assertIn('variable', sp[1][1].dims)
self.assertEqual(sp[1][1].coords['variable'].values.tolist(),
['u', 'v'])
self.assertEqual(sp[0][0].shape, ds.t2m.isel(lon=1).shape)
self.assertEqual(list(sp[0][1].shape),
[2] + list(ds.u.isel(lon=1).shape))
self.assertEqual(sp[1][0].shape, ds.t2m.isel(lon=2).shape)
def test_plot_creation_07_list_and_dims(self):
# use dimensions which should result in one list with 4 arrays,
# t2m, t2m, u, u
# psyplot.project.Project([arr3: psyplot.data.InteractiveList([
# arr0: 3-dim DataArray of t2m, with
# (time, lev, lat)=(5, 4, 96), lon=1.875,
# arr1: 3-dim DataArray of t2m, with
# (time, lev, lat)=(5, 4, 96), lon=3.75,
# arr2: 3-dim DataArray of u, with
# (time, lev, lat)=(5, 4, 96), lon=1.875,
# arr3: 3-dim DataArray of u, with
# (time, lev, lat)=(5, 4, 96), lon=3.75])])
psy.register_plotter('test_plotter',
import_plotter=True, module='test_plotter',
plotter_name='TestPlotter', prefer_list=True)
ds = psy.open_dataset(bt.get_file('test-t2m-u-v.nc'))
sp = psy.plot.test_plotter(ds, name=['t2m', 'u'], x=[1, 2])
self.assertEqual(len(sp), 1)
self.assertEqual(len(sp[0]), 4)
self.assertEqual(sp[0][0].name, 't2m')
self.assertEqual(sp[0][1].name, 't2m')
self.assertEqual(sp[0][2].name, 'u')
self.assertEqual(sp[0][3].name, 'u')
self.assertEqual(sp[0][0].shape, ds.t2m.isel(lon=1).shape)
self.assertEqual(sp[0][1].shape, ds.t2m.isel(lon=2).shape)
self.assertEqual(sp[0][2].shape, ds.u.isel(lon=1).shape)
self.assertEqual(sp[0][3].shape, ds.u.isel(lon=2).shape)
self.assertEqual(sp[0][0].values.tolist(),
ds.t2m.isel(lon=1).values.tolist())
self.assertEqual(sp[0][1].values.tolist(),
ds.t2m.isel(lon=2).values.tolist())
self.assertEqual(sp[0][2].values.tolist(),
def test_save_and_load_07_sharedx(self):
"""Test whether shared x- and y-axis are restored correctly"""
psy.register_plotter('test_plotter', import_plotter=True,
module='test_plotter', plotter_name='TestPlotter')
ds = psy.open_dataset(bt.get_file('test-t2m-u-v.nc'))
plt.close('all')
fig, axes = plt.subplots(1, 3, sharex=True)
sp = psy.plot.test_plotter(ds, name=['t2m', 'u', 'v'], x=0, y=4,
ax=axes)
axes[0].set_xlim(5, 10)
self.assertEqual(list(axes[1].get_xlim()), [5, 10])
# save the project
fname = 'test.pkl'
self._created_files.add(fname)
sp.save_project(fname)
psy.close('all')
# load the project
sp = psy.Project.load_project(fname)
self.assertEqual(len(sp.axes), 3, msg=sp.axes)
sp[0].psy.ax.set_xlim(10, 15)
# test list of list of arrays
# psyplot.project.Project([
# arr0: psyplot.data.InteractiveList([
# arr0: 3-dim DataArray of t2m, with
# (time, lev, lat)=(5, 4, 96), lon=1.875,
# arr1: 3-dim DataArray of u, with #
# (time, lev, lat)=(5, 4, 96), lon=1.875]),
# arr1: psyplot.data.InteractiveList([
# arr0: 3-dim DataArray of t2m, with
# (time, lev, lat)=(5, 4, 96), lon=3.75,
# arr1: 3-dim DataArray of u, with
# (time, lev, lat)=(5, 4, 96), lon=3.75])])
psy.register_plotter('test_plotter',
import_plotter=True, module='test_plotter',
plotter_name='TestPlotter', prefer_list=True)
ds = psy.open_dataset(bt.get_file('test-t2m-u-v.nc'))
sp = psy.plot.test_plotter(bt.get_file('test-t2m-u-v.nc'),
name=[['t2m', 'u']], x=[1, 2])
self.assertEqual(len(sp), 2)
self.assertEqual(len(sp[0]), 2)
self.assertEqual(len(sp[1]), 2)
self.assertEqual(sp[0][0].name, 't2m')
self.assertEqual(sp[0][1].name, 'u')
self.assertEqual(sp[1][0].name, 't2m')
self.assertEqual(sp[1][1].name, 'u')
self.assertEqual(sp[0][0].shape, ds.t2m.isel(lon=1).shape)
self.assertEqual(sp[0][1].shape, ds.u.isel(lon=1).shape)
self.assertEqual(sp[1][0].shape, ds.t2m.isel(lon=2).shape)
self.assertEqual(sp[1][1].shape, ds.u.isel(lon=2).shape)
self.assertEqual(sp[0][0].values.tolist(),
ds.t2m.isel(lon=1).values.tolist())
self.assertEqual(sp[0][1].values.tolist(),