Friday 14 April 2023

Python Walrus Operator

The walrus operator aka Assignment Expressions was added to Python a few years ago, and it seems pretty interesting to me as I had never seen it before in other languages. As the name says, it allows us to have expressions that assign to a value. The main sample in the PEP-572 documentation gives us some useful use case:


# Handle a matched regex
if (match := pattern.search(data)) is not None:
    # Do something with match

# A loop that can't be trivially rewritten using 2-arg iter()
while chunk := file.read(8192):
   process(chunk)

# Reuse a value that's expensive to compute
[y := f(x), y**2, y**3]

# Share a subexpression between a comprehension filter clause and its output
filtered_data = [y for x in data if (y := f(x)) is not None]

I've started to use it in conditionals like this:


if (us := get_user()) and us.startswith("a"):
    print(f"yes, {us} starts with a")
else:
    print("no")

Given that in Kotlin almost everything is an expression I was expecting it to treat assignments as expressions also, but that's not the case. When searching about that I found out to my surprise that Java provides this feature!.

Though I'm still going through the process of interiorizing when/how to use Kotlin's scope functions and similar ones like takeIf, I've managed to come up to rewrite the conditional above like this:


    getUser()?.takeIf {
	    it.startsWith("a")
    }?.let { println("yes, ${it} starts with a") }

We can use walrus in Python when passing parameters to a function (even with named parameters):


def say(txt):
    print(txt)   

say(msg := "hi")
# hi
print(msg)
# 'hi'

say(txt=(msg := "hey"))
# hey
print(msg)
# hey

We can use it also as an indexer


items = [1,2,3,4,5,6]
pos = 0
items[pos:(pos:=pos+2)]
Out[46]: [1, 2]
pos
Out[47]: 2

Which seems pretty useful to write code like this:


def chunks_generator_fn(chunk_size, items):
    cur_pos = 0
    while cur_pos < len(items):
        next_chunk = items[cur_pos:(cur_pos := cour_pos + chunk_size]
        yield next_chunk

And one more use case:


if (index := index + 1) > end:
        index = 0

# looks better to me than:
index += 1
if index > end:
        index = 0        

No comments:

Post a Comment