Skip to content

Commit

Permalink
fix(csv): Ensure df_to_escaped_csv handles NULL (apache#20151)
Browse files Browse the repository at this point in the history
Co-authored-by: John Bodley <john.bodley@airbnb.com>
  • Loading branch information
2 people authored and philipher29 committed Jun 9, 2022
1 parent 9bf64de commit 6aa14bc
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 2 deletions.
9 changes: 7 additions & 2 deletions superset/utils/csv.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from typing import Any, Dict, Optional
from urllib.error import URLError

import numpy as np
import pandas as pd
import simplejson

Expand Down Expand Up @@ -64,8 +65,12 @@ def df_to_escaped_csv(df: pd.DataFrame, **kwargs: Any) -> Any:
# Escape csv headers
df = df.rename(columns=escape_values)

# Escape csv rows
df = df.applymap(escape_values)
# Escape csv values
for name, column in df.items():
if column.dtype == np.dtype(object):
for idx, value in enumerate(column.values):
if isinstance(value, str):
df.at[idx, name] = escape_value(value)

return df.to_csv(**kwargs)

Expand Down
4 changes: 4 additions & 0 deletions tests/integration_tests/utils/csv_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import io

import pandas as pd
import pyarrow as pa
import pytest

from superset.utils import csv
Expand Down Expand Up @@ -77,3 +78,6 @@ def test_df_to_escaped_csv():
["a", "'=b"], # pandas seems to be removing the leading ""
["' =a", "b"],
]

df = pa.array([1, None]).to_pandas(integer_object_nulls=True).to_frame()
assert csv.df_to_escaped_csv(df, encoding="utf8", index=False) == '0\n1\n""\n'

0 comments on commit 6aa14bc

Please sign in to comment.