Roll Your Own Power Function in Python

As you might already know, Python has the two ways to compute power: the power operator, `**`, and the power function, `pow()`. Today, we’re going to try to replicate the behavior of the power function.

Problem Description

For today’s challenge, we’re going to implement our own power function equivalent to the `pow()` function. As trivial as this sounds (after all, `**` exists), I don’t want to take any shortcuts today.

To start, let’s talk about the function definition:

Return base to the power exp; if mod is present, return base to the power exp, modulo mod (computed more efficiently than `pow(base, exp) % mod`). The two-argument form `pow(base, exp)` is equivalent to using the power operator: `base**exp`.

The arguments must have numeric types. With mixed operand types, the coercion rules for binary arithmetic operators apply. For `int` operands, the result has the same type as the operands (after coercion) unless the second argument is negative; in that case, all arguments are converted to float and a float result is delivered. For example, `10**2` returns `100`, but `10**-2` returns `0.01`.

For `int` operands base and exp, if mod is present, mod must also be of integer type and mod must be nonzero. If mod is present and exp is negative, base must be relatively prime to mod. In that case, `pow(inv_base, -exp, mod)` is returned, where inv_base is an inverse to base modulo mod.

Source: Python Docs

Based on this description, it’s clear we’re not going to be implementing a straightforward power function. In fact, we’re actually implementing power mod, a handy function for cryptography.

However, there are a lot of nasty issues we can run into while trying to implement this solution, so I figured we could set some reasonable bounds.

First, let’s assume that all exponents are integers. As it turns out, it’s quite straightforward to compute power given these constraints because we don’t have to deal with fractional exponents.

Second, let’s not worry too much about performance. I’m aware that there are tricks to computing power that save on computation, but we’re going to stick with a linear solution.

Finally, I’m going to ignore that last paragraph about relative primes. If it turns out this is too easy, I’ll circle back and try to implement it. Otherwise, let’s get going!

Thought Process

Based on the problem description, our own power function will need to support three inputs:

• Base
• Exponent
• Mod (Optional)

Here’s what that looks like as a function header (feel free to throw this in a file called `roll_your_own.py`):

```def power(base, exp, mod=None):
pass```

From there, we need to figure out how to compute power given these values.

My first though was to leverage a quick and dirty solution which treats the calculation as a form of repeated multiplication. For example, if we had 2 to the 4th power, we could compute the result by repeatedly multiplying 2 until we’ve done that 4 times. Here’s what that might look like as pseudocode:

1. Create variable to store result
2. Loop from 0 to exponent
1. On each iteration, multiply the result by the base
3. Return result

Unfortunately, there are a couple things that we have to consider. What happens if we have a negative power? In that case, we’ll need to invert the result. Here’s what the algorithm looks like with this change:

1. Create variable to store result
2. Loop from 0 to exponent
1. On each iteration, multiply the result by the base
3. Check if exponent is negative
1. True: return inverted result
2. False: return result

Likewise, we have to consider things like whether or not the user provides the `mod` argument. As a result, we’ll need to include a branch based on the status of `mod`. That said, I suspect various issues like this to pop up when we get to writing our code. For now, what we have is a nice start. If we need to come back and modify our algorithm, we can. In the meantime, let’s move on to testing.

Testing

As far as I can tell, this function is pretty straightforward to test. As a result, I think our usual “first, middle, last, zero, one, many” routine should do the trick.

• First: N/A (not dealing with any data structures)
• Middle: N/A (see above)
• Last: N/A (see above)
• Zero: raise to the power of zero; raise a base of zero to any power
• One: raise to the power of one; raise a base of one to any power
• Many: the usual case

In addition, folks sometimes like to test errors. One of the possible errors here is modding by zero:

• Oops: mod by zero

With these basic test cases out of the way, let’s get to writing them in code. Before we do that, however, I wanted to acknowledge that there’s likely a better way to do testing than we’ve been doing before. After all, since we’re trying to replicate an existing function, we might as well run the same inputs through both and verify that the results are the same. Here’s what that looks like:

```import unittest
import importlib

roll_your_own = importlib.import_module("roll_your_own")

class TestPower(unittest.TestCase):

def assert_helper(self, base, exp, mod):
self.assertEqual(
roll_your_own.power(base, exp),
pow(base, exp),
f"Failed to compute {base}^{exp}"
)
self.assertEqual(
roll_your_own.power(base, exp, mod),
pow(base, exp, mod),
f"Failed to compute {base}^{exp} % {mod}"
)

def test_power_zero_exponent(self):
self.assert_helper(5, 0, 7)

def test_power_zero_base(self):
self.assert_helper(0, 5, 7)

def test_power_zero_both(self):
self.assert_helper(0, 0, 7)

def test_power_one_exponent(self):
self.assert_helper(5, 1, 7)

def test_power_one_base(self):
self.assert_helper(1, 5, 7)

def test_power_one_both(self):
self.assert_helper(1, 1, 7)

def test_power_many(self):
self.assert_helper(3, 5, 7)

def test_power_oops(self):
self.assertRaises(ValueError, roll_your_own.power, 4, 4, 0)

if __name__ == '__main__':
unittest.main()```

Also, this time around, I included a bit of a helper method to cut down on code a bit.

Solution

As a reminder, here’s the function header we’ll be using:

```def power(base, exp, mod=None):
pass```

As long as our tests are in good shape, we should see the following when executing this empty method:

```FFFFFFFF
======================================================================
FAIL: test_power_many (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(3, 5, 7)
self.assertEqual(
AssertionError: None != 243 : Failed to compute 3^5

======================================================================
FAIL: test_power_one_base (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(1, 5, 7)
self.assertEqual(
AssertionError: None != 1 : Failed to compute 1^5

======================================================================
FAIL: test_power_one_both (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(1, 1, 7)
self.assertEqual(
AssertionError: None != 1 : Failed to compute 1^1

======================================================================
FAIL: test_power_one_exponent (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(5, 1, 7)
self.assertEqual(
AssertionError: None != 5 : Failed to compute 5^1

======================================================================
FAIL: test_power_oops (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assertRaises(ValueError, roll_your_own.power, 4, 4, 0)
AssertionError: ValueError not raised by power

======================================================================
FAIL: test_power_zero_base (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(0, 5, 7)
self.assertEqual(
AssertionError: None != 0 : Failed to compute 0^5

======================================================================
FAIL: test_power_zero_both (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(0, 0, 7)
self.assertEqual(
AssertionError: None != 1 : Failed to compute 0^0

======================================================================
FAIL: test_power_zero_exponent (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(5, 0, 7)
self.assertEqual(
AssertionError: None != 1 : Failed to compute 5^0

----------------------------------------------------------------------
Ran 8 tests in 0.068s

FAILED (failures=8)```

To get the results that we want, we’ll need to implement our function properly. Here’s the steps I took:

Perform Repeated Multiplication

As I mentioned before, one of the quickest ways to perform power is to treat it like repeated multiplication. To do that, we can use a loop:

```def power(base, exp, mod=None):
result = 1
for i in range(exp):
result *= base
return result```

The way this works is we create a variable that holds a value of one. The reason for this is that we’re going to perform multiplication over and over again. As a result, the initial value cannot be zero or it would cancel out any product we try to compute.

Also, one happens to be the agreed upon value for our edge case where the exponent is zero. That makes it very easy to return before we do any work.

At any rate, I call this the quick and dirty solution because it only works for integers, and it’s not the fastest solution. That said, it’s pretty easy to read, and it gets the job done.

Now if we run our tests, we should get some different results:

```F...F...
======================================================================
FAIL: test_power_many (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(3, 5, 7)
self.assertEqual(
AssertionError: 243 != 5 : Failed to compute 3^5 % 7

======================================================================
FAIL: test_power_oops (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assertRaises(ValueError, roll_your_own.power, 4, 4, 0)
AssertionError: ValueError not raised by power

----------------------------------------------------------------------
Ran 8 tests in 0.011s

FAILED (failures=2)```

And just like that, six of our tests passed! Now, these results are a bit misleading because all of the mod tests just so happen to work out. If we were a bit more careful in the creation of our test cases, we would probably still fail all eight of them (or at least six of them).

That said, our “many” test case caught the mod issue, so let’s modify our code to support it.

Now, mod is a bit tricky to incorporate because it’s an optional parameter. As a result, we have to make sure it exists before we try to apply it. One quick way we could do that is to take the final result and apply mod if and only if the argument exists:

```def power(base, exp, mod=None):
result = 1
for i in range(exp):
result *= base
if mod:
result %= mod
return result```

Again, this is sort of quick and dirty because it would be preferable to apply the mod after every multiplication. That way, we don’t let our integers grow too large. That said, for our purposes, I think this is a nice comprise. Let’s see how it stacks up to testing:

```....F...
======================================================================
FAIL: test_power_oops (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assertRaises(ValueError, roll_your_own.power, 4, 4, 0)
AssertionError: ValueError not raised by power

----------------------------------------------------------------------
Ran 8 tests in 0.055s

FAILED (failures=1)```

Awesome! We didn’t break any of our tests, and we even fixed the mod issue. Now, all that’s left is to fix up this ValueError.

Throw Appropriate Errors

Personally, I’m not a huge fan of exceptions. That said, if we’re going to try to replicate power as close as possible, we’re going to need to throw errors when appropriate. Fortunately, this error is pretty easy to raise:

```def power(base, exp, mod=None):
if mod == 0:
raise ValueError("power() 3rd argument cannot be 0")
result = 1
for i in range(exp):
result *= base
if mod:
result %= mod
return result```

In other words, if `mod` is zero, we can throw the ValueError. Otherwise, we compute power as usual.

Upon completion, we end up with the following test results.

```........
----------------------------------------------------------------------
Ran 8 tests in 0.069s

OK```

Despite these satisfying results, I’m not sure we fully implemented power. In the next section, we’ll take a look at tying up some loose ends.

Tying Up Loose Ends

After implementing power up to this point, I realized that I failed to consider two possible cases: negative bases and negative exponents. As a result, I’ve updated the test suite to include both cases:

```def assert_helper(self, base, exp, mod):
# 2 argument test
self.assertEqual(
roll_your_own.power(base, exp),
pow(base, exp),
f"Failed to compute {base}^{exp}"
)
# 3 argument test
self.assertEqual(
roll_your_own.power(base, exp, mod),
pow(base, exp, mod),
f"Failed to compute {base}^{exp} % {mod}"
)
# negative base test
self.assertEqual(
roll_your_own.power(-base, exp),
pow(-base, exp),
f"Failed to compute -{base}^{exp}"
)
# negative exponent test
if base != 0:
self.assertEqual(
roll_your_own.power(base, -exp),
pow(base, -exp),
f"Failed to compute {base}^-{exp}"
) ```

Now, whenever we write a test case using our helper, we should test each combination four different times:

• Once for the 2 argument version
• Once for the 3 argument version
• Once for negative base
• Once for negative exponents

Once these tests are in place, we get the following result using our current solution:

```F..F....
======================================================================
FAIL: test_power_many (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(3, 5, 7)
self.assertEqual(
AssertionError: 1 != 0.00411522633744856 : Failed to compute 3^-5

======================================================================
FAIL: test_power_one_exponent (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(5, 1, 7)
self.assertEqual(
AssertionError: 1 != 0.2 : Failed to compute 5^-1

----------------------------------------------------------------------
Ran 8 tests in 0.067s

FAILED (failures=2)```

Clearly, the only issue we’re running into is negative exponents. Fortunately, this is also a quick fix. I chose to apply absolute value to the range calculation first:

```def power(base, exp, mod=None):
if mod == 0:
raise ValueError("power() 3rd argument cannot be 0")
result = 1
for i in range(abs(exp)):
result *= base
if mod:
result %= mod
return result```

This will at least ensure that the power is computed. As a result, the test results change slightly:

```F..F....
======================================================================
FAIL: test_power_many (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(3, 5, 7)
self.assertEqual(
AssertionError: 243 != 0.00411522633744856 : Failed to compute 3^-5

======================================================================
FAIL: test_power_one_exponent (__main__.TestPower)
----------------------------------------------------------------------
Traceback (most recent call last):
self.assert_helper(5, 1, 7)
self.assertEqual(
AssertionError: 5 != 0.2 : Failed to compute 5^-1

----------------------------------------------------------------------
Ran 8 tests in 0.053s

FAILED (failures=2)```

From here, we need another branch to verify that `exp` is negative. If it is, we can invert the result before returning it:

```def power(base, exp, mod=None):
if mod == 0:
raise ValueError("power() 3rd argument cannot be 0")
result = 1
for i in range(abs(exp)):
result *= base
if exp < 0:
result = 1 / result
if mod:
result %= mod
return result```

And now with these changes in place, we get roughly the expected behavior. Here’s the proof!

```........
----------------------------------------------------------------------
Ran 8 tests in 0.062s

OK```

Surely, we could approximate the behavior a bit better, but I think this is a decent place to stop. If you’d like to take this code a bit further, feel free to share it with me on Twitter using #RenegadePython.

Hiccup Harvest

As I always say in the closing section, I write these articles because I love to teach. These sort of articles in particular are my opportunity to show myself making mistakes, so you can learn from them. Surely, I could write and polish up the code to be as clean and performant as possible, but that’s not the point. The point is to show you the exact process that I might go through to implement a function.

To be as transparent as possible, I actually write these articles more or less top to bottom, so you can see when and I why choose to move on to the next step. For example, when I develop the algorithm in the “Thought Process” section, I can’t possibly anticipate all the issues that may arise. As a result, I like to reach “good enough” status with my plans before trying to write some tests and ultimately trying to write a solution.

Naturally, the entire development process is iterative, so it makes sense that we have to revise our plans as run into issues. This article is fairly linear, so I don’t get a chance to share every hiccup along the way. That said, I’m happy to share some of the hiccups I ran into while writing this article.

• Early in the testing phase, I had assumed that the `mod` parameter was not optional, so testing failed for every method. To fix this, I went back to revise the function header to include the default value.
• Around the same time as the previous bullet, I remembered that floating point values exist and that it wouldn’t be trivial to revise my algorithm. As a result, I added an additional assumption to the problem description.

If as a community we can begin to share our mistakes, we’ll be a lot better off.

The purpose of these roll your own articles is threefold:

First, they allow me to take some time to practice my Python, and it’s fun trying to reverse engineering common Python functions and methods.

Second, they allow me to demonstrate the thought process of an experienced programmer to newer programmers.

Finally, they give me yet another way for folks in the community to contribute. If you’d like to share your own solution to this problem, head on over to Twitter and share your solution with #RenegadePython. Alternatively, I’m happy to check out your solutions in our Discord.

Beyond that, I’m happy to say we’ve rolled yet another function of our own. In the future, we’ll be looking to keep the series going by replicating common Python functions. Feel free to share your favorites with me over on Twitter.

Likewise, you’re welcome to stick around to check some of these related articles:

In addition, here are some resources from the folks at Amazon (#ad):

Otherwise, thanks again for sticking around! I appreciate your time, and I hope you’ll come back again soon.