Merge pull request #549 from jitseniesen/polars

PR: Support Polars dataframes

Author: ccordoba12

requirements/tests.txt

@@ -5,6 +5,8 @@ matplotlib
 mock
 numpy
 pandas
+polars
+pyarrow
 pytest
 pytest-cov
 scipy

setup.py

@@ -58,6 +58,8 @@ def get_version(module='spyder_kernels'):
     "mock",
     "numpy",
     "pandas",
+    "polars",
+    "pyarrow",
     "pytest",
     "pytest-cov",
     "scipy",

spyder_kernels/comms/commbase.py

@@ -108,6 +108,7 @@ def to_json(self):
             "call_id": self.call_id,
             "etype": self.etype.__name__,
             "args": self.error.args,
+            "error_name": getattr(self.error, "name", None),
             "tb": stacksummary_to_json(self.tb)
         }
 
@@ -124,6 +125,8 @@ def from_json(cls, json_data):
             type(etype, (Exception,), {})
         )
         instance.error = instance.etype(*json_data["args"])
+        if json_data["error_name"]:
+            instance.error.name = json_data["error_name"]
         instance.tb = staksummary_from_json(json_data["tb"])
         return instance
 

spyder_kernels/console/kernel.py

@@ -375,6 +375,11 @@ def get_value(self, name, encoded=False):
         """Get the value of a variable"""
         ns = self.shell._get_current_namespace()
         value = ns[name]
+
+        if str(type(value)) == "<class 'polars.dataframe.frame.DataFrame'>":
+            # Convert polars dataframes to pandas
+            value = value.to_pandas()
+
         if encoded:
             # Encode with cloudpickle
             value = cloudpickle.dumps(value)

spyder_kernels/console/tests/test_console_kernel.py

@@ -346,6 +346,17 @@ def test_get_value(kernel):
     assert kernel.get_value(name) == 124
 
 
+def test_get_value_with_polars(kernel):
+    """Test getting the value of a Polars DataFrame."""
+    import pandas
+    from pandas.testing import assert_frame_equal
+
+    command = "import polars; polars_df = polars.DataFrame({'a': [10]})"
+    asyncio.run(kernel.do_execute(command, True))
+    pandas_df = pandas.DataFrame({'a': [10]})
+    assert_frame_equal(kernel.get_value('polars_df'), pandas_df)
+
+
 def test_set_value(kernel):
     """Test setting the value of a variable."""
     name = 'a'

spyder_kernels/utils/nsview.py

@@ -380,7 +380,10 @@ def value_to_display(value, minmax=False, level=0):
                 display = '%s  Mode: %s' % (address(value), value.mode)
             else:
                 display = 'Image'
-        elif isinstance(value, pd.DataFrame):
+        elif (
+            isinstance(value, pd.DataFrame)
+            or str(type(value)) == "<class 'polars.dataframe.frame.DataFrame'>"
+        ):
             if level == 0:
                 cols = value.columns
                 cols = [str(c) for c in cols]
@@ -558,7 +561,10 @@ def get_human_readable_type(item):
             return "Image"
         else:
             text = get_type_string(item)
-            return text[text.find('.')+1:]
+            if text == 'polars.dataframe.frame.DataFrame':
+                return 'Polars DataFrame'
+            else:
+                return text[text.find('.') + 1:]
     except Exception:
         return 'Unknown'
 

spyder_kernels/utils/tests/test_nsview.py

@@ -25,9 +25,16 @@
 
 # Local imports
 from spyder_kernels.utils.nsview import (
-    sort_against, is_supported, value_to_display, get_size,
-    get_supported_types, get_type_string, get_numpy_type_string,
-    is_editable_type)
+    get_human_readable_type,
+    get_numpy_type_string,
+    get_size,
+    get_supported_types,
+    get_type_string,
+    is_editable_type,
+    is_supported,
+    sort_against,
+    value_to_display,
+)
 
 
 def generate_complex_object():
@@ -479,5 +486,13 @@ def test_get_numpy_type():
     assert get_numpy_type_string(df) == 'Unknown'
 
 
+def test_polars_dataframe():
+    import polars
+    df = polars.DataFrame({'name': ['Alice', 'Bob'], 'height': [1.77, 1.69]})
+    assert value_to_display(df) == 'Column names: name, height'
+    assert value_to_display(df, level=1) == 'Dataframe'
+    assert get_human_readable_type(df) == 'Polars DataFrame'
+
+
 if __name__ == "__main__":
     pytest.main()