ENH: better detection of categorical columns in plot and explore (#2470)

This commit is contained in:
keirayuki310
2022-09-02 08:32:07 +10:00
committed by GitHub
parent 4d29fcfeb0
commit 4f4784b372
4 changed files with 34 additions and 8 deletions
+3 -2
View File
@@ -378,8 +378,9 @@ GON (((180.00000 -16.06713, 180.00000...
)
categorical = True
elif (
gdf[column].dtype is np.dtype("O")
or gdf[column].dtype is np.dtype(bool)
pd.api.types.is_object_dtype(gdf[column])
or pd.api.types.is_bool_dtype(gdf[column])
or pd.api.types.is_string_dtype(gdf[column])
or categories
):
categorical = True
+6 -1
View File
@@ -733,7 +733,12 @@ GON (((-122.84000 49.00000, -120.0000...
"Cannot specify 'categories' when column has categorical dtype"
)
categorical = True
elif values.dtype is np.dtype("O") or categories:
elif (
pd.api.types.is_object_dtype(values.dtype)
or pd.api.types.is_bool_dtype(values.dtype)
or pd.api.types.is_string_dtype(values.dtype)
or categories
):
categorical = True
nan_idx = np.asarray(pd.isna(values), dtype="bool")
+17 -3
View File
@@ -256,10 +256,24 @@ class TestExplore:
def test_bool(self):
df = self.nybb.copy()
df["bool"] = [True, False, True, False, True]
m = df.explore("bool")
df["bool_extension"] = pd.array([True, False, True, False, True])
m1 = df.explore("bool")
m2 = df.explore("bool_extension")
out1_str = self._fetch_map_string(m1)
assert '"__folium_color":"#9edae5","bool":true' in out1_str
assert '"__folium_color":"#1f77b4","bool":false' in out1_str
out2_str = self._fetch_map_string(m2)
assert '"__folium_color":"#9edae5","bool":true' in out2_str
assert '"__folium_color":"#1f77b4","bool":false' in out2_str
def test_string(self):
df = self.nybb.copy()
df["string"] = pd.array([1, 2, 3, 4, 5], dtype="string")
m = df.explore("string")
out_str = self._fetch_map_string(m)
assert '"__folium_color":"#9edae5","bool":true' in out_str
assert '"__folium_color":"#1f77b4","bool":false' in out_str
assert '"__folium_color":"#9edae5","string":"5"' in out_str
def test_column_values(self):
"""
+8 -2
View File
@@ -374,6 +374,9 @@ class TestPointPlotting:
self.df["cats_ordered"] = pd.Categorical(
["cat2", "cat1"] * 5, categories=["cat2", "cat1"]
)
self.df["bool"] = [False, True] * 5
self.df["bool_extension"] = pd.array([False, True] * 5)
self.df["cats_string"] = pd.array(["cat1", "cat2"] * 5, dtype="string")
ax1 = self.df.plot("cats_object", legend=True)
ax2 = self.df.plot("cats", legend=True)
@@ -381,14 +384,17 @@ class TestPointPlotting:
ax4 = self.df.plot("singlecat", legend=True)
ax5 = self.df.plot("cats_ordered", legend=True)
ax6 = self.df.plot("nums", categories=[1, 2], legend=True)
ax7 = self.df.plot("bool", legend=True)
ax8 = self.df.plot("bool_extension", legend=True)
ax9 = self.df.plot("cats_string", legend=True)
point_colors1 = ax1.collections[0].get_facecolors()
for ax in [ax2, ax3, ax4, ax5, ax6]:
for ax in [ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9]:
point_colors2 = ax.collections[0].get_facecolors()
np.testing.assert_array_equal(point_colors1[1], point_colors2[1])
legend1 = [x.get_markerfacecolor() for x in ax1.get_legend().get_lines()]
for ax in [ax2, ax3, ax4, ax5, ax6]:
for ax in [ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9]:
legend2 = [x.get_markerfacecolor() for x in ax.get_legend().get_lines()]
np.testing.assert_array_equal(legend1, legend2)