generalize transformer to also handle partial returns (#11)

Co-authored-by: Pavel Zwerschke <[email protected]>

README.md

@@ -42,6 +42,7 @@ This gets transformed into:
 def signum(x: pl.Expr) -> pl.Expr:
     return pl.when(x > 0).then(1).otherwise(pl.when(x < 0).then(-1).otherwise(0))
 ```
+
 ### Handling Multiple Statements
 
 polarIFy can also handle multiple statements like:
@@ -125,6 +126,38 @@ print(result)
 # └─────┴─────────┘
 ```
 
+### Displaying the transpiled polars expression
+
+You can also display the transpiled polars expression by calling the `transform_func_to_new_source` method:
+
+```python
+from polarify import transform_func_to_new_source
+
+def signum(x):
+    s = 0
+    if x > 0:
+        s = 1
+    elif x < 0:
+        s = -1
+    return s
+
+
+print(f"Original function:\n{inspect.getsource(signum)}")
+# Original function:
+# def signum(x):
+#     s = 0
+#     if x > 0:
+#         s = 1
+#     elif x < 0:
+#         s = -1
+#     return s
+print(f"Transformed function:\n{transform_func_to_new_source(signum)}")
+# Transformed function:
+# def signum_polarified(x):
+#     import polars as pl
+#     return pl.when(x > 0).then(1).otherwise(pl.when(x < 0).then(-1).otherwise(0))
+```
+
 TODO: complicated example with nested functions
 
 ## ⚙️ How It Works

pixi.toml

@@ -2,7 +2,7 @@
 # https://github.com/prefix-dev/pixi/issues/79
 [project]
 name = "polarify"
-version = "0.1.1"
+version = "0.1.2"
 description = "Simplifying conditional Polars Expressions with Python 🐍 🐻‍❄️"
 authors = ["Bela Stoyan <[email protected]>", "Pavel Zwerschke <[email protected]>"]
 channels = ["conda-forge"]

polarify/__init__.py

@@ -2,14 +2,16 @@
 import inspect
 from functools import wraps
 
-from .main import parse_body
+from .main import parse_body, transform_tree_into_expr
 
 
 def transform_func_to_new_source(func) -> str:
     source = inspect.getsource(func)
     tree = ast.parse(source)
     func_def: ast.FunctionDef = tree.body[0]  # type: ignore
-    expr = parse_body(func_def.body)
+    root_node = parse_body(func_def.body)
+
+    expr = transform_tree_into_expr(root_node)
 
     # Replace the body of the function with the parsed expr
     # Also import polars as pl since this is used in the generated code