Commit 7a059a0d authored by Martin Bergemann's avatar Martin Bergemann
Browse files

Make get_timestamp method more flexible

parent c0143b4c
......@@ -10,6 +10,7 @@ from tempfile import NamedTemporaryFile, TemporaryDirectory
from cartopy import crs
import cftime
import cartopy
import matplotlib
matplotlib.use('Agg')
......@@ -49,7 +50,16 @@ def get_filenames(fname, start, end):
========
fname (str, pathlib.Path)
"""
if isinstance(start, str) and start:
try:
start = pd.Timestamp(start)
except:
start = None
if isinstance(end, str) and end:
try:
end = pd.Timestamp(end)
except:
end = None
try:
start = start.year
except AttributeError:
......@@ -62,14 +72,14 @@ def get_filenames(fname, start, end):
fname2 = Path(fname).with_suffix('').name
except ValueError:
return
if not fname.strip():
if not str(fname).strip():
return
try:
ts = pd.DatetimeIndex(fname2.split('_')[-1].strip().split('-'))
ts = pd.DatetimeIndex(str(fname2).split('_')[-1].strip().split('-'))
except:
return fname.strip()
if ts[0].year >= start and ts[-1].year <= end and fname.strip():
return fname.strip()
return str(fname).strip()
if ts[0].year >= start and ts[-1].year <= end and str(fname).strip():
return str(fname).strip()
class MapCreator:
......@@ -311,19 +321,26 @@ class Plotter:
self.cbar.set_label(cbar_label, fontsize=self.fontsize)
return im
def get_timestamp(self, dset, fmt=' (%Y/%m/%d %H:%M)'):
@staticmethod
def get_timestamp(dset, fmt=' (%Y/%m/%d %H:%M)', tz='utc'):
"""Construt the string repr. of a timestep from a given xr.DataArray."""
if dset.time.values == None:
return ''
try:
time = pd.DatetimeIndex(np.array([dset.time.data]), tz='utc').to_pydatetime()[0]
except AttributeError:
return ''
time_lt = time.astimezone(pytz.timezone(self.tz))
except TypeError:
units = 'seconds since 1970-01-01:00:00'
tt = cftime.date2num(dset.time.values, units)
return cftime.num2date(tt, units).strftime(fmt)
time_lt = time.astimezone(pytz.timezone(tz))
return time_lt.strftime(fmt)
def get_cbar_label(self, dset):
"""Construct the label of the colorbar."""
cbar_label = self._cbar_label
return self._cbar_label + self.get_timestamp(dset)
return self._cbar_label + self.get_timestamp(dset, tz=self.tz)
def plot_map(dset, lon, lat, *args, dirname=None, **kwargs):
......@@ -353,7 +370,7 @@ def plot_map(dset, lon, lat, *args, dirname=None, **kwargs):
plot.im = getattr(plot, plot_method)(dset, **plot._kwargs)
if not dirname:
return plot
ts = plot.get_timestamp(dset, fmt='_%Y%m%d_%H%M')
ts = plot.get_timestamp(dset, fmt='_%Y%m%d_%H%M', tz=plot.tz)
fn = f'{dset.name}'
if ts:
fn += ts
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment