mirror of
https://github.com/apache/superset.git
synced 2026-09-01 13:01:33 +00:00
211 lines
6.9 KiB
Python
211 lines
6.9 KiB
Python
# 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_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",
|
|
]
|