A Simple Connect 4 game in Python I have made a simple command line version of connect 4. It is my first time using classes and importing one part of my project into another so any feedback on how I've done and any way it could be changed or improved would be helpful.

This code file (connect4.py) contains the basic functions I have made:

class connect4:
    def __init__(self, height=6, width=7):
        self.height = height
        self.width = width
        self.board = [['' for x in range(width)] for i in range(height)]

    def return_board(self):
        return self.board

    def get_column(self, index):
        return [i[index] for i in self.board]

    def get_row(self, index):
        return self.board[index]

    def get_diagonals(self):
        diagonals = []
        for p in range(self.height + self.width - 1):
            diagonals.append([])
            for q in range(max(p - self.height + 1, 0),
                           min(p + 1, self.height)):
                diagonals[p].append(self.board[self.height - p + q - 1][q])
        for p in range(self.height + self.width - 1):
            diagonals.append([])
            for q in range(max(p - self.height + 1, 0),
                           min(p + 1, self.height)):
                diagonals[p].append(self.board[p - q][q])
        return diagonals

    def make_move(self, team, col):
        if '' not in self.get_column(col):
            return self.board
        i = self.height - 1
        while self.board[i][col] != '':
            i -= 1
        self.board[i][col] = team
        return self.board

    def check_win(self):
        for i in range(self.height):  # check rows
            for x in range(self.width - 3):
                if self.get_row(i)[x:x + 4] in [['0', '0',
                                                 '0', '0'], ['1', '1', '1', '1']]:
                    return self.board[i][x]
        for i in range(self.width):  # check columns
            for x in range(self.height - 3):
                if self.get_column(
                        i)[x:x + 4] in [['0', '0', '0', '0'], ['1', '1', '1', '1']]:
                    return self.board[x][i]
        for i in self.get_diagonals():
            for x in range(len(i)):
                if i[x:x + 4] in [['0', '0', '0', '0'], ['1', '1', '1', '1']]:
                    return i[x]

        return None


This file is what I run and imports the connect4.py file:

import connect4

b = connect4.connect4()

while True:
    for i in b.return_board():
        print(i)
    if b.check_win() != None:
        break
    col = int(input('Team 0 choose column: '))-1
    b.make_move('0' , col)
    for i in b.return_board():
        print(i)
    if b.check_win() != None:
        break
    col = int(input('Team 1 choose column: '))-1
    b.make_move('1' , col)

print('Thank you for playing')
Docstrings

You should include a docstring at the beginning of every method, class, and module you write. This will allow any documentation to identify what your code is supposed to do. Pulling from GeeksforGeeks, this is what a docstring should be:


What should a docstring look like?


The doc string line should begin with a capital letter and end with a period.
The first line should be a short description.
If there are more lines in the documentation string, the second line should be blank, visually separating the summary from the rest of the description.
The following lines should be one or more paragraphs describing the object’s calling conventions, its side effects, etc.



ClassNamingShouldBePascalCase

Class names should normally use the CapWords convention. The naming convention for functions may be used instead in cases where the interface is documented and used primarily as a callable. Note that there is a separate convention for builtin names: most builtin names are single words (or two words run together), with the CapWords convention used only for exception names and builtin constants. [source]

In your case, you have a class named connect4. This does not conform to PEP-8 standards. It should instead be ConnectFour.

Variable Naming

You variable names should be descriptive and thoughtful enough for you to understand what that variable is holding. In your test file, you use b as the variable to hold your connect4 object. If you end up writing a larger program from this, how will you remember b is the game? What if you decide to use c, d, e, etc? Then you end up with variable names spanning the entire alphabet. This is become confusing very fast. Changing b to game makes it very clear that that's the game object.

Accessor Method Naming

Having accessor and mutator (getters and setters) methods aren't really common in python, at least not as much as in java or c#. When dealing with accessor methods, the name should be get_...(). This makes it clear that you're getting the thing from the object you're calling it on. For example, get_board() as opposed to return_board(). Not convinced? Consider this:

...
if b.get_board() is not None:
    do_stuff()
...


...
if b.return_board() is not None:
    do_stuff()
...


In this context, which is what you're using, it seems like return_board is returning a boolean, True or False, if the board was "returned successfully". get_board makes it clear that you're getting a board.

Unnecessary accessor methods

On the subject of get_board(), the whole method is unnecessary. In fact, all of the get_...() are unnecessary, but some of them are there so the code isn't all messy. Which I can appreciate. However, the get_board() is not needed. Just do game.board instead of game.get_board(). 

enumerate vs range(len())

Many python beginners, whether they be coming from other languages like C, or brand new to programming, inevitably go through a phase where they try to do the same thing.

This results in things like this:

for i in range(len(sequence)):
    print(sequence[i])


This is never the right answer.  In python we have the power and ability to loop directly over the items in a sequence.

We can instead do this:

for item in sequence:
    print(item)


It is simpler to type; simpler to read; and most importantly makes more sense.  As we were never concerned with the index to begin with, we don't need to bother with it.

There are however cases in which one does want the index as well as the item.  In such cases enumerate is the tool of choice.

Our loop becomes:

for i, item in enumerate(sequence):
    print("{} : {}".format(i,item)


In short, if you ever catch yourself writing range(len(...)) rethink your goal and figure out what it really is you are interested in. You can apply this principle to your code. 

is not vs !=

This StackOverflow answer explains this amazingly:

== is an equality test. It checks whether the right hand side and the left hand side are equal objects (according to their __eq__ or __cmp__ methods.)

is is an identity test. It checks whether the right hand side and the left hand side are the very same object. No methodcalls are done, objects can't influence the is operation.

You use is (and is not) for singletons, like None, where you don't care about objects that might want to pretend to be None or where you want to protect against objects breaking when being compared against None.

_ for unused loop variables

_ has 4 main conventional uses in Python:


To hold the result of the last executed expression(/statement) in an interactive
interpreter session. This precedent was set by the standard CPython
interpreter, and other interpreters have followed suit
For translation lookup in i18n (see the
gettext
documentation for example), as in code like:
raise forms.ValidationError(_("Please enter a correct username"))
As a general purpose "throwaway" variable name to indicate that part
of a function result is being deliberately ignored (Conceptually, it is being discarded.), as in code like:
label, has_label, _ = text.partition(':').
As part of a function definition (using either def or lambda), where
the signature is fixed (e.g. by a callback or parent class API), but
this particular function implementation doesn't need all of the
parameters, as in code like: callback = lambda _: True


(For a long time this answer only listed the first three use cases, but the fourth case came up often enough, as noted here, to be worth listing explicitly)

The latter "throwaway variable or parameter name" uses cases can conflict with the translation lookup use case, so it is necessary to avoid using _ as a throwaway variable in any code block that also uses it for i18n translation (many folks prefer a double-underscore, __, as their throwaway variable for exactly this reason).

Use i / j for double loops

This one is my opinion / possibly a common practice. When using two nested for loops, I like to use i and j as the outer and inner loop variables, respectively. i and j have typically been used as subscripts in quite a bit of math for quite some time (e.g., even in papers that predate higher-level languages, you frequently see things like "Xi,j", especially in things like a summation). Most people seem to have seen little reason to change that.

Multiple anonymous array checks

You have this code:

...
for i in range(self.height):  # check rows
    for x in range(self.width - 3):
        if self.get_row(i)[x:x + 4] in [['0', '0','0', '0'], ['1', '1', '1', '1']]:
            return self.board[i][x]
for i in range(self.width):  # check columns
    for x in range(self.height - 3):
        if self.get_column(i)[x:x + 4] in [['0', '0', '0', '0'], ['1', '1', '1', '1']]:
            return self.board[x][i]
for i in self.get_diagonals():
    for x in range(len(i)):
        if i[x:x + 4] in [['0', '0', '0', '0'], ['1', '1', '1', '1']]:
            return i[x]
...


You create the anonymous array [['0', '0', '0', '0'], ['1', '1', '1', '1']] three times! Instead of this repetition, you should assign this array to a variable, and check against that variable. The updated code reflects these changes.

'' => ' ' formatting

You have a couple of formatting problems when outputting the board to the console. When there's a column a couple stacks high in the middle of the board, it does this:

['', '', '', '', '', '', '']
['', '', '', '', '', '', '']
['', '', '', '', '', '', '']
['', '1', '', '', '', '', '']
['0', '1', '', '', '', '', '']
['0', '1', '0', '', '', '', '']


Even after just one input, it gets shifted:

['', '', '', '', '', '', '']
['', '', '', '', '', '', '']
['', '', '', '', '', '', '']
['', '', '', '', '', '', '']
['', '', '', '', '', '', '']
['0', '', '', '', '', '', '']


You should allocate this space in the beginning. In __init__, self.board should look like this:

self.board = [[' ' for x in range(width)] for i in range(height)]


and make_move(self, team, col) should now look like this:

def make_move(self, team, col):
    """
    Simulates a move and puts a 0/1 in the specified column
    """
    if ' ' not in self.get_column(col):
        return self.board
    i = self.height - 1
    while self.board[i][col] != ' ':
        i -= 1
    self.board[i][col] = team
    return self.board


Updated Code

"""
Module Docstring
A description of your program/class goes here
"""

class ConnectFour:
    """
    Class for creating a connect four game
    """
    def __init__(self, height=6, width=7):
        self.height = height
        self.width = width
        self.board = [[' ' for x in range(width)] for i in range(height)]

    def get_column(self, index):
        """
        Returns a column at the specified index

        :param index: Index at which column will be returned
        """
        return [i[index] for i in self.board]

    def get_row(self, index):
        """
        Returns a row at the specified index

        :param index: Index at which row will be returned
        """
        return self.board[index]

    def get_diagonals(self):
        """
        Returns all the diagonals in the game
        """

        diagonals = []

        for i in range(self.height + self.width - 1):
            diagonals.append([])
            for j in range(max(i - self.height + 1, 0), min(i + 1, self.height)):
                diagonals[i].append(self.board[self.height - i + j - 1][j])

        for i in range(self.height + self.width - 1):
            diagonals.append([])
            for j in range(max(i - self.height + 1, 0), min(i + 1, self.height)):
                diagonals[i].append(self.board[i - j][j])

        return diagonals

    def make_move(self, team, col):
        """
        Simulates a move and puts a 0/1 in the specified column
        """
        if ' ' not in self.get_column(col):
            return self.board
        i = self.height - 1
        while self.board[i][col] != ' ':
            i -= 1
        self.board[i][col] = team
        return self.board

    def check_win(self):
        """
        Checks self.board if either user has four in a row
        """

        four_in_a_row = [['0', '0', '0', '0'], ['1', '1', '1', '1']]

        #Check rows
        for i in range(self.height):
            for j in range(self.width - 3):
                if self.get_row(i)[j:j + 4] in four_in_a_row:
                    return self.board[i][j]

        #Check columns
        for i in range(self.width):
            for j in range(self.height - 3):
                if self.get_column(i)[j:j + 4] in four_in_a_row:
                    return self.board[j][i]

        #Check diagonals
        for i in self.get_diagonals():
            for j, _ in enumerate(i):
                if i[j:j + 4] in four_in_a_row:
                    return i[j]

        return None

def start_game():
    """
    Starts a game of ConnectFour
    """
    game = ConnectFour()

    while True:

        for i in game.board:
            print(i)
        if game.check_win() is not None:
            break

        col = int(input('Team 0 choose column: ')) - 1
        game.make_move('0', col)

        for i in game.board:
            print(i)
        if game.check_win() is not None:
            break

        col = int(input('Team 1 choose column: ')) - 1
        game.make_move('1', col)

    print('Thank you for playing')

if __name__ == '__main__':
    start_game()