The State Pattern in Python - I Like How This Turned Out

Here's an implementation of the traditional state design pattern using classes and inheritance. I have a payment state protocol class. So at least slightly modern and it is the common interface for all the states. So it has authorize, capture, fail and refund. Then I have a payment data class which has a payment ID. It has a list of audits like a log and there is the current state which is a payment state. So that's actually an instance of uh this class following this protocol. And then I have different methods like authorize, capture, fail, and refund. And these basically call the methods on the particular state that we're in. And then what this pattern allows to do is create concrete states that model the actual behavior that should happen if you do these uh method calls, for example, for new payments. So we can authorize it. Uh but we can't capture it. It can fail. Uh we can't refund it because there is no payment yet. Uh then we have authorized. So this has different behavior. So authorize again is not possible because it's already authorized. Uh we can capture it. Uh it can fail and can be refunded. And then we have failed. So this has again basically failed is like an end state. So you can't do anything. It just raises all of these runtime errors. There's one more refunded. Uh that's basically the same as failed. So you can't really do anything after that as well. Uh so then in the main function I create a payment. I set it to the state new and then I call a couple of uh methods here like authorizing it, capturing it, refunding it and then we get a final state and we get an audit. So when I run this then this is what we get. The final state is refunded and here's the audit. And what this then allows to do, let's say I do another refund after I've already refunded this payment, then you see that we will get an error, namely that uh this was already refunded. So that's the traditional state pattern. That's what it looks like. But I don't really like it though. I think this hinges too much on inheritance for my taste. And there's a lot of duplicate code. And maybe that's due to this particular example. But I think there's many cases where uh you would have uh certain methods in these states that actually have the same behavior. And if you model it like this, you're going to duplicate all of that or you're going to need to add more inheritance. Now, I'm going to build this in a very different way using fancy Python features obviously because this is a Python channel. And by the end of this video, you'll have a very clean engine that you can drop into any of your own projects that need some sort of state machine, which is like really neat. Now, what I've learned over the years is that with patterns like this, it's not just important that you know they exist, but that you also understand the design trade-offs behind them. That's exactly what I go deep into software design mastery. This is not just another course with patterns and principles. It's a comprehensive system for thinking about architecture, boundaries, trade-offs, and long-term maintainability. If you want to become truly confident in your design decisions, the program is opening this year. Join the waiting list so you'll be the first to know at ar. go/mastery. The link is also in the video description. Now, why do we need the

state pattern at all? Actually, when you think of it, state logic is in many systems. So, payment flows is just one example, but you may have order processing or logistics. Uh parsers will need states. Automations in general, if you're building an automation system, it's going to have different states. And depending on the state you're in, you want to do a different task. And unfortunately, most code bases don't really uh handle it very well. They just add like a bunch of boolean flags everywhere or there's maybe illegal transitions that are slipping through as behavior scattered across modules. And the state pattern basically tries to solve this by making these transitions explicit, encapsulating the behavior per state and preventing invalid flows. Now

I've already shown you this traditional class-based state pattern that we have here. Each state class implements the same methods, and some of these methods transition the payment to a different state. For example, captured, if you do refund, it changes the payment state to refunded. It can also raise an error if some uh method is not possible in a particular state. The nice thing about this is that you don't have giants uh if blocks and behavior is separated by state. But like I said, it also results in a lot of duplication, lots of these tiny classes, many raise runtime error methods, and the transitions are kind of scattered across the different files and classes that you see here. So, it's hard to see what the actual state machine actually looks like in one

place. When you zoom out and think about what this actually is, actually a state machine is simply a lookup table going from the current state and some event to the next state with a particular action. And that's just data. And the class-based pattern that we have here hides that data structure inside polymorphism. But in Python, we can make that structure explicit. So instead of

doing this, we can create a state machine class and then have a couple of types that it's going to use. So in order to do that, I'm going to create a state machine file where I'm going to define this class. And then let's try to make this a bit generic. So from data classes, I'm going to import uh data class and I'm also going to use fields. So we'll have a class state machine. And what I'd like to do is make this generic. And in essence, a state machine needs three things. It needs some type to represent states. It needs to know a type to represent events. And it needs a context. Like in our case, for example, we have a payment object. And the state machine needs to know about that. So we have these three things. So we have a state type, we have an event type, and we have a context type. And we can make this a bit more specific by saying that well state that should be an enome and event we can also use an enome for that. You don't have to but uh you could do that. And this is by the way a nice feature of generics where you can specify these types of constraints in the generic types. So we have a state machine and actually I'm going to use a data class for this and then we're going to have the uh transitions and this is going to be a dictionary and like I said mapping from a tupole containing a state and an event to another tupil containing the next state and then the action that needs to be taken and the action is a function. that should be called that gets the context. So let's also define that. So this is a callable which gets the context and returns let's say none and we import that from the typing module. So that action we're going to include here. So transitions is a dictionary that maps a tuple of the current state and an event to the next state and an action function. And of course action is a generic type. So we also need to indicate that here. And then this is going to be a field where the default factory is going to be the dictionary type that we've defined right here. Like so. In essence, this is the only data structure that we need. But we can add a few convenient methods. Uh for example, we could have an add transition method. And this gets a from state which is of type s. It gets an event. It gets a to state and it gets an action function like so. And the only thing this is going to do is add it to the transitions. So given the from state and event tuple, we're going to store the two state and the action function like so. And now let's go back to our main file and start using this state machine. So in order to do that, we're going to need a couple of types to represent the different uh things that we have. one is the state enome and we have an event enome. So from enum I'm going to import the enum type and auto which is a helpful thing and then we're going to have a class pay state which will be an enome and this is going to have the various states. So we have new and we have authorized captured failed and refunded. These are the different states that we have. But we also have an event which is also an email. So these are the four types of occurrences that we can have. An authorized event, a capture event, a fail event, and a refund event. And we can simply create enom values for these four different types of events. I'm going to remove this payment state protocol because we're no longer going to use that. The payments, that's going to be the context. Maybe we should also rename it to payment context. And this is going to have the payment ID and the audit. And it's not going to have the state anymore. So we remove all of this uh boiler plate code that we have here. Let's also remove the comment because it doesn't really add much useful information. And of course uh we need to make sure that this is a field with the default factory of list of string like so. And now from the uh state machine module I'm going to import the state machine and then I'm going to create that here. So that's the pay state machine and that's going to be a state machine that will have as types the pay state. It's going to have a pay event and it will have the payment context and we simply initialize that to an empty object and that should be a column like so. So now we have our state machine. Now what we can do is define the transition actions and we can remove a lot of these duplications. So for example we can have an authorize function. So I won't need self uh this is going to get the context. Let's just call it like that. And the only thing uh this does is that it appends the payment ID to the audit in the context. It doesn't change the stick because that's what the state machine is going to do for us. And similarly, we're going to have capture, fail, and refund like so. So, we now have these four functions. And since we now created the state machine, we can actually store the transitions. So for example, one thing you could do is say well the pay state machine add transition and we can make a transition from pay state dot new. So if that's the current state and we have a pay event dot authorize then the next state is going to be authorized and the action is the authorize function like so. So now we've added this transition and we can do the same thing for all of the other transitions as well. So for each possible transition we basically encoded here by calling the add transition method. And then as a result we don't need any of these classes anymore because they're already uh there in our state machine. So we don't need that. So the next what we can do we have this payment context which is the payment ID and the audit the data related to the payment. Well, then we can still have a payment class and this is going to have the context and this is where we will store the state. And let's say initially that is going to be a new payment like so. And now what we want to do is add a method here called handle. And this is going to get an event. So what this method then will do is use our state machine to handle the occurrence of this event given the current state. Now in order to do that on for the time being I'll just write pass here. We still need to add a bit of extra behavior to our state machine because currently this only represents transitions but we actually want to handle a particular event. So in order to do that, let's add a few methods to state machine to support that. So let's say we will have a handle method and this is going to get some context. It's going to get the current state and a particular event that happened. And what this will do is it will return the next state. So what we now need to do is that given the state and the event what we now need to do is look up in our transitions dictionary uh given this state and an event what the next state is going to be and what action we should execute. So to make that a bit easier let's make a next transition method that does that work for us. So this gets a state and this gets an event. And this is going to return a tupole containing the next state and an action like so. And what this is going to do, I'll just paste it here is that it will try to return a particular transition from dictionary. If not, it's going to raise an invalid transition. And let's add that as an exception class here. And I'll leave this empty for the time being. Now that we have the next transition method, the handle method is actually really easy to build. The only thing this needs is a next state and an action and it gets that from self. ext transition and we pass it the state and the event. Then it calls the action pass in the context and it's going to return the next state. And now in our main file this becomes really easy. So we now have all of these transitions, we have the functions, we have the state machine, we have the context. So now we can create the payment object and we pass it a context which is a payment context with payment ID P1 like so. And the state is already set. And now instead of calling these methods directly, what we can do is p. andle. And then we simply supply the particular event. For example, payvent dot authorize. And the only thing of course we still need to do here is because handle is empty. So we need to call the method from the state machine. This now has a handle method. We're going to pass the context, the current state, and the event that occurred. And the result we're going to store as the new state. That's the only thing we need to do here. So now we can handle authorize. Uh we can handle other events like capture and we can also handle a refund like so and then we can simply print what the final state is and we can also print the audit like so. So when we run this code then this is what we get. So the final state is payday. refunded refunded. And as you can see, the audit is there. We authorized, we captured, and we refunded. And a nice thing about this is that now we don't have all these duplicated functions because we just have the ones that are important. Authorize, capture, fail, and refund. And we simply explicitly define what the transitions are. And that is happening in this part. So this is a much in my opinion cleaner way of doing it than having all of these different classes with all of these duplicated methods.

Now, you might say, well, this is a bit annoying that we have to call manually all of these add transition method calls here. And you would be absolutely right. But that's when we can rely on some fancy Python features because instead of doing this manually like I'm doing here, we can also define a decorator that does that job for us. And then we can basically write the decorator above each of these uh functions that we have here. And it's going to be even better. So the decorator we are going to add to our state machine and let's say I'm going to call that decorator transition and this is going to transition uh from a from state which is of type s uh we're going to have some event of type e and we're going to have a to state which is also of type s. And this decorator is going to be put on top of the action function. So we don't need to pass it here as an argument. So what we'll do is define the decorator and that gets the action function as an argument and it will also return the action function. So what this will do is self add transition s event to state and the function like so. And then we're going to return the decorator and this should be from state like so. And then finally the decorator should return the action function like so. So this is our decorator and then instead of doing this manual thing that we have here. So for example for this first one we have a new pay state and the event is authorize then the pay state needs to be authorized and we need to call authorize. So what you can do then is go here and then we're going to do pay state machine dot transition. I need to type that correctly obviously and then we simply can pass all the information. So the from state that's past state new the event is authorized and the two state is authorized like so. And then we can simply do this. And as you can see this is way shorter. And now that we made this change let's just make sure this still works. And this seems to still work, which is nice. So now, instead of doing all of these manual transitions that we write here, we can simply write them as decorators above our functions. Now, one more thing I want to show you, but before I do, if you enjoy deep dives into patterns, architecture, and clean design, you want more content like this, give this video a like and subscribe to the channel. It helps me a lot. Now, there's one

problem. If you take a look at these transitions, uh there's multiple transitions that happen with a single function like fail for example. It can go from the new pay state. Uh it can also come from the authorized pay state and uh same for refund. So refund can also come after authorized but also after captured. And the problem is currently we can't express that cleanly because this just gets a single state. But fortunately there's a very easy fix to this. We can actually let the transition method here, the transition decorator to accept not only a single state but also an iterable and that means we can now supply it with a tupole or a list of different from states. The only thing we need to do is of course handle that here. So what we can do is if not is instance from state iterable. So if it's a single state then we are going to turn it into an iterable by turning it into a tuple like so. And then here the only thing we need to do is create a for loop because for each of these from states we're going to need to add a transition like so. And now what we can do in our main function is for the uh fail operation we're going to add the decorator like so. So what will happen is that if we have pay state new or authorized. So we can then create a tupil here like so. The event is fail and then the next state is going to be failed like so. And now I can remove this manual one and as well. Let's just make sure this still works. It's still refunded. And to test, let's say we uh are going to also handle a fail event. And then of course this doesn't happen. So when we run it like so, we see now is the base is failed. We authorized and we get failed in the audit. So that works. And of course, as a final step, you can also take these manual calls to add transition and turn them into a decorator as well to make it really simple. Here's the final version of that code. So, we have all of these nice decorators on top of each of the states. So, what is the key design when here? Well, if we want to add a new transition, we simply write a new function. We don't modify existing code. It's really the open closed principle and action that you're seeing here. When you take a look at the final version almost like a spec. We have transitions from new and authorized to fail. And then this is the action that's going to be execute. We have a transition from new and to authorized if this event occurs. And then this action is going to be called. So this is really easy to see. Now you can review the entire state machine just by scrolling through the decorators. Another thing that's nice is

that the payment class uh actually also becomes really boring and simple. It just stores the current uh context and the current state. Now you might wonder why not store the state in the state machine. Actually the nice thing about doing it like this is that you can have multiple payments and you use the same state machine which means that you can keep using the decorator in this way. Otherwise, it's going to be a mess because then with the decorator is associated to this particular object, but then you may have multiple payments. Uh it's not going to work. So, having the state being part of the payment object actually makes more sense here. And the cool thing is we now have this generic state machine that you can apply in all sorts of different places. And in the main fold, it's just the business logic and specifying how we want state transitions to happen. So, this is more reusable. You can create multiple independent state machines and each instance is going to be isolated. It's easy to test because we can test the state machine engine separately from the rest of the code. We can test the transitions uh separately from the rest of the code. It's easy to extend because we can add new states without modifying any existing classes. We can add transitions just by writing a few functions. There's no inheritance. This in my opinion is Python relying on its strengths. first class functions, generics, decorators, datadriven design.

Now, there are a few cases where you might want to stick with the original class-based approach that I showed you. For example, if each state has very complex internal behavior that relies on data stored in that state. Uh, or maybe you have state specific methods that contain significant pieces of business logic. In most other cases though, my go-to is to use a state machine pattern like I showed you today.

But I'd like to hear from you. Do you prefer the classic class-based approach I showed you in the beginning or this decorated driven generic engine? Or would you solve this in some other way? How would you design a system like this? Have you used this in your own projects? Let me know in the comments. Now, if you want to explore more software design patterns like this one, check out my design patterns playlist right here. Thanks for watching and see you next