Fixed bug that can occur if nested arrays/dicts are passed to columnd…

…atasource
rotemgor · Sep 27, 2018 · 237e7c8 · 237e7c8
1 parent 4f5a233
commit 237e7c8
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 5 deletions.
diff --git a/chartify/_core/plot.py b/chartify/_core/plot.py
@@ -328,10 +328,16 @@ def line(self,
         for color_value, color in zip(color_values, colors):
 
             if color_column is None:  # Single line
-                sliced_data = data_frame[[x_column, y_column]]
+                sliced_data = data_frame
             else:
                 sliced_data = data_frame[data_frame[color_column] ==
                                          color_value]
+            # Filter to only relevant columns.
+            sliced_data = (
+                sliced_data[
+                    [col for col in sliced_data.columns
+                        if col in (
+                            x_column, y_column, color_column)]])
 
             cast_data = self._cast_datetime_axis(sliced_data, x_column)
 
@@ -402,7 +408,12 @@ def scatter(self,
             else:
                 sliced_data = data_frame[data_frame[color_column] ==
                                          color_value]
-
+            # Filter to only relevant columns.
+            sliced_data = (
+                sliced_data[
+                    [col for col in sliced_data.columns
+                        if col in (
+                            x_column, y_column, size_column, color_column)]])
             cast_data = self._cast_datetime_axis(sliced_data, x_column)
 
             source = self._named_column_data_source(
@@ -476,11 +487,16 @@ def text(self,
         for color_value, color in zip(color_values, colors):
 
             if color_column is None:  # Single series
-                sliced_data = data_frame[[x_column, y_column, text_column]]
+                sliced_data = data_frame
             else:
                 sliced_data = data_frame[data_frame[color_column] ==
                                          color_value]
-                sliced_data = sliced_data[[x_column, y_column, text_column]]
+            # Filter to only relevant columns.
+            sliced_data = (
+                sliced_data[
+                    [col for col in sliced_data.columns
+                        if col in (
+                            x_column, y_column, text_column, color_column)]])
             cast_data = self._cast_datetime_axis(sliced_data, x_column)
 
             source = self._named_column_data_source(

diff --git a/requirements_dev.txt b/requirements_dev.txt
@@ -1,5 +1,5 @@
 -e .
-pip==10.0.0
+pip==18.0.0
 wheel==0.31.1
 setuptools==40.3.0
 twine==1.11.0
@@ -10,6 +10,7 @@ coverage==4.1
 Sphinx==1.7.7
 commonmark==0.5.4
 recommonmark==0.4.0
+prompt-toolkit==1.0.15
 
 pytest==3.2.3
 pytest-cov==2.5.1