A Pythonic Switch Statement

I came up with a perfect use case for switch statements. In my Bittorrent client, all incoming messages from peers have a message type that is referenced in the fifth byte. Knowing the message type tells me how to parse the rest of the message. I have classes for each message type that can deal with parsing the incoming bytestring, but I have to determine the message (and therefore class) type first. Looks like a great place to use a switch statement.

Only one problem: python doesn’t have switch statements. What?! But this is such a good place for one… (*sigh).

Instead, I can do nested if/else statements which would look something like this:
msg_id = bytestring[4] #the fifth byte is the message id
if msg_id = 0:
message_obj = Choke(response=bytestring)
elif msg_id = 1:
message_obj = Unchoke(response=bytestring)
elif msg_id = 2:
message_obj = Interested(response=bytestring),
elif msg_id = 3:
message_obj = NotInterested(response=bytestring)

Workable? Yes. Elegant? Not so much.

A bit of google searching turned up this much more elegant way to write my switch statement:

msg_id = bytestring[4] #the fifth byte is the message id
message_obj = {
0: lambda: Choke(response=bytestring),
1: lambda: Unchoke(response=bytestring),
2: lambda: Interested(response=bytestring),
3: lambda: NotInterested(response=bytestring),
4: lambda: Have(response=bytestring),
5: lambda: Bitfield(response=bytestring),
6: lambda: Request(response=bytestring),
7: lambda: Piece(response=bytestring),
8: lambda: Cancel(response=bytestring),
9: lambda: Port(response=bytestring),
}[msg_id]()
return message_obj

Sooo much prettier. 🙂

How does it work?
A dictionary is constructed with:
keys = all the possible values of my control variable
values = lambda functions to create the specified message object

This dictionary is called immediately after it is constructed with an input value of the fifth byte of the message (msg_id). A simpler example illustrating this is shown below.

# this:
{1:’a’, 2:’b’, 3:’c’}[3]

# is equivalent to this:
my_dict = {1:’a’, 2:’b’, 3:’c’}
my_dict[3]

Once the key is determined, a lambda function is returned as the value. If it weren’t for the ‘( )’ at the end of this code, the code would return a function object every time this was called rather than executing the function. The parentheses at the very end cause the lambda to be evaluated, and thus return the appropriate message object.

Why use lambdas at all though?
Lambda functions give us lazy evaluation, which means in this case that no message objects will be created until a lambda is evaluated. Without lambdas, the code will try to create one of each of these message object types with the bytestring provided and then return the one we ask for when we give it the msg_id. Nine of these message object creations will fail because the msg_id does not match what the object expects; these will throw exceptions. With the lambda functions, the program will create ten (lambda) functions instead of ten objects. When the dictionary is accessed with the msg_id, a single lambda function is selected. The parentheses after the msg_id cause the lambda function selected to be evaluated, yielding the message object we want. Because the other nine functions are never evaluated, they don’t throw any exceptions.

In the end, I can simply return message_obj and know that I will have an object of the appropriate type for this message.

Lambda functions also allow you to call the same function multiple times and receive different objects, stored in different memory locations, rather than references to the same memory address. This doesn’t actually matter here, as the function is only called once for a given bytestring, but this is another useful property of lambda functions for other applications.

One thought on “A Pythonic Switch Statement

  1. A great post describing how to avoid long python elif blocks! Many people consider switch statements kludgey (http://c2.com/cgi/wiki?SwitchStatementsSmell), so consensus was to leave it out of Python: http://www.python.org/dev/peps/pep-3103/

    In this case, since you’re always passing the same parameter, you can skip the lambdas altogether:

    msg_id = bytestring[4] # The fifth byte is the message id
    message_classes = {
    0: Choke,
    1: Unchoke,
    2: Interested,
    3: NotInterested,
    4: Have,
    5: Bitfield,
    6: Request,
    7: Piece,
    8: Cancel,
    9: Port,
    }
    return message_classes[msg_id](response=bytestring)

Leave a Reply

Your email address will not be published. Required fields are marked *