1# -*- coding: utf-8 -*-
2# Licensed under a 3-clause BSD style license - see LICENSE.rst
3
4import pytest
5
6pytest.importorskip('matplotlib')  # noqa
7
8import matplotlib.pyplot as plt
9import matplotlib.dates
10from contextlib import nullcontext
11from erfa import ErfaWarning
12
13from astropy.time import Time
14from astropy.visualization.time import time_support
15
16# Matplotlib 3.3 added a settable epoch for plot dates and changed the default
17# from 0000-12-31 to 1970-01-01. This can be checked by the existence of
18# get_epoch() in matplotlib.dates.
19MPL_EPOCH_1970 = hasattr(matplotlib.dates, 'get_epoch')
20
21# Since some of the examples below use times/dates in the future, we use the
22# TAI time scale to avoid ERFA warnings about dubious years.
23DEFAULT_SCALE = 'tai'
24
25
26def get_ticklabels(axis):
27    axis.figure.canvas.draw()
28    return [x.get_text() for x in axis.get_ticklabels()]
29
30
31def teardown_function(function):
32    plt.close('all')
33
34
35# We first check that we get the expected labels for different time intervals
36# for standard ISO formatting. This is a way to check both the locator and
37# formatter code.
38
39RANGE_CASES = [
40
41    # Interval of many years
42    (('2014-03-22T12:30:30.9', '2077-03-22T12:30:32.1'),
43     ['2020-01-01',
44      '2040-01-01',
45      '2060-01-01']),
46
47    # Interval of a few years
48    (('2014-03-22T12:30:30.9', '2017-03-22T12:30:32.1'),
49     ['2015-01-01',
50      '2016-01-01',
51      '2017-01-01']),
52
53    # Interval of just under a year
54    (('2014-03-22T12:30:30.9', '2015-01-22T12:30:32.1'),
55     ['2014-05-01',
56      '2014-10-01']),
57
58    # Interval of a few months
59    (('2014-11-22T12:30:30.9', '2015-02-22T12:30:32.1'),
60     ['2014-12-01',
61      '2015-01-01',
62      '2015-02-01']),
63
64    # Interval of just over a month
65    (('2014-03-22T12:30:30.9', '2014-04-23T12:30:32.1'),
66     ['2014-04-01']),
67
68    # Interval of just under a month
69    (('2014-03-22T12:30:30.9', '2014-04-21T12:30:32.1'),
70     ['2014-03-24',
71      '2014-04-03',
72      '2014-04-13']),
73
74    # Interval of just over an hour
75    (('2014-03-22T12:30:30.9', '2014-03-22T13:31:30.9'),
76     ['2014-03-22T12:40:00.000',
77      '2014-03-22T13:00:00.000',
78      '2014-03-22T13:20:00.000']),
79
80    # Interval of just under an hour
81    (('2014-03-22T12:30:30.9', '2014-03-22T13:28:30.9'),
82     ['2014-03-22T12:40:00.000',
83      '2014-03-22T13:00:00.000',
84      '2014-03-22T13:20:00.000']),
85
86    # Interval of a few minutes
87    (('2014-03-22T12:30:30.9', '2014-03-22T12:38:30.9'),
88     ['2014-03-22T12:33:00.000',
89      '2014-03-22T12:36:00.000']),
90
91    # Interval of a few seconds
92    (('2014-03-22T12:30:30.9', '2014-03-22T12:30:40.9'),
93     ['2014-03-22T12:30:33.000',
94      '2014-03-22T12:30:36.000',
95      '2014-03-22T12:30:39.000']),
96
97    # Interval of a couple of seconds
98    (('2014-03-22T12:30:30.9', '2014-03-22T12:30:32.1'),
99     ['2014-03-22T12:30:31.000',
100      '2014-03-22T12:30:31.500',
101      '2014-03-22T12:30:32.000']),
102
103    # Interval of under a second
104    (('2014-03-22T12:30:30.89', '2014-03-22T12:30:31.19'),
105     ['2014-03-22T12:30:30.900',
106      '2014-03-22T12:30:31.000',
107      '2014-03-22T12:30:31.100']),
108
109]
110
111
112@pytest.mark.parametrize(('interval', 'expected'), RANGE_CASES)
113def test_formatter_locator(interval, expected):
114
115    # Check that the ticks and labels returned for the above cases are correct.
116
117    with time_support():
118        fig = plt.figure()
119        ax = fig.add_subplot(1, 1, 1)
120        ax.set_xlim(Time(interval[0], scale=DEFAULT_SCALE),
121                    Time(interval[1], scale=DEFAULT_SCALE))
122        assert get_ticklabels(ax.xaxis) == expected
123
124
125FORMAT_CASES = [
126  ('byear', ['2020', '2040', '2060']),
127  ('byear_str', ['B2020.000', 'B2040.000', 'B2060.000']),
128  ('cxcsec', ['1000000000', '1500000000', '2000000000', '2500000000']),
129  ('decimalyear', ['2020', '2040', '2060']),
130  ('fits', ['2020-01-01T00:00:00.000', '2040-01-01T00:00:00.000', '2060-01-01T00:00:00.000']),
131  ('gps', ['1500000000', '2000000000', '2500000000', '3000000000']),
132  ('iso', ['2020-01-01 00:00:00.000', '2040-01-01 00:00:00.000', '2060-01-01 00:00:00.000']),
133  ('isot', ['2020-01-01T00:00:00.000', '2040-01-01T00:00:00.000', '2060-01-01T00:00:00.000']),
134  ('jd', ['2458000', '2464000', '2470000', '2476000']),
135  ('jyear', ['2020', '2040', '2060']),
136  ('jyear_str', ['J2020.000', 'J2040.000', 'J2060.000']),
137  ('mjd', ['60000', '66000', '72000', '78000']),
138  ('plot_date', (['18000', '24000', '30000', '36000'] if MPL_EPOCH_1970 else
139                 ['738000', '744000', '750000', '756000'])),
140  ('unix', ['1500000000', '2000000000', '2500000000', '3000000000']),
141  ('yday', ['2020:001:00:00:00.000', '2040:001:00:00:00.000', '2060:001:00:00:00.000']),
142]
143
144
145@pytest.mark.parametrize(('format', 'expected'), FORMAT_CASES)
146def test_formats(format, expected):
147    # Check that the locators/formatters work fine for all time formats
148    with time_support(format=format, simplify=False):
149        fig = plt.figure()
150        ax = fig.add_subplot(1, 1, 1)
151        # Getting unix time and plot_date requires going through a scale for
152        # which ERFA emits a warning about the date being dubious
153        with pytest.warns(ErfaWarning) if format in ['unix', 'plot_date'] else nullcontext():
154            ax.set_xlim(Time('2014-03-22T12:30:30.9', scale=DEFAULT_SCALE),
155                        Time('2077-03-22T12:30:32.1', scale=DEFAULT_SCALE))
156        assert get_ticklabels(ax.xaxis) == expected
157        ax.get_xlabel() == f'Time ({format})'
158
159
160@pytest.mark.parametrize(('format', 'expected'), FORMAT_CASES)
161def test_auto_formats(format, expected):
162    # Check that the format/scale is taken from the first time used.
163    with time_support(simplify=False):
164        fig = plt.figure()
165        ax = fig.add_subplot(1, 1, 1)
166        # Getting unix time and plot_date requires going through a scale for
167        # which ERFA emits a warning about the date being dubious
168        with pytest.warns(ErfaWarning) if format in ['unix', 'plot_date'] else nullcontext():
169            ax.set_xlim(Time(Time('2014-03-22T12:30:30.9', scale=DEFAULT_SCALE), format=format),
170                        Time('2077-03-22T12:30:32.1', scale=DEFAULT_SCALE))
171        assert get_ticklabels(ax.xaxis) == expected
172        ax.get_xlabel() == f'Time ({format})'
173
174
175FORMAT_CASES_SIMPLIFY = [
176  ('fits', ['2020-01-01', '2040-01-01', '2060-01-01']),
177  ('iso', ['2020-01-01', '2040-01-01', '2060-01-01']),
178  ('isot', ['2020-01-01', '2040-01-01', '2060-01-01']),
179  ('yday', ['2020', '2040', '2060']),
180]
181
182
183@pytest.mark.parametrize(('format', 'expected'), FORMAT_CASES_SIMPLIFY)
184def test_formats_simplify(format, expected):
185    # Check the use of the simplify= option
186    with time_support(format=format, simplify=True):
187        fig = plt.figure()
188        ax = fig.add_subplot(1, 1, 1)
189        ax.set_xlim(Time('2014-03-22T12:30:30.9', scale=DEFAULT_SCALE),
190                    Time('2077-03-22T12:30:32.1', scale=DEFAULT_SCALE))
191        assert get_ticklabels(ax.xaxis) == expected
192
193
194def test_plot():
195    # Make sure that plot() works properly
196    with time_support():
197        fig = plt.figure()
198        ax = fig.add_subplot(1, 1, 1)
199        ax.set_xlim(Time('2014-03-22T12:30:30.9', scale=DEFAULT_SCALE),
200                    Time('2077-03-22T12:30:32.1', scale=DEFAULT_SCALE))
201        ax.plot(Time(['2015-03-22T12:30:30.9',
202                      '2018-03-22T12:30:30.9',
203                      '2021-03-22T12:30:30.9'], scale=DEFAULT_SCALE))
204
205
206def test_nested():
207
208    with time_support(format='iso', simplify=False):
209
210        with time_support(format='yday', simplify=True):
211
212            fig = plt.figure()
213            ax = fig.add_subplot(1, 1, 1)
214            ax.set_xlim(Time('2014-03-22T12:30:30.9', scale=DEFAULT_SCALE),
215                        Time('2077-03-22T12:30:32.1', scale=DEFAULT_SCALE))
216            assert get_ticklabels(ax.xaxis) == ['2020', '2040', '2060']
217
218        fig = plt.figure()
219        ax = fig.add_subplot(1, 1, 1)
220        ax.set_xlim(Time('2014-03-22T12:30:30.9', scale=DEFAULT_SCALE),
221                    Time('2077-03-22T12:30:32.1', scale=DEFAULT_SCALE))
222        assert get_ticklabels(ax.xaxis) == ['2020-01-01 00:00:00.000',
223                                            '2040-01-01 00:00:00.000',
224                                            '2060-01-01 00:00:00.000']
225