diff --git a/exercises/grains/example.py b/exercises/grains/example.py
index 0d42623f67..0b406e9dd4 100644
--- a/exercises/grains/example.py
+++ b/exercises/grains/example.py
@@ -1,20 +1,19 @@
-def on_square(square):
-    check_square_input(square)
-    return 2 ** (square - 1)
+def square(number):
+    check_square_input(number)
 
+    return 2 ** (number - 1)
 
-def total_after(square):
-    check_square_input(square)
-    return sum(
-        on_square(n + 1) for n
-        in range(square)
-    )
 
+def total(number):
+    check_square_input(number)
 
-def check_square_input(square):
-    if square == 0:
+    return sum(square(n + 1) for n in range(number))
+
+
+def check_square_input(number):
+    if number == 0:
         raise ValueError("Square input of zero is invalid.")
-    elif square < 0:
+    elif number < 0:
         raise ValueError("Negative square input is invalid.")
-    elif square > 64:
+    elif number > 64:
         raise ValueError("Square input greater than 64 is invalid.")
diff --git a/exercises/grains/grains.py b/exercises/grains/grains.py
index 89275f4579..194198fa68 100644
--- a/exercises/grains/grains.py
+++ b/exercises/grains/grains.py
@@ -1,6 +1,6 @@
-def on_square(integer_number):
+def square(number):
     pass
 
 
-def total_after(integer_number):
+def total(number):
     pass
diff --git a/exercises/grains/grains_test.py b/exercises/grains/grains_test.py
index 2c91354640..0f97d02c2e 100644
--- a/exercises/grains/grains_test.py
+++ b/exercises/grains/grains_test.py
@@ -1,55 +1,52 @@
 import unittest
 
-from grains import (
-    on_square,
-    total_after,
-)
+from grains import square, total
 
 
 # Tests adapted from `problem-specifications//canonical-data.json` @ v1.2.0
 
 class GrainsTest(unittest.TestCase):
     def test_square_1(self):
-        self.assertEqual(on_square(1), 1)
+        self.assertEqual(square(1), 1)
 
     def test_square_2(self):
-        self.assertEqual(on_square(2), 2)
+        self.assertEqual(square(2), 2)
 
     def test_square_3(self):
-        self.assertEqual(on_square(3), 4)
+        self.assertEqual(square(3), 4)
 
     def test_square_4(self):
-        self.assertEqual(on_square(4), 8)
+        self.assertEqual(square(4), 8)
 
     def test_square_16(self):
-        self.assertEqual(on_square(16), 32768)
+        self.assertEqual(square(16), 32768)
 
     def test_square_32(self):
-        self.assertEqual(on_square(32), 2147483648)
+        self.assertEqual(square(32), 2147483648)
 
     def test_square_64(self):
-        self.assertEqual(on_square(64), 9223372036854775808)
+        self.assertEqual(square(64), 9223372036854775808)
 
     def test_square_0_raises_exception(self):
         with self.assertRaisesWithMessage(ValueError):
-            on_square(0)
+            square(0)
         with self.assertRaisesWithMessage(ValueError):
-            total_after(0)
+            total(0)
 
     def test_square_negative_raises_exception(self):
         with self.assertRaisesWithMessage(ValueError):
-            on_square(-1)
+            square(-1)
         with self.assertRaisesWithMessage(ValueError):
-            total_after(-1)
+            total(-1)
 
     def test_square_gt_64_raises_exception(self):
         with self.assertRaisesWithMessage(ValueError):
-            on_square(65)
+            square(65)
         with self.assertRaisesWithMessage(ValueError):
-            total_after(65)
+            total(65)
 
     def test_total(self):
-        self.assertEqual(total_after(64), 18446744073709551615)
+        self.assertEqual(total(64), 18446744073709551615)
 
     # Utility functions
     def setUp(self):