Python Collatz Sequence as in OEIS A006577

collatz.py

#!/usr/bin/env python3

def collatz_sequence(n):
    """Generate the collatz sequence for n (see OEIS A006577)."""

    c = int(n)
    if c < 1 or c != n:
        raise ValueError(f"Argument n ({n}) must be a positive integer")

    while c != 1:
        yield c
        if (c % 2) == 0:
            c = c // 2
        else:
            c = (3 * c) + 1
    yield 1


def collatz_steps(n):
    """Returns # of collatz steps as defined in OEIS A006577."""
    return len(list(collatz_sequence(n))) - 1


if __name__ == "__main__":
    import unittest

    class TestMethods(unittest.TestCase):
        def test1(self):
            # test data literally copy-pasted from https://oeis.org/A006577
            testdata = (0, 1, 7, 2, 5, 8, 16, 3, 19, 6, 14, 9, 9, 17, 17,
                        4, 12, 20, 20, 7, 7, 15, 15, 10, 23, 10, 111, 18,
                        18, 18, 106, 5, 26, 13, 13, 21, 21, 21, 34, 8, 109,
                        8, 29, 16, 16, 16, 104)
            for n, k in enumerate(testdata):
                with self.subTest(n=n, k=k):
                    self.assertEqual(collatz_steps(n+1), k)

        def test2(self):
            badvals = (0, -1, 2.5)
            for bad in badvals:
                with self.subTest(bad=bad):
                    with self.assertRaises(ValueError):
                        next(collatz_sequence(bad))

        def test3(self):
            # arguably perhaps this should be an error
            self.assertEqual(collatz_steps(4.0), 2)

    unittest.main()