The Invent with Python Blog

Writings from the author of Automate the Boring Stuff.

Source Code Makeover: Square Shooter, Part 1

Thu 09 August 2012    Al Sweigart

UPDATE: Felix, the original author of Square Shooter, has made some comments in the comment section that are pretty insightful. I recommend you check them out.

In this blog post, I'm taking a game off of Pygame.org and going through it to make it more readable and extend its functionality. You'll see an example of how to take code that works and changes to improve it's design (but more importantly, I explain why I make those changes). This is an intermediate level tutorial and assumes some familiarity with Python. This can be pretty helpful if you know programming basics but want to know, "How can I write better code?"

(Part 2 and Part 3 are now up as well.)

This blog post covers Square Shooter, made by Felix Pleșoianu. You can download the original source or view my changes on GitHub.

I only check in working code, so you can download the code at any point in history and see how it looks. I make a lot more check ins than I normally would, but I want you to see the gradual changes that I make.

(UPDATE: A note to everyone who mentions that I break PEP-8, see this link. Also, see this link.)

Let's play it. You'll need to have Python (a version 2, not version 3) and Pygame installed on your computer to run the square-shooter_original.py file. It looks like a standard Asteroids clone. I notice that the circles dont wrap around the map until the center of the circle goes past the edge. Also, the bullets don't wrap around either, they just stop once they get to the edge.

Now take a look at the original source code. Try to get an idea of what each of the parts do, but you don't need to fully understand everything before continuing. (I didn't understand all the code when I started cleaning it up.)

First we'll do just a code clean up, and then we'll look at changing the the game to improve play and user experience.

The computer doesn't care what your code looks like. It will run the most unreadable code just fine. But the purpose of cleaning up code to be readable is so that it is easier to understand for humans so they can easily change code either for new features or bug fixes. This becomes especially important if multiple people are working on the same source code.

Also keep in mind that "readable" is a bit subjective. The changes I make here aren't necessarily the changes someone else would make. Feel free to leave your own suggestives in the blog post comments.

First I'll take a rough look through the original code. Just a note, all line numbers in this tutorial refer to the original program's line numbers, not the line numbers of my makeover program.

There are very few comments in this script (though I'm bad about this also.) Much of this makeover will be adding comments and documentation.

We'll start at the top of the program and just move down. Each time I make some changes, I constantly retest the program to make sure my changes don't cause any new bugs.

Spacing and removing duplicating variables

(These changes can be seen in the github history.)

When there is nearly identical code, I like to make the spacing match up so it easier to compare them with each other. In this case, you can see that they are all the color variables are tuples of three integers.

BLACK  = (  0,   0,   0)

GREEN  = (  0, 204,   0)

RED    = (255,   0,   0)

SILVER = (204, 204, 204)

WHITE  = (255, 255, 255)

Note that powerup_types is both a global list defined on line 35 and as a static list in the GameWorld class on line 138. But if you keep grepping for "powerup_types" you'll see it's only used once in the program. Let's get rid of this duplication by deleting lines 35 and 138, and just putting a tuple with the three strings in the one place where it is used.

If we used this list of powerups in more than one place, then putting it in a variable would make sense. But since it's only used in one place, we shouldn't take the extra step of putting it into a variable.

And since this list of three strings is never going to be modified, let's use a tuple instead of a list because tuples are immutable data types.

I'm going to go ahead and commit my changes now with the log message, "Aligned the color constants and got rid of the powerup_types variable."

Getting rid of magic numbers

(These changes can be seen in the github history.)

The scale_and_round() function can be improved. 480 is a magic number. Magic numbers are bad. The number 480 means something to the program, but as a human looking at it the number doesn't tell me anything (and there's no comments to explain them either). Using a small program called PixelRuler I can confirm that the game map (not the entire window) is 480 x 480 pixels. That's what 480 in scale_and_round() must refer to. Let's change those 480 ints to global variables MAP_WIDTH and MAP_HEIGHT.

MAP_WIDTH = 480

MAP_HEIGHT = 480



...



self.bglayer.fill(GREEN, (MAP_WIDTH, 0, 160, MAP_HEIGHT))

There are some 480 integer values in the program that don't refer to a width or height so much as a generic size. Since the map is a square, it doesn't matter. But if we ever want to make the map dimensions flexible, we should put this in a variable as well. Let's add a MAP_SIZE global constant as well, and just set it to the larger of the MAP_WIDTH and MAP_HEIGHT constants using max().

MAP_SIZE = max(MAP_WIDTH, MAP_HEIGHT)

I also notice on line 522 that the midbottom of some text for the pause screen is placed at the X coordinate 240 (in the middle of the screen's width). This is another magic number that we can get rid of. Let's replace it with a MAP_HALF_WIDTH constant (and you can see that 120 and 360 are also used, so we can add MAP_QUARTER_WIDTH and MAP_THREE_QUARTER_WIDTH constants too, as well as constants for height). We don't use all these values, but for completeness let's just add them anyway in case one day we do change the program to use them.

But what we do need is a WINDOW_WIDTH and WINDOW_HEIGHT for the pygame.display.set_mode() call. Even though WINDOW_HEIGHT and MAP_HEIGHT are both 480, they don't describe the same thing. This is why we use two separate constants for them. (Think of it like this: We would want two constants ONE_DOZEN and NUMBER_OF_DUCKLINGS even though they may both contain be 12.)

Note that the MAP_SIZE, MAP_QUARTER_WIDTH, MAP_HALF_WIDTH, and MAP_THREE_QUARTER_WIDTH are not directly set with values, but calculated from MAP_WIDTH and MAP_HEIGHT. This way, if we are playing around with these values we only have to update MAP_WIDTH and MAP_HEIGHT, and all the others are automatically updated.

model.shoot_at(x / 480.0, y / 480.0)

model.thrust_at(x / 480.0, y / 480.0)

One final thing about converting magic numbers to constants. Note that on line 560 and 561, you see the float values 480.0 instead of the integer values 480. This is because in Python 2, dividing two integer values does integer division (e.g. 22 / 7 is 3), which is too imprecise for what we want. This is why Felix used floats, which will force the / operator to do float division (e.g. 22 / 7.0 is 3.142857142857143). One of the improvements of Python 3 was to always make the / do float division.

But if we put MAP_WIDTH and MAP_HEIGHT on these lines, the integer in these constants will cause this to do integer division and cause a weird bug where the ship can only thrust and shoot towards the top left corner of the map. So we wrap them in float() calls to convert them to float values, which ensures that the / operator does float division.

model.shoot_at(x / float(MAP_WIDTH), y / float(MAP_HEIGHT))

model.thrust_at(x / float(MAP_WIDTH), y / float(MAP_HEIGHT))

Also, there are a few places where the magic number 500 is used that I will replace with MAP_WIDTH + 20. At this point, we should be able to modify the MAP_WIDTH, MAP_HEIGHT, WINDOW_WIDTH, and WINDOW_HEIGHT variables and the rest of the program scales automatically. (Though the font sizes are based on the window height, so they make look a bit awkward. We can deal with that later.)

I'm going to commit my changes now to git source control with the log message, "Replaced the magic numbers for the map and window dimensions with constant variables."

Improve scale_and_round()

(These changes can be seen in the github history.)

The scale_and_round() function can be improved again. First, we'll add a comment saying what it does. Then we can shorten it into a return statement one-liner. (The one liner is just my own preference.)

def scale_and_round(x, y):

    """Returns x and y coordinates from 0.0 to 1.0 scaled to 0 to MAP_WIDTH or MAP_HEIGHT."""

    return int(round(x * MAP_WIDTH)), int(round(y * MAP_HEIGHT))

Also, at this time I'm going to go ahead and convert all the tabs to 4 spaces with my editor's Find-and-Replace. (See github.) And then I will make some other magic number replacements also. (See github.)

Add comments to the Bubble2D class

(These changes can be seen in the github history.)

The Vector2D class looks good, and is straightforward enough that it doesn't need comments. I'll be lazy and skip them.

The Bubble2D class is used for every object in this game: asteroids, the player's ship, bullets, powerups, and explosions. Maybe "bubble" isn't the best name to use. Anyway, this fact should be made more apparent in the comments. The methods also need commenting.

class Bubble2D:

    """Represents a circular object on the game map with position, radius, and velocity."""

I'll check in these changes with the log message, "Added some comments for the Bubble2D class's methods."

Rename position to pos

(These changes can be seen in the github history.)

The Bubble2D class's position member can be shortened to pos. The word "pos" is generally accepted as a short form of "position" (or, less often, "positive"). You can see in the wrap_around() and is_out() methods that Felix creates a shorter-named pos variable to stand in the place of self.position. Let's cut this middle man out.

def __init__(self, radius):

    self.pos = Vector2D(0, 0)

Checked in with the log message, "Renamed Bubble2D's position member to pos."

Improving Bubble2D.update()

(These changes can be seen in the github history.)

There's a slight bug with wrap_around() method. Say the horizontal velocity of an object is 0.25 and the object is already at the 0.0 X position. In the next one-second update, it will be set to -0.25. If we think of the map as a toroid ring (i.e. going past the left edge moves you to the right edge) this should put the bubble at the 0.75 X position. But the current wrap_around() method will put it at 1 (the edge of the map).

The logical bug is that instead of setting the position to 0 or 1, it should actually add or subtract to the current position. I've changed the code to this:

def wrap_around(self):

    """Change the position of the bubble to toroidally "wrap around" if it goes off one edge of the map."""

    if self.pos.x < 0: self.pos.x += 1

    if self.pos.y < 0: self.pos.y += 1

    if self.pos.x > 1: self.pos.x -= 1

    if self.pos.y > 1: self.pos.y -= 1

Also, it's kind of odd that the Bubble2D object isn't smart enough to call it's own wrap_around() method after the code in update() runs. This means that the management of Bubble2D objects will have to rely on other code outside the Bubble2D class, which isn't good object-oriented design.

The reason Felix has the code set up this way is because while he wants every object to wrap except for bullets. But we can fix this. Let's have the update() method call the wrap_around() method, but return True if the object has wrapped around. That way the code that detects when to delete the bullet can know if the bullet's Bubble2D object has wrapped around or not.

Checked in with the log message, "Updated Bubble2D's update() method. How meta." (Okay, that's a bad habit. You really should keep source control log messages professional and joke-free, especially if you are working on code at your job. It's distracting, and not everyone speaks English as a first language or will get that you are just joking. Also, you might not be as funny as you think you are.)

Improve Bubble2D.is_out()

(These changes can be seen in the github history.)

The is_out() method is a pretty big one liner. If we did something like 0 < x < 1 and 0 < y < 1, that would check if the x and y was on the map. And if we put a not operator in front, the logic would be the same as the current is_out(), except a bit more readable.

return not (0 < self.pos.x < 1 and 0 < self.pos.y < 1)

Checked in with the log message, "Simplified the Bubble2D is_out() method."

Impove Bubble2D.collides_with()

(These changes can be seen in the github history.)

The collides_with() method is pretty good, but let's change the name of the bubble parameter to "other", so that we don't confuse it with the current Bubble2D object. Also, since we are squaring the a and b parameters for the Pythagorean Theorem calculation on line 85, there's no reason to get the absolute value of the differences. Even if they are negative, a negative times a negative is positive. So we can get rid of the abs() calls.

Checked in with log message, "Minor changes to the collides_with() method."

random_speed() and random_position() changes

(These changes can be seen in the github history.)

The random_position() and random_speed() functions desperately need some documentation. I've added a few comments. Better yet, instead of doing all this math on random.random() (which returns a float between 0.0 and 1.0), let's just use Python's random.uniform() to specify those values directly.

def random_position():

    """Returns a random float value that will be near the edge of the map"""

    if random.randint(0, 1) == 0:

        return random.uniform(0.0, 0.25)

    else:

        return random.uniform(0.75, 1.0)



def random_speed(magnitude):

    """Returns a random float value between -magnitude and magnitude."""

    return random.uniform(-magnitude, magnitude)

Actually, looking at it now, it doesn't make sense that the random position would be from -1.0 to 2.0, rather than 0.0 to 1.0. This might have been a bug which was masked by the fact that the wrap_around() method will correct any numbers less than 0.0 or greater than 1.0. Then again, wrap_around()'s old behavior (before our changes) would just set it at the boundary. If you look at the starting position of bubbles, you'll notice they did have a propensity for starting at the edges. We can duplicate this behavior with some of our own code, but make it more explicit.

We can also get rid of the superfluous semicolons at the end of these statements. They aren't needed in Python, and just look unpythonic.

Checking in with log message, "Simplifying random_speed() and correcting random_position(), also adding documentation."

Rename Bubble2D and simplify

(These changes can be seen in the github history.)

The make_bubble() function looks like some kind of factory function (that is, a function that is not an object's constructor function that creates and customizes objects) for the asteroids (which in the source code are called "bubbles".) Since the Bubble2D class is used not only for bubbles but also the ship, bullets, and power ups, this can get kind of confusing. Let's rename the Bubble2D class to something more generic like "ObjectOnMap" ("Object" is a bit too generic for my taste.) This way the make_bubble() function (which we should add documentation to) will be less ambiguous.

Also, since random_speed() is a one-liner and only used in the make_bubble() function, let's just replace the random_speed() calls with random.uniform() calls and get rid of the random_speed() function altogether.

Checked in with log message, "Renamed Bubble2D class to ObjectOnMap and got rid of random_speed()."

Improve make_bubble()

(These changes can be seen in the github history.)

Now let's improve the code in make_bubble(). Just to make it clear that make_bubble() is a factory function for bubbles, let's rename it to bubble_factory(). The assignment statements for the size and speed variables across line 95 and 103 (remember, these are lines 95 and 103 in the original source file) are not very pythonic. Let's use a dictionary and multiple assignment to simplify them:

#                   (size, speed)

kinds  = {'big':    (0.1,   0.1),

          'medium': (0.075, 0.15),

          'small':  (0.05,  0.25)}



size, speed = kinds[kind]

This way it'll be more obvious and easy how to add new "kinds" if we ever want, say, "extrabig" bubbles.

Checked in with log message, "Improve make_bubble(), and renamed it to bubble_factory()."

Making child classes

(These changes can be seen in the github history.)

Actually, now that I look at it, we should just make a Bubble subclass that is the child class of ObjectOnMap. The code in bubble_factory() can be moved to Bubble's __init__() function. (All bubbles, after all, are objects on the map.) While we're at it, let's replace random_position() with a single random.random() call (the player has shields at the beginning anyway.)

Note that we should probably fold the spawn_bubbles() method and GameScreen's bubble_colors member into Bubble2D. This type of info is more relevant to the Bubble objects, not the game world or any screen. We can have spawn_bubbles() be a method (simply named spawn(), since it is already in the Bubble class) that will return a two-item tuple. The first item will be a list of new Bubble objects created and the second will be a list of new Powerup items created.

While making the spawn() method, I noticed that powerups should probably also be a subclass on ObjectOnMap, so I made a Powerup class. The __init__() function for this class can mostly be taken from the spawn_powerup() method.

A minor note, it's kind of odd that spawn_bubbles() uses a variable named new_type instead of keeping consistent naming and calling it new_kind.

Also note, in order to make the super() calls of the Bubble and Powerup child classes to work, we need to change ObjectOnMap to a "new-style" class. This is done by specifying "object" as the class that ObjectOnMap inherits from:

class ObjectOnMap(object):

Wow, after all those changes, the game actually still works. Checking in these changes with the log message, "#!/usr/bin/python" (because I screwed up when pasting the log message, but it should have been "Created Bubble and Powerup child classes.")

Adding more comments

(These changes can be seen in the github history.)

Now I should add some comments while the code is still fresh in my head (and take out that debugging print statement I accidentally left in in the last check-in.) One thing to notice in my last change is that now that spawn() puts all the new Bubble objects in a list called spawned_bubbles and creates new Bubble objects in a for loop, it's trivial to have Bubbles create any number of new Bubble objects when hit. (Just for fun, try replacing for i in range(2): with for i in range(20):)

Checking in with log message, "Added some more comments to Bubble and Powerup classes."

Create a Ship class

(These changes can be seen in the github history.)

Next, I notice that a lot of things about the ship (which right now is simply a ObjectOnMap object with some extra attributes added on, or with state tracked by variables in GameState and even GameScreen) that would fit better into a Ship class.

I'll create such a class now. I can move the accel_x, accel_y, and ship_shield_timer variables into this class. I can also move the thrust_at() and stop_thrust() methods into this Ship class.

class Ship(ObjectOnMap):

    def __init__(self):

        super(Ship, self).__init__(0.04) # all Ships are the same size.

        self.shield_timer = 0

        self.accel_x = 0 # acceleration rate of the ship

        self.accel_y = 0

If we change the ObjectOnMap's init function to set objects at the XY coordinates of 0.5, 0.5 instead of 0, 0, then we can save a line of code for the ship. (The Bubble and Powerup init functions set the starting position for those types of objects, and the middle of the screen is a sensible default to have.)

def __init__(self, radius):

    self.pos = Vector2D(0.5, 0.5)

Also, the magic number 0.99 can be replaced with a constant named DECELERATION. I also notice that move_timer isn't used for anything and can be removed altogether.

DECELERATION = 0.99

And finally, I've removed some code that resets the Ship object's position and velocity on a new level, so the Ship continues to move as it did at the end of the last level. I did this accidentally, but kind of like it and will keep it in even though this goes beyond a mere code cleanup and is actually changing the way the game behaves.

Checked in with log message, "Created a Ship class and moved code into it."

Get rid of the shield_timer variable and add Ship methods instead

(These changes can be seen in the github history.)

Technically, it's bad OOP design to have outside code directly manipulating the Ship object's shield_timer member. Remember that we want to abstract away implementation details for a class. We don't want the logic of how the shield works to exist in code outside the class, because then it could be used inconsistently and there would be more places to change code if we wanted to change the behavior of the shields. Instead of increasing shield_timer by 6 and checking if shield_timer is greater than 0, let's add an add_shield() and has_shield() class.

def add_shield(self, secs=6):

    """Extends the time on the ship's shield by secs seconds."""

    self._shield_timer += secs



def has_shield(self):

    """Returns True if this ship currently has a shield, False if it does not have a shield."""

    return self._shield_timer > 0

This way, code that uses the Ship class doesn't need to know how the Ship class implements it's shield (admittedly in this case it's pretty simple since it's just a single integer value) but only needs to know about add_shield() and has_shield() methods.

On line 177 the value of shield_timer is reduced as part of GameWorld's update() method (the shield reduces over time). We don't need to create a reduce_shield() method for our Ship class, we can just have the logic for reducing the shield a part of the Ship's update() method. This design ensures that we never accidentally call the Ship object's update() but then forget to reduce the ship's shield.

In the Square Shooter game, the shields are always increased by 6 seconds. Let's add a parameter to the add_shield() method so that you can specify how many seconds to add. (This way we could add a Super Shield powerup that extends the shield by, say, 12 seconds.) But we can make this parameter have a default value of 6 so that the normal case can just call add_shield() without passing anything.

In Python, there are no private and public members or methods. Everything is public. But to imply that a member or method is private and that a programmer shouldn't directly access or call it, you can add an underscore to the beginning of the name. Let's rename shield_timer to _shield_timer.

def __init__(self):

    super(Ship, self).__init__(0.04) # all Ships are the same size.

    self._shield_timer = 0

Let's check in these changes with the log message, "Converted shield_timer member to a bunch of methods." after testing to make sure it works.

Create a Bullet class

(These changes can be seen in the github history.)

And to round everything out, let's add a new class for bullets called Bullet. Currently the bullets are just plain old ObjectOnMap objects with the shoot_at() and custom update behavior implemented outside of ObjectOnMap. Our new class will fix that.

From looking at this code, I realize that bullet_shield_timer is something that should go in the Ship class, not the Bullet class. In Square Shooter, a bullet shield lets a bullet pierce a bubble and continue instead of ceasing to exist once it hits a bubble. This is the timer for how long the Ship can fire bullets that have shields on them, not how long an individual Bullet object has a shield on it. The Bullet class will simply have a boolean shield member to say whether it has a shield or not.

It's kind of confusing for a ship to have "shields" (for the ship) and "bullet shields" (for the bullet), so I'm going to rename the bullet shield to be "Super Bullets".

def add_super_bullets(self, secs=6):

    """Extends the time on the ship's super bullets by secs seconds."""

    self._super_bullet_timer += secs



def has_super_bullets(self):

    """Returns True if this ship currently has a super bullets, False if it does not have a super bullets."""

    return self._super_bullet_timer > 0

Checking in these changes with the log comment, "Created a Bullet class."

Move shoot_at() to the Ship class

(These changes can be seen in the github history.)

Also, I'm going to move the shoot_at() method into the Ship class, since it's the ship that does the shooting. This method will return a list of Bullet objects that are created when shoot_at() is called. (For now the list will just contain one Bullet object, but this lets us be flexible in how many Bullet objects are created.)

Creating the absx and absy variables aren't necessary, because they're only used once. Let's get rid of lines 321 and 322 and just use abs(x) and abs(y) directly in the if statement on line 323.

if abs(x) < 0.1 and abs(y) < 0.1:

Because shoot_at() is now in the Ship class, we can get rid of these lines (307 and 308):

if self.bullet != None or self.ship == None:

    return

But keep in mind that this logic also prevents the player from firing a bullet if one exists on the screen, so we'll have to implement that by replacing line 560 with:

if model.bullet == None:

    model.bullet = model.ship.shoot_at(x / float(MAP_WIDTH), y / float(MAP_HEIGHT))[0]

We can probably find a better place to implement this logic later.

Checking in with log message, "Moved shoot_at() to the Ship class."

Move the freeze_timer to the Ship class

(These changes can be seen in the github history.)

At this point the only things that use just ObjectOnMap instead of a class that is derived from ObjectOnMap are the explosions. Let's hold off making a class for them right now, since an "explosion" in Square Shooter is just a circle with a position and radius. ObjectOnMap can handle that just fine.

Let's move the freeze timer logic into the Ship class, just like we did with the shield timer:

def add_freeze(self, secs=6):

    """Extends the time on the ship's freeze powerup by secs seconds."""

    self._freeze_timer += secs



def has_freeze(self):

    """Returns True if this ship currently has a freeze powerup, False if it does not have a shield."""

    return self._freeze_timer > 0

Commiting with log message, "Moved freeze timer logic into the Ship class, and some other minor changes." (This is lazy of me, I should specify what "other minor changes" is exactly. I'll probably regret this vague log message in the future.)

End of Part 1

That's all for now. I'll continue making changes in Part 2.


Learn to program for free with my books for beginners:

Sign up for my "Automate the Boring Stuff with Python" online course with this discount link.

Email | Mastodon | Twitter | Twitch | YouTube | GitHub | Blog | Patreon | LinkedIn | Personal Site