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

Key Takeaways

Here's an implementation of the traditional state design pattern using classes and inheritance. I have a payment state protocol class, so it's at least slightly modern, and it is a common interface for all the states. So, it has authorized, 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 authorized, capture, fail, and refund. And these basically call the methods on the particular state that we're in. And then, what this pattern allows us to do is create concrete states that model the actual behavior that should happen if you do these method calls. For example, for new payment, 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, authorized 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 uh basically failed is like an end state, so you can't do anything. It just raises all of these uh 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 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 you 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 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. Now, maybe that's due to this particular example, but I think there's many cases where you would have 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 we need. 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 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. That's arjan.codes/mastery. The link is also in the video description. Now, why do we need the state pattern at all? Actually, when we think of it, state logic is in many systems. So, payment flows is just one example, but you may have order processing or logistics, parsers all 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 codebases don't really uh it very well. They just add like a bunch of Boolean flags everywhere, or there's maybe illegal transitions that are slipping through. This behavior is 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 a some uh method is not possible in a particular state. The nice thing about this is that you don't have uh giant uh if-elif blocks, and behavior is separate per state. But, uh 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 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 that 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 field. 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 type to represent 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 enum. And event, we can also use an enum for that. You don't have to, but uh you could do that. And this is, by the way, 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, this is going to be a mapping from a tuple containing a state and an event to another tuple 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 So, 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 two 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 things that we have. One is the state enum, and we have an event enum. 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 enum. 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 enum. 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 enum 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 payment, 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 boilerplate code that we have here. Uh let's also remove the comment because it doesn't really add much useful information. And of course, we need to make sure that this is a field. Just the default factory of list of string, like so. And now, from the 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 states. 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 colon, 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 this does is that it appends the payment ID to the audit in the context. It doesn't change the state 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, as 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 authorized 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 encode it here by calling the add transition method. And then, as a result, we don't need any of these classes any more because they're already uh there in our state machine. So, we don't need that. So, the next part we can do, we have this payment context, which is the payment ID and the order, 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 it is 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 it's going to get a particular event that happens. 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 given the state and the 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 the state and this gets the event. And this is going to return a tuple 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.next_transition and we pass it the state and the event. Then it calls the action passing 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.handle. And then we simply supply the particular event. For example, pay event. 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 pay state.refunded. And as you can see, the audit is there. We authorized, we captured, and we refunded. And the 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, it's a bit annoying that we have to call manually all of these at transition method calls here." And you would be absolutely right. But, that's where 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 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 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 two 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, two 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. Need to type that correctly, obviously. And then, we simply can pass all the information. So, the from state, that's pay state dot new. The event is authorize. 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 it 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 refunds. 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 tuple 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 fail operation, we're going to add the decorator, like so. So, what will happen is that if we have pay state new or if we have pay state authorized, so we can then create a tuple 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 I can remove this manual one 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 as the payment 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 these 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 in action that you're seeing here. When you take a look at the final version, almost reads 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 actually also becomes really boring and simple. It just stores the current context and the current state. And 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, 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 file, there's just the business logic and specifying how we want state transitions to happen. So, this is more usable. 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 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, data-driven 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 or that contains 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 decorator-driven generic engine, or would you solve this in some other way? How How would you design a system like this? Have you used this in your own project? 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 time.