# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import io from datetime import datetime, timezone import pandas as pd from openpyxl import load_workbook from pandas.api.types import is_numeric_dtype from superset.utils.core import GenericDataType from superset.utils.excel import ( apply_column_types, df_to_excel, NEUTRAL_TIMESTAMP, quote_formulas, ) def test_timezone_conversion() -> None: """ Test that columns with timezones are converted to a string. """ df = pd.DataFrame({"dt": [datetime(2023, 1, 1, 0, 0, tzinfo=timezone.utc)]}) apply_column_types(df, [GenericDataType.TEMPORAL]) contents = df_to_excel(df) assert pd.read_excel(contents)["dt"][0] == "2023-01-01 00:00:00+00:00" def test_quote_formulas() -> None: """ Test that formulas are quoted in Excel. """ df = pd.DataFrame({"formula": ["=SUM(A1:A2)", "normal", "@SUM(A1:A2)"]}) contents = df_to_excel(df) assert pd.read_excel(contents)["formula"].tolist() == [ "'=SUM(A1:A2)", "normal", "'@SUM(A1:A2)", ] def test_quote_formulas_in_headers_and_index() -> None: """ Test that formulas in column headers and index labels are quoted too. Pivot exports promote data values into the column MultiIndex and row index, so hostile warehouse strings can end up there. """ df = pd.DataFrame( {"=SUM(A1:A2)": ["normal"]}, index=pd.Index(['=cmd|" /C calc"!A0'], name="label"), ) contents = df_to_excel(df) result = pd.read_excel(contents, index_col=0) assert result.columns.tolist() == ["'=SUM(A1:A2)"] assert result.index.tolist() == ['\'=cmd|" /C calc"!A0'] def test_quote_formulas_in_axis_names() -> None: """ Test that formula-triggering axis *names* are quoted too, not just axis labels/values. Pivot exports can promote a warehouse-controlled column name to ``df.index.name`` (or a MultiIndex level name), and pandas writes those names into the sheet as header cells. ``rename`` alone leaves these untouched since they are a separate attribute from the axis labels/values, so this is asserted directly against the ``quote_formulas`` output rather than round-tripped through ``read_excel``, whose header parsing doesn't reliably preserve the columns-axis name. """ df = pd.DataFrame( {"value": ["normal"]}, index=pd.Index(["row"], name="=SUM(A1:A2)"), ) df.columns.name = "+cmd" result = quote_formulas(df) assert result.index.name == "'=SUM(A1:A2)" assert result.columns.name == "'+cmd" # exercised end-to-end to confirm it doesn't error when the axis names # are written out as sheet header cells df_to_excel(df) def test_document_properties_are_neutral() -> None: """ Test that exported workbooks do not carry identifying document properties. """ df = pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}) contents = df_to_excel(df, index=False) workbook = load_workbook(io.BytesIO(contents)) properties = workbook.properties # Authoring/descriptive fields are cleared. for field in ( "creator", "lastModifiedBy", "title", "subject", "description", "keywords", "category", ): value = getattr(properties, field) assert value in (None, ""), f"{field} should be empty, got {value!r}" # Timestamps are pinned to a fixed, neutral value rather than the # actual generation time. assert properties.created == NEUTRAL_TIMESTAMP assert properties.modified == NEUTRAL_TIMESTAMP def test_column_data_types_with_one_numeric_column(): df = pd.DataFrame( { "col0": ["123", "1", "2", "3"], "col1": ["456", "5.67", "0", ".45"], "col2": [ datetime(2023, 1, 1, 0, 0, tzinfo=timezone.utc), datetime(2023, 1, 2, 0, 0, tzinfo=timezone.utc), datetime(2023, 1, 3, 0, 0, tzinfo=timezone.utc), datetime(2023, 1, 4, 0, 0, tzinfo=timezone.utc), ], "col3": ["True", "False", "True", "False"], } ) coltypes: list[GenericDataType] = [ GenericDataType.STRING, GenericDataType.NUMERIC, GenericDataType.TEMPORAL, GenericDataType.BOOLEAN, ] # only col1 should be converted to numeric, according to coltypes definition assert not is_numeric_dtype(df["col1"]) apply_column_types(df, coltypes) assert not is_numeric_dtype(df["col0"]) assert is_numeric_dtype(df["col1"]) assert not is_numeric_dtype(df["col2"]) assert not is_numeric_dtype(df["col3"]) def test_column_data_types_with_failing_conversion(): df = pd.DataFrame( { "col0": ["123", "1", "2", "3"], "col1": ["456", "non_numeric_value", "0", ".45"], "col2": [ datetime(2023, 1, 1, 0, 0, tzinfo=timezone.utc), datetime(2023, 1, 2, 0, 0, tzinfo=timezone.utc), datetime(2023, 1, 3, 0, 0, tzinfo=timezone.utc), datetime(2023, 1, 4, 0, 0, tzinfo=timezone.utc), ], "col3": ["True", "False", "True", "False"], } ) coltypes: list[GenericDataType] = [ GenericDataType.STRING, GenericDataType.NUMERIC, GenericDataType.TEMPORAL, GenericDataType.BOOLEAN, ] # should not fail neither convert assert not is_numeric_dtype(df["col1"]) apply_column_types(df, coltypes) assert not is_numeric_dtype(df["col0"]) assert not is_numeric_dtype(df["col1"]) assert not is_numeric_dtype(df["col2"]) assert not is_numeric_dtype(df["col3"]) def test_apply_column_types_with_duplicate_column_labels() -> None: """ Test that duplicate column labels do not break the export. The verbose_map rename in QueryContextProcessor.get_data can collapse two columns onto the same label, which used to raise "'DataFrame' object has no attribute 'dtype'". """ df = pd.DataFrame( [ ["1", datetime(2023, 1, 1, 0, 0, tzinfo=timezone.utc), "2"], ["3", datetime(2023, 1, 2, 0, 0, tzinfo=timezone.utc), "4"], ], columns=["dupe", "dupe", "other"], ) coltypes: list[GenericDataType] = [ GenericDataType.STRING, GenericDataType.TEMPORAL, GenericDataType.NUMERIC, ] apply_column_types(df, coltypes) # each position is typed independently, despite sharing a label assert not is_numeric_dtype(df.iloc[:, 0]) assert df.iloc[:, 1].tolist() == [ "2023-01-01 00:00:00+00:00", "2023-01-02 00:00:00+00:00", ] assert is_numeric_dtype(df.iloc[:, 2]) contents = df_to_excel(df, index=False) assert pd.read_excel(contents).shape == (2, 3) def test_quote_formulas_with_duplicate_column_labels() -> None: """ Test that formulas are quoted even when column labels are duplicated. """ df = pd.DataFrame( [["=SUM(A1:A2)", "@SUM(A1:A2)", "normal"]], columns=["dupe", "dupe", "other"], ) result = quote_formulas(df) assert result.iloc[0].tolist() == ["'=SUM(A1:A2)", "'@SUM(A1:A2)", "normal"] def test_quote_formulas_with_dedicated_string_dtype() -> None: """ Test that formulas are quoted in columns using the dedicated string dtype. pandas 3 gives string columns a ``str`` dtype rather than ``object``, so an object-only dtype check would skip them and leave formulas unquoted. """ df = pd.DataFrame({"formula": pd.array(["=SUM(A1:A2)", "normal"], dtype="string")}) result = quote_formulas(df) assert result["formula"].tolist() == ["'=SUM(A1:A2)", "normal"] def test_column_data_types_with_large_numeric_values(): df = pd.DataFrame( { "big_number": [ 10**14, 999999999999999, 10**15 + 1, 10**16, 1100108628127863, 2**54, ], } ) apply_column_types(df, [GenericDataType.NUMERIC]) assert df["big_number"].tolist() == [ 100000000000000, 999999999999999, "1000000000000001", "10000000000000000", "1100108628127863", "18014398509481984", ]