From 952787281997875214159e5a4bcfde0d1303b6d1 Mon Sep 17 00:00:00 2001 From: John Bodley Date: Fri, 22 Jun 2018 18:17:12 -0700 Subject: [PATCH] [get_df] Fix datetime conversion --- superset/viz.py | 8 ++--- tests/viz_tests.py | 81 +++++++++++++++++++++------------------------- 2 files changed, 40 insertions(+), 49 deletions(-) diff --git a/superset/viz.py b/superset/viz.py index 89b806596a50..a595e49b68a9 100644 --- a/superset/viz.py +++ b/superset/viz.py @@ -197,13 +197,11 @@ def get_df(self, query_obj=None): # be considered as the default ISO date format # If the datetime format is unix, the parse will use the corresponding # parsing logic. - if df is None or df.empty: - return pd.DataFrame() - else: + if not df.empty: if DTTM_ALIAS in df.columns: if timestamp_format in ('epoch_s', 'epoch_ms'): - df[DTTM_ALIAS] = pd.to_datetime( - df[DTTM_ALIAS], utc=False, unit=timestamp_format[6:]) + # Column has already been formatted as a timestamp. + df[DTTM_ALIAS] = df[DTTM_ALIAS].apply(pd.Timestamp) else: df[DTTM_ALIAS] = pd.to_datetime( df[DTTM_ALIAS], utc=False, format=timestamp_format) diff --git a/tests/viz_tests.py b/tests/viz_tests.py index 30c37c84f1af..1bc6e62a9744 100644 --- a/tests/viz_tests.py +++ b/tests/viz_tests.py @@ -38,17 +38,13 @@ def test_get_fillna_returns_default_on_null_columns(self): def test_get_df_returns_empty_df(self): datasource = Mock() datasource.type = 'table' - mock_dttm_col = Mock() - mock_dttm_col.python_date_format = Mock() - datasource.get_col = Mock(return_value=mock_dttm_col) form_data = {'dummy': 123} query_obj = {'granularity': 'day'} results = Mock() results.query = Mock() results.status = Mock() results.error_message = None - results.df = Mock() - results.df.empty = True + results.df = pd.DataFrame() datasource.query = Mock(return_value=results) test_viz = viz.BaseViz(datasource, form_data) result = test_viz.get_df(query_obj) @@ -56,55 +52,52 @@ def test_get_df_returns_empty_df(self): self.assertTrue(result.empty) def test_get_df_handles_dttm_col(self): - datasource = Mock() - datasource.type = 'table' - datasource.offset = 1 - mock_dttm_col = Mock() - mock_dttm_col.python_date_format = 'epoch_ms' - datasource.get_col = Mock(return_value=mock_dttm_col) form_data = {'dummy': 123} query_obj = {'granularity': 'day'} results = Mock() results.query = Mock() results.status = Mock() results.error_message = Mock() - df = Mock() - df.columns = [DTTM_ALIAS] - f_datetime = datetime(1960, 1, 1, 5, 0) - df.__getitem__ = Mock(return_value=pd.Series([f_datetime])) - df.__setitem__ = Mock() - df.replace = Mock() - df.fillna = Mock() - results.df = df - results.df.empty = False + datasource = Mock() + datasource.type = 'table' datasource.query = Mock(return_value=results) + mock_dttm_col = Mock() + datasource.get_col = Mock(return_value=mock_dttm_col) test_viz = viz.BaseViz(datasource, form_data) - test_viz.df_metrics_to_num = Mock() test_viz.get_fillna_for_columns = Mock(return_value=0) - test_viz.get_df(query_obj) - mock_call = df.__setitem__.mock_calls[0] - self.assertEqual(mock_call[1][0], DTTM_ALIAS) - self.assertFalse(mock_call[1][1].empty) - self.assertEqual(mock_call[1][1][0], f_datetime) - mock_call = df.__setitem__.mock_calls[1] - self.assertEqual(mock_call[1][0], DTTM_ALIAS) - self.assertEqual(mock_call[1][1][0].hour, 6) - self.assertEqual(mock_call[1][1].dtype, 'datetime64[ns]') - mock_dttm_col.python_date_format = 'utc' - test_viz.get_df(query_obj) - mock_call = df.__setitem__.mock_calls[2] - self.assertEqual(mock_call[1][0], DTTM_ALIAS) - self.assertFalse(mock_call[1][1].empty) - self.assertEqual(mock_call[1][1][0].hour, 7) - mock_call = df.__setitem__.mock_calls[3] - self.assertEqual(mock_call[1][0], DTTM_ALIAS) - self.assertEqual(mock_call[1][1][0].hour, 6) - self.assertEqual(mock_call[1][1].dtype, 'datetime64[ns]') - mock_call = df.__setitem__.mock_calls[4] - self.assertEqual(mock_call[1][0], DTTM_ALIAS) - self.assertEqual(mock_call[1][1][0].hour, 7) - self.assertEqual(mock_call[1][1].dtype, 'datetime64[ns]') + + results.df = pd.DataFrame(data={DTTM_ALIAS: ['1960-01-01 05:00:00']}) + datasource.offset = 0 + mock_dttm_col.python_date_format = 'epoch_ms' + result = test_viz.get_df(query_obj) + pd.testing.assert_series_equal( + result[DTTM_ALIAS], + pd.Series([datetime(1960, 1, 1, 5, 0)], name=DTTM_ALIAS), + ) + + mock_dttm_col.python_date_format = None + result = test_viz.get_df(query_obj) + pd.testing.assert_series_equal( + result[DTTM_ALIAS], + pd.Series([datetime(1960, 1, 1, 5, 0)], name=DTTM_ALIAS), + ) + + datasource.offset = 1 + result = test_viz.get_df(query_obj) + pd.testing.assert_series_equal( + result[DTTM_ALIAS], + pd.Series([datetime(1960, 1, 1, 6, 0)], name=DTTM_ALIAS), + ) + + datasource.offset = 0 + results.df = pd.DataFrame(data={DTTM_ALIAS: ['1960-01-01']}) + mock_dttm_col.python_date_format = '%Y-%m-%d' + result = test_viz.get_df(query_obj) + pd.testing.assert_series_equal( + result[DTTM_ALIAS], + pd.Series([datetime(1960, 1, 1, 0, 0)], name=DTTM_ALIAS), + ) def test_cache_timeout(self): datasource = Mock()