feat(python):Add return_as_string arg to DF.glimpse; default=False (#…

…6678)

py-polars/polars/internals/dataframe/frame.py

@@ -2758,12 +2758,25 @@ def filter(
             .collect(no_optimization=True)
         )
 
-    def glimpse(self: DF) -> str:
+    @overload
+    def glimpse(self: DF, return_as_string: Literal[False]) -> None:
+        ...
+
+    @overload
+    def glimpse(self: DF, return_as_string: Literal[True]) -> str:
+        ...
+
+    def glimpse(self: DF, return_as_string: bool = False) -> str | None:
         """
-        Print a dense preview of the dataframe.
+        Return a dense preview of the dataframe.
 
-        Printing is done one line per column, so wide dataframes show nicely. Each
-        line will show the column name, the data type and the first few values.
+        The formatting is done one line per column, so wide dataframes show nicely.
+        Each line will show the column name, the data type and the first few values.
+
+        Parameters
+        ----------
+        return_as_string
+            If True, return as string rather than printing to stdout.
 
         See Also
         --------
@@ -2782,7 +2795,7 @@ def glimpse(self: DF) -> str:
         ...         "f": [date(2020, 1, 1), date(2021, 1, 2), date(2022, 1, 1)],
         ...     }
         ... )
-        >>> print(df.glimpse())
+        >>> df.glimpse()
         Rows: 3
         Columns: 6
         $ a  <f64> 1.0, 2.8, 3.0
@@ -2831,7 +2844,13 @@ def _parse_column(col_name: str, dtype: PolarsDataType) -> tuple[str, str, str]:
                 f" {val_str:<{min(len(val_str), max_col_values)}}\n"
             )
 
-        return output.getvalue()
+        s = output.getvalue()
+
+        if not return_as_string:
+            print(s, end=None)
+            return None
+        else:
+            return s
 
     def describe(self: DF) -> DF:
         """

py-polars/tests/unit/test_df.py

@@ -2832,7 +2832,7 @@ def test_floordiv_truediv(divop: Callable[..., Any]) -> None:
             assert divop(elem1, elem2) == df_div[i][j]
 
 
-def test_glimpse() -> None:
+def test_glimpse(capsys: Any) -> None:
     df = pl.DataFrame(
         {
             "a": [1.0, 2.8, 3.0],
@@ -2854,7 +2854,7 @@ def test_glimpse() -> None:
             "k": [["A", "a"], ["B", "b"], ["C", "c"]],
         }
     )
-    result = df.glimpse()
+    result = df.glimpse(return_as_string=True)
 
     expected = textwrap.dedent(
         """\
@@ -2875,6 +2875,11 @@ def test_glimpse() -> None:
     )
     assert result == expected
 
+    # the default is to print to the console
+    df.glimpse(return_as_string=False)
+    # remove the last newline on the capsys
+    assert capsys.readouterr().out[:-1] == expected
+
 
 def test_item() -> None:
     df = pl.DataFrame({"a": [1]})