Тестирование скорости кода python

Тестирование скорости кода python

Часто, когда мы создаем какую-то программу, у нас может стоять выбор между несколькими подходами в написании кода. Например, один вариант использует цикл for и словари, а второй списковые включения и элементы функционального программирования. Оба варианта выдают правильный ответ, но скорость работы может отличаться.

И правильное решение в этом случае — провести замеры производительности, хотя бы элементарные.

Либо, еще один из сценариев, программа уже выложена и мы знаем что какая-то её часть работает медленно и мы бы хотели её ускорить. И после написания нового оптимизированного кода, нам, опять же, нужно сравнить скорость выполнения двух, а то и трёх вариантов.

И в Python для этого используют встроенный модуль timeit, который предназначен для тестирования производительности небольших фрагментов кода.

Есть разные способы применения этого модуля, в том числе из командной строки, я же покажу вариант, которым регулярно пользуюсь сам. А также раскрою некоторые особенности и подводные камни.

Для примера возьмем два способа подсчета суммы элементов списка: первый с помощью обычного цикла for, а второй с помощью функции sum():

lst = [1, 2, 3, 4, 5]

def f1(numbers):
    total = 0
    for value in numbers:
        total += value
    return total

def f2(numbers):
    return sum(numbers)

Очевидно, что второй вариант занимает гораздо меньше места, но быстрее ли он работает?

Для того чтобы провести тесты необходимо предварительно поместить оба варианта внутрь функций: у меня это f1() и f2(). Также вверху программы я создал список из 5 элементов.

Теперь, под функциями следует написать такое выражение:

assert f1(lst) == f2(lst)

С помощью этого кода, а точнее с помощью инструкции assert происходит предварительный запуск функций и проверка того, что их ответы совпадают. Если ответы равны, то ничего не происходит и программа продолжает выполняться дальше.

Если же результаты не совпадают, то возникает исключение AssertionError и программа останавливается. Такой код добавляется для проверки того, что если после внесения в какую-то из функций изменений, её работа не изменится.

Бывает что-то исправишь и раз, функция начинает работать в 2 раза быстрее, но неправильно. Такие моменты лучше сразу отлавливать.

Далее (после функций) надо добавить такой код:

import timeit
print(timeit.timeit('f1(lst)', globals=globals()))
print(timeit.timeit('f2(lst)', globals=globals()))

Итак, мы импортируем модуль timeit и далее вызываем одноименную функцию timeit, которая как раз и занимается подсчетом производительности.

Первым параметром функция принимает код, который нужно протестировать. Код в виде строки! То есть фактически она запустит функции f1() и f2() с передачей списка lst.

Так как код мы записываем в виде строки, а на просто запускаем как выше, то нам нужно дополнительно предать в функцию сами имена f1, f2, lst. Для этого мы через аргумент globals передаем глобальное пространство имен. В нём как раз будут и функции, и списки, которые timeit проверит.

Что ж, давайте запустим программу. Сперва через IDE:

0.175016834
0.11417374999999999

Как видите первый вариант с циклом работает в 1,5 раза медленней чем вариант с использованием sum(). При этом сейчас программа выполняется в Python версии 3.9

Но если запустить тот же код в Python 3.14, то ут вторая функция работает быстрее почти в два раза. И при этом каждый из вариантов работает быстрее чем в Python 3.9:

0.10523716604802758
0.0612475840607658

То есть, чтобы выжать из вашего кода максимум, нужно не только выбрать второй вариант, но еще и обновиться до последней версии питона.

Как работает timeit

Но давайте вернемся к timeit и поговорим о том, как функция работает. И фактически она просто берет ваш код и выполняет его 1 миллион раз подряд и считает суммарное время выполнения.

Но миллион раз — это не всегда удобное значение. Иногда код работает так быстро, что даже 1 миллион запусков проходит почти моментально. Например, сейчас наши функции суммарно выполнялись менее секунды.

Я предпочитаю, чтобы время выполнения каждого теста было в районе 1-3 секунд. Для этого необходимо добавить в timeit дополнительный параметр number, который отвечает за количество запусков:

print(timeit.timeit('f1(lst)', globals=globals(), number=10_000_000))
print(timeit.timeit('f2(lst)', globals=globals(), number=10_000_000))

После запуска кода мы получим значения в районе одной секунды:

1.430678042001091
0.9172235419973731

Тест крайних значений

Также при тестах производительности, нужно проверить как функции ведут себя на данных разной длины. Если длина входящего списка может колебаться от 5 до 500 значений, то надо, как минимум, протестировать крайние варианты. Для этого вверху программы можно добавить количество элементов с помощью функции range():

lst = list(range(1, 501))

После запуска такого кода программа зависнет, потому что 10 млн запусков для суммирования 500 элементов занимает в 100 раз больше времени, чем суммирование 5 элементов.

Поэтому нужно прервать выполнение и уменьшить number до 1 миллиона. Иногда приходится несколько раз изменять number, чтобы подобрать нужное время выполнения: не слишком длинное и не чересчур короткое.

В этот раз мы смогли дождаться результата и обратите внимание, что вторая функция работает практически в 5 раз быстрее чем первая.

7.006657332996838
1.359301874996163

То есть размер входящих данных также имеет значение. Несмотря на то, что обе функции выполняются за время O(N).

Использование функции print() в timeit

Теперь давайте в самих функциях заменим return на print():

def f1(numbers):
    total = 0
    for value in numbers:
        total += value
    print(total)

def f2(numbers):
    print(sum(numbers))

А количество проверок можно уменьшить в 10 раз. После запуска кода, в консоль будут падать результаты работы функций. По 100 000 раз на каждую функцию.

Это одна из причин, почему использовать print() с выводом в консоль на тестах производительности нельзя, так как мы физически не сможем отыскать скорость работы первой функции среди большого количества мусора.

Вторая причина — сам print() также влияет на скорость выполнения функции, но по факту он не несет никакой ценности.

То есть print() не является узким звеном алгоритма, который мы хотим ускорить, print() просто выводит данные, но при этом он тормозит программу.

Даже если ваша финальная функция подразумевает print(), на тестах производительности лучше использовать return.

Работа с файлами в timeit

И еще будьте аккуратны при тестировании работы с файлами. Чтение и запись в файлы всегда будет медленней, чем работа с оперативной памятью. Одно дело прочитать список, а другое дело прочитать файл. Еще медленнее, если нам нужно записывать данные.

В этом случае количество тестов нужно еще сильнее уменьшать, иногда вплоть до 100 или 1000.

100 тестов тоже будет достаточно для того, чтобы сравнить несколько вариантов кода.

Финальный код

# Данные
lst = list(range(1, 501))

# Функции с разным кодом
def f1(numbers):
    total = 0
    for value in numbers:
        total += value
    return total

def f2(numbers):
    return sum(numbers)

# Проверка, что функции работают одинаково
assert f1(lst) == f2(lst)

# Тестирование скорости кода функций
import timeit
print(timeit.timeit('f1(lst)', globals=globals(), number=1_000_000))
print(timeit.timeit('f2(lst)', globals=globals(), number=1_000_000))