with_walrus

with_walrus.py

import sys
import polars as pl
"""
This method is a work around for two annoyances.

1. If one wants to reuse an earlier column definition then the 
walrus operator can be used but then it looks awkward because it
has the python variable on the left but then still needs an 
alias at the end such as
`with_columns(a:=(pl.col('b')+1).alias('a'), (a*2).alias('c'))`

2. It mimics the kwargs feature without requiring all variables 
after the first kwarg to still be kwargs. For instance, this doesn't work
```
with_columns(a=pl.col('b')+1, pl.exclude('b')-1)
```

Instead the usage is all *args like this
```
with_walrus(a:=pl.col('b')+1, pl.exclude('b')-1, z:=a*2)
```

This doesn't come without its own cons. The way it works is to look at the last n
local variables where n is the number of args. It then compares the contents of those
variables to all the args. When the arg matches the local variable, it assigns that
name to that arg. It then deletes that variable from the local space. 
"""
def with_walrus(self, *args):
    globals = sys._getframe(1).f_locals
    recents = list(globals)[::-1][0 : len(args)][::-1]
    new_args = []
    for arg in args:
        matched_walrus = False
        for i, recent in enumerate(recents):
            arg_match = arg is globals[recent]
            if arg_match:
                matched_walrus = True
                if isinstance(arg, str):
                    arg = pl.col(arg)
                elif not isinstance(arg, pl.Expr):
                    arg = pl.lit(arg)
                arg = arg.alias(recent)
                _ = recents.pop(i)
                new_args.append(arg)
                del sys._getframe(1).f_locals[recent]
                break
        if not matched_walrus:
            new_args.append(arg)
    return self.with_columns(*new_args)
pl.DataFrame.with_walrus=with_walrus


### Example usage
df=pl.DataFrame({'a':[1,2,3], 'b':[2,3,4]})

df.with_walrus(pl.all()-1, c:=pl.col('a')+pl.col('b'), d:=c+5)