Hi folks, I wish to add Polars support to the python formatting added in 9.7.

Polars is a simple and fantastic library, akin to pandas, but with a stronger
foundation in relational logic, and far fewer arcane whiums.

Is there appetite for something like this patch?  Please note, I have only
just hacked this out, and am yet to run/extend the test-suite.

Yours,
  Samuel
 
(org-babel-python--def-format-value):

  Reduce cyclomatic-complexity, and plainly state
  invariants (i.e. prefer explicit and early returns)

  Add formatting for polars.DataFrame and polars.Series.

  Add :colname if available, for pandas.Series and polars.Series.

TODO provide formatting to execution methods other than
'org-babel-python-evaluate-external-process (see
'org-babel-execute:python)
---
 lisp/ob-python.el | 122 ++++++++++++++++++++++++++++++++--------------
 1 file changed, 85 insertions(+), 37 deletions(-)

diff --git a/lisp/ob-python.el b/lisp/ob-python.el
index 1cb683492..dcee656bf 100644
--- a/lisp/ob-python.el
+++ b/lisp/ob-python.el
@@ -160,48 +160,96 @@ file to save the graphics to.")
 
 (defconst org-babel-python--def-format-value "\
 def __org_babel_python_format_value(result, result_file, result_params):
     with open(result_file, 'w') as __org_babel_python_tmpfile:
+        text = __org_babel_python_tmpfile.write
         if 'graphics' in result_params:
-            result.savefig(result_file)
-        elif 'pp' in result_params:
+            return result.savefig(result_file)
+
+        if 'pp' in result_params:
             import pprint
-            __org_babel_python_tmpfile.write(pprint.pformat(result))
-        elif 'list' in result_params and isinstance(result, dict):
-            __org_babel_python_tmpfile.write(str(['{} :: {}'.format(k, v) for 
k, v in result.items()]))
+
+            return text(pprint.pformat(result))
+        if 'list' in result_params and isinstance(result, dict):
+            return __org_babel_python_tmpfile.write(
+                str(['{} :: {}'.format(k, v) for k, v in result.items()])
+            )
+
+        if set(result_params).intersection(['scalar', 'verbatim', 'raw']):
+            return text(str(result))
+
+        if 'table' not in result_params:
+            return text(repr(result))
+
+        def dict2table(res):
+            if isinstance(res, dict):
+                return [(k, dict2table(v)) for k, v in res.items()]
+            elif isinstance(res, list) or isinstance(res, tuple):
+                return [dict2table(x) for x in res]
+            else:
+                return res
+
+        def tbl(cols: list[str], rows: list, index=False):
+            return repr(
+                [
+                    cols,
+                    None,
+                    *(
+                        [
+                            row if (
+                                isinstance(row, list) or isinstance(row, tuple)
+                            ) else [row]
+                            for row in rows
+                        ]
+                        if not index
+                        else [[i] + list(row) for i, row in rows]
+                    ),
+                ]
+            )
+
+        result = dict2table(result)
+
+        try:
+            import pandas
+        except ImportError:
+            pass
         else:
-            if not set(result_params).intersection(\
-['scalar', 'verbatim', 'raw']):
-                def dict2table(res):
-                    if isinstance(res, dict):
-                        return [(k, dict2table(v)) for k, v in res.items()]
-                    elif isinstance(res, list) or isinstance(res, tuple):
-                        return [dict2table(x) for x in res]
-                    else:
-                        return res
-                if 'table' in result_params:
-                    result = dict2table(result)
-                try:
-                    import pandas
-                except ImportError:
-                    pass
-                else:
-                    if isinstance(result, pandas.DataFrame) and 'table' in 
result_params:
-                        result = [[result.index.name or ''] + 
list(result.columns)] + \
-[None] + [[i] + list(row) for i, row in result.iterrows()]
-                    elif isinstance(result, pandas.Series) and 'table' in 
result_params:
-                        result = list(result.items())
-                try:
-                    import numpy
-                except ImportError:
-                    pass
-                else:
-                    if isinstance(result, numpy.ndarray):
-                        if 'table' in result_params:
-                            result = result.tolist()
-                        else:
-                            result = repr(result)
-            __org_babel_python_tmpfile.write(str(result))"
+            if isinstance(result, pandas.DataFrame):
+                return text(
+                    tbl(
+                        cols=[result.index.name or ''] + list(result.columns),
+                        rows=result.iterrows(),
+                        index=True,
+                    )
+                )
+
+            elif isinstance(result, pandas.Series):
+                return text(
+                    tbl(
+                        cols=[result.index.name or '', result.name or ''],
+                        rows=result.items(),
+                    )
+                )
+
+        try:
+            import polars
+        except ImportError:
+            pass
+        else:
+            if isinstance(result, polars.DataFrame):
+                return text(tbl(cols=result.columns, rows=result.rows()))
+            elif isinstance(result, polars.Series):
+                return text(tbl(cols=[result.name], rows=result.to_list()))
+
+        try:
+            import numpy
+        except ImportError:
+            pass
+        else:
+            if isinstance(result, numpy.ndarray):
+                return text(result.tolist())
+
+        return text(repr(result))"
   "Python function to format value result and save it to file.")
 
 (defun org-babel-variable-assignments:python (params)
   "Return a list of Python statements assigning the block's variables.
-- 
2.52.0

Attachment: signature.asc
Description: PGP signature

Reply via email to