Update main.py

[badr-elmazaz]

made it more pythonic

[carlospolop]

Hi @badr-elmazaz, does this work? removing all the warriors doesn't look like it will

[badr-elmazaz]

Hi @carlospolop sorry, I had some notification issues that now I have solved, I just checked your reply. Actually it does not work. My real solution is to use a factory pattern here instead of using an elif structure, what do you think about it? It will be more robust and less timeconsuming when we will add new class. Let me know what you think about it, in this way I can contribuite with this idea.