Overloading the [ ] operator for indexing and slicing

The square bracket operator [ ] is used by built-in collection data types such as lists, tuples and dictionaries. The item placed inside the brackets can be a single object such as an int or string. In this case, it is used as an index into the data object. For example, a list such as myList can be indexed using the notation myList[2], which accesses the third element of the list (remember that indexes start at 0). In a dictionary, each entry consists of a key, which can be a number or a string, and an associated value. In this case, the key is the object placed inside the brackets if we wish to access the associated value in the dictionary.

With some collection types such as lists, the contents of the brackets can be a slice object, which consists of three parts: the start, stop and step. We’ve looked at some examples of slicing in an earlier post.

Like most operators in Python, the brackets can be overridden to provide indexing and slicing features in user-defined classes. To explore this, we’ll use the following class.

from math import sqrt, floor

class Primes:
    def __init__(self, max):
        self.__primes = Primes.genPrimes(max)
        self.__numPrimes = len(self.primes)
        self.__gapStats = Primes.gaps(self.primes)

    def primes(self):
        return self.__primes

    def numPrimes(self):
        return self.__numPrimes

    def gapStats(self):
        return self.__gapStats

    def genPrimes(max):
        primes = [2]
        for n in range(3, max + 1, 2):
            maxCheck = floor(sqrt(n))
            isPrime = True
            for fac in primes[1:]:
                if fac > maxCheck: break
                if n % fac == 0:
                    isPrime = False
            if isPrime:
                primes += [n]
        return primes

    def gaps(primes):
        lastPrime = 2
        gapStats = {}
        for n in primes:
            gap = n - lastPrime
            if gap in gapStats:
                gapStats[gap] += 1
                gapStats[gap] = 1
            lastPrime = n
        return gapStats

    def isint(num):
            return True
        except ValueError:
            return False

    def __getitem__(self, key):
        key = key.split()
        if len(key) > 3:  return 'Invalid entry'
        i = 0
        for n in key:
            if not Primes.isint(n): return 'Invalid entry'
            key[i] = int(key[i])
            i += 1
        if len(key) == 1:
            if key[0] in self.gapStats:
                return f'Number of gaps: {self.gapStats[key[0]]}'
                return 'Number of gaps: 0'
            key = slice(key[0], key[1], None if len(key) == 2 else key[2])
            return f'Primes in range {key}: {self.primes[key]}'
        return 'Invalid entry'

This class allows the user to build a list of all the prime numbers up to some maximum value, and also to count the number of times each gap between successive primes occurs. It’s a modification of the code we wrote earlier to illustrate generator functions. See the earlier post for an explanation of the algorithm.

The difference here is that we’ve removed the generator feature, so that when an instance of the Primes class is created, it takes the maximum value as its argument and then builds a list of the primes up to that value by calling the genPrimes() method. As this method doesn’t require access to any of the class’s instance variables, I’ve declared it to be a static method. The method is declared to be static by using the built-in @staticmethod decorator on line 21. This is just an alternative to the technique we used in the earlier post to declare a static method.

We also define a second static method gaps() on line 36 which counts the number of times each gap between successive primes occurs, and stores the result in a dictionary.

With this code out of the way, we can now look at how to overload the [ ] operator. The example here is a bit artificial, but it illustrates most of the features required. The idea is that if the user enters a single int, that value is used as a key in the dictionary containing the data on the frequency of gaps between primes. If the dictionary contains that key, we want to return the number of times that gap was found; if it doesn’t contain the key, we should return 0.

If the user enters 2 or 3 values as the key, this is to be interpreted as a slice from the list of primes, and the corresponding slice should be returned.

The [ ] operator (for reading the values of one or more items) is implemented in a class by calling the class’s __getitem__() built-in method. To overload the [ ] operator, we therefore need to write our own __getitem__() method in our class.

This is done starting on line 57. __getitem__() requires the usual self argument, and one additional argument which is the key, which is the object inside the brackets. In principle, the key can be any data type, but the two most common types are a single integer (for indexing) and a slice object (for accessing a range of objects within the collection). The example here illustrates both.

This example assumes that the user will enter the key as a string of one, two or three ints separated by blanks. The first thing we need to do is split the input into a list of the items, which is done on line 58. As we’re allowing up to 3 elements in a key, we check that the list is no longer than 3 items on line 59.

Next, we need to check that all the entries in the key are valid integers, and if so, to convert them from strings to ints. In a more mature program, we would do this by using exceptions, but to keep things simple here, I’ll just return the message ‘Invalid entry’ if something goes wrong.

As far as I can tell, Python has no built-in function for testing if a string contains a valid int (there is the isdigit() function, but this won’t accept negative numbers). Thus I’ve included a little static method on line 49 which checks if a string num is a valid int. It does this by attempting to cast num from a string to an int. If this fails, a ValueError is raised which is used to return False.

Using isint() we can verify that the components of key are valid ints with the loop on line 61. If each element of key is an int, we cast it from a string to an int so it’s ready to be used in our overload of [ ].

On line 65, we deal with the case where the key consists of a single int. This is to be used as a key in the gapStats dictionary, so we check to see if the dictionary contains that key (line 66) and if so, we return the corresponding value. If it doesn’t, we return 0 (lineĀ  69).

The else on line 70 deals with the cases where key contains 2 or 3 elements, and thus is to be interpreted as a slice. We use the built-in slice() function on line 71 to create the slice object, and then just pass this to the primes list on line 72. Note that we don’t try to implement the inner workings of a slice; we just pass the slice along to the list and let the code in the built-in list class deal with the slice. The behaviour of a slice is fairly complicated, as it has to deal with 2 or 3 components, negative integers, variable step sizes and so on, so attempting to duplicate all this wouldn’t be easy. Usually when we overload the [ ] operator in our own class, that class contains some built-in collection object which can implement slicing for us. We just need to pass a valid slice object to that object.

We can test the code using a main program such as this:

from Primes import *

pList = Primes(200)

print(f'Number of primes = {pList.numPrimes}')

while True:
    key = input('Enter single int for numGaps or slice for prime range: ')
    if key == 'quit': break

We print out the list of primes and the dictionary of gap frequencies, and then enter a loop asking the user to enter the elements of a key. A typical session might look like this:

[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199]
Number of primes = 46
{0: 1, 1: 1, 2: 15, 4: 13, 6: 12, 8: 1, 14: 1, 10: 2}
Enter single int for numGaps or slice for prime range: 14
Number of gaps: 1
Enter single int for numGaps or slice for prime range: 5 -1
Primes in range slice(5, -1, None): [13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197]
Enter single int for numGaps or slice for prime range: 5 20 3
Primes in range slice(5, 20, 3): [13, 23, 37, 47, 61]
Enter single int for numGaps or slice for prime range: quit
Press any key to continue . . .

We see that if we enter a single int (such as 14 in the example), we get the message ‘Number of gaps: 1’ as the reply, indicating there is one gap of 14 in the prime list.

Entering 5 -1 results in the slice [5:-1] which means that the sixth through to the second-to-last element in the primes list should be selected, which we can see printed out.

Entering 5 20 3 means that we obtain the slice of the primes list from element [5] to element [19] in steps of 3, which again checks out.

Leave a Reply

This site uses Akismet to reduce spam. Learn how your comment data is processed.