Krok Do Pythonicznego Kodu – Comprehension

Jarosław Piszczała
14 May, 2020
14 May, 2020
- Collections

W społeczności Pythonowej istnieje takie określenie kodu, że jest on Pythoniczny. Oznacza to, że wykorzystuje on możliwości i cechy charakterystyczne tego języka. Jednym z takich elementów są tak zwane złożenia (ang. comprehension), nazywane też wyrażeniami (np. wyrażenia listowe ang. list comprehension). Ich rodzaje oraz sposoby użycia powinien znać każdy fan Pythona!

Comprehension

Aby zrozumieć, jak wygląda implementacja złożeń, oraz jak zastąpić zwykłą pętlę tą konstrukcją, przejdźmy do przykładu. Stworzymy trzy funkcje, które w tym pomogą. Pierwsza to map_func, której zadaniem mogłoby być przekształcenie naszych danych w pewien sposób. Kolejna to filter_func, której zadaniem byłoby określenie, czy dana zmienna powinna się znaleźć w nowej kolekcji, czy nie. Funkcja ta zwraca bool, ponieważ oczekujemy konkretnej odpowiedzi: True albo False. Ostatnia funkcja to iterable, która zwraca nam obiekt iterowalny - po jego elementach będzie tworzyć pętlę.

1
    from typing import Iterable, Any
2

3

4
    def map_func(obj) -> Any:
5
        return obj
6

7

8
    def filter_func(obj) -> bool:
9
        return bool(obj)
10

11

12
    def iterable() -> Iterable:
13
        return range(-2, 5)

Gdy mamy już zdefiniowane takie elementy, to możemy podstawić je pod standardową konstrukcję pythonową. Nie korzystając ze złożeń, musielibyśmy najpierw utworzyć pustą kolekcję, a następnie poprzez pętle sprawdzać kolejne elementy i dodawać tylko te, które są poprawne (z naszego punktu widzenia).

1
    result = []
2
    for obj in iterable():
3
        if filter_func(obj):
4
            result.append(map_func(obj))

Złożenia Pythonowe bardzo uproszczą nam ten kod. Cała logika mieści się nam w jednej linii kodu. Zamykamy ją w odpowiednich nawiasach (w zależności od typu złożenia). Wynik operacji będzie automatycznie dodawany do listy. Poniższe to przykład list comprehension.

1
    result = [map_func(obj) for obj in iterable() if filter_func(obj)]
2

3
    print(result)

1
    [-2, -1, 1, 2, 3, 4]

Aby dokładniej przyjrzeć się, jak zmieniła się nasza logika, przeformatujmy nasze złożenie. Po pierwsze nie definiujemy już pustej kolekcji. Po drugie zmieniła się trochę kolejność czytania kodu. Poprzednie rozwiązanie mogliśmy czytać jako:

Dla każdego obiektu w obiekcie iterowalnym, jeżeli spełnia on pewne założenie, to przekonwertuj go i dodaj do kolekcji.

Nowe rozwiązanie będziemy czytać jako:

Przekonwertuj i dodaj do nowej kolekcji każdy obiekt w obiekcie iterowalnym, jeżeli spełnia on pewne założenia.

1
    result = [
2
        map_func(obj)
3
        for obj in iterable()
4
        if filter_func(obj)
5
    ]

Dlaczego korzystamy w takim razie z tej konstrukcji w naszym kodzie? Jest to czytelna i szybka metoda przekształcania, filtrowania i tworzenia nowych kolekcji, na podstawie już istniejących.

Aby nie było wątpliwości co do powyższych przykładów. Gdyby nie chęć pokazania przykładu z opisami poszczególnych elementów, nasze złożenie wyglądałoby w tym przypadku następująco:

1
    result = [num for num in range(-2, 5) if num]

Przypadek użycia

Aby uatrakcyjnić nasze przykłady, skupmy się na jednym przypadku użycia. Użytkownik na wejściu do naszej aplikacji przekazuje string liczb oddzielonych białym znakiem. Naszym celem jest odpowiednio przekształcić to wejście w pewną kolekcję. Aby kod był czytelniejszy, utworzymy zmienną input_iter, która będzie kolekcją kolejnych liczb przekazanych przez użytkownika, wciąż w formacie str.

1
    user_input = "123 52 43 -11 1 9 4 0 23 1 1 -2"
2
    input_iter = user_input.split()

List comprehension

Najprostszym i najczęściej spotykanym przykładem jest list comprehension. Poniższy przykład pokazuje, w jaki sposób przekształcić naszą kolekcję liczb w formacie str do int. Jest to dość prosty zabieg, więc widzimy zdecydowanie, dlaczego developerzy tak bardzo lubią tę konsturkcję :)

1
    numbers = []
2
    for number in input_iter:
3
        numbers.append(int(number))
4

5
    numbers = [int(number) for number in input_iter]
6
    print(numbers)

1
    [123, 52, 43, -11, 1, 9, 4, 0, 23, 1, 1, -2]

Dict comprehension

Gdy potrzebujemy przekształcić kolekcję na słownik, z pomocą przychodzi dict comprehension. Nasz nowy przykład pokazuje mapowanie liczb jako kluczy do wartości, którą jest bool definiujący czy przed liczbą stoi znak. Jak widać, różne typy złożeń definiujemy poprzez zmianę nawiasów.

1
    numbers_sign = {}
2
    for number in input_iter:
3
        numbers_sign[int(number)] = number[0] == "-"
4

5
    numbers_sign = {int(number): number[0] == "-" for number in input_iter}
6
    print(numbers_sign)

1
    {
2
        123: False, 52: False, 43: False, -11: True, 1: False,
3
        9: False, 4: False, 0: False, 23: False, -2: True
4
    }

Set comprehension

W sytuacji, gdy zachodzi potrzeba utworzenia kolekcji unikalnych elementów, przydać się może set comprehension. Dzięki niemu dość szybko stworzymy kolekcje unikalnych liczb podanych przez użytkownika.

1
    unique_numbers = set()
2
    for number in input_iter:
3
        unique_numbers.add(int(number))
4

5
    unique_numbers = {int(number) for number in input_iter}
6
    print(unique_numbers)

1
    {0, 1, 4, 9, 43, 52, -11, 23, 123, -2}

If

Złożenia oferują nam także możliwość filtrowania danych w kolekcjach. Aby to zrobić, w list comprehension należy na końcu po pętli dodać if z konkretnym predykatem. Na nasze potrzeby chcemy utworzyć listę wyłącznie dodatnich liczb.

1
    positive_numbers = []
2
    for number in input_iter:
3
        if int(number) >= 0:
4
            positive_numbers.append(int(number))
5

6
    positive_numbers = [int(number) for number in input_iter if int(number) >= 0]
7
    print(positive_numbers)

1
    [123, 52, 43, 1, 9, 4, 0, 23, 1, 1]

If else

Nie zawsze poprzez if pragniemy odfiltrować dane. Zdarza się, że w różnych przypadkach chcemy otrzymać inną wartość. Wtedy możemy skorzystać z if else w formie tzw. conditional expression. W tym przypadku każdą ujemną wartość będziemy przekształcać na dodatnią.

1
    positive_numbers = []
2
    for number in input_iter:
3
        if int(number) >= 0:
4
            positive_numbers.append(int(number))
5
        else:
6
            positive_numbers.append(0 - int(number))
7

8
    positive_numbers = [
9
        int(number) if int(number) >= 0 else 0 - int(number) for number in input_iter
10
    ]
11
    print(positive_numbers)

1
    [123, 52, 43, 11, 1, 9, 4, 0, 23, 1, 1, 2]

For x,y

Tak jak w prostych pętlach, tak i tutaj możemy iterować po kilku elementach poprzez rozpakowanie tupli. W ten sposób możemy, chociażby iterować po kluczach i wartościach słowników, lub odpowiednio do naszych potrzeb rozpakowywać dane.

1
    last_digits = []
2
    for *digits, last_digit in input_iter:
3
        last_digits.append(last_digit)
4

5
    last_digits = [last_digit for *digits, last_digit in input_iter]
6
    print(last_digits)

1
    ['3', '2', '3', '1', '1', '9', '4', '0', '3', '1', '1', '2']

Nested

Złożenia nie są ograniczone. Dzięki temu możemy pozwolić także na zagnieżdżone pętle w złożeniach! Poniższy przykład utworzy kolekcje unikalnych cyfr użytych przez użytkownika!

1
    unique_numbers = set()
2
    for number in input_iter:
3
        for char in number:
4
            if char != '-':
5
                unique_numbers.add(int(char))
6

7
    unique_numbers = {int(char) for number in input_iter for char in number if char != "-"}
8
    print(unique_numbers)

1
    {0, 1, 2, 3, 4, 5, 9}

Generator expression

Jak wiemy, nie zawsze tworzenie kolekcji ma sens. Gdy chcemy przetworzyć dane, aby je zaraz wykorzystać, dużo lepszym rozwiązaniem jest utworzenie generatora. Jednak tworzenie zwykłego generatora oznacza tworzenie funkcji, która poprzez yield zwracać będzie kolejne wartości.

1
    def numbers(user_input):
2
        for number in user_input:
3
            yield int(number)
4

5

6
    print(numbers)
7
    numbers_sum = sum(numbers(input_iter))

1
    <function numbers at 0x7fd5b9b661f0>

W tym wypadku możemy taką konstrukcję zastąpić poprzez tzw. generator expression które niczym nie różni się od możliwości, które zaprezentowane zostały powyżej przy innych złożeniach. Taka konstrukcja wymaga zamknięcia naszego list comprehension nawiasami okrągłymi, jak przy definicji tuple.

1
    numbers = (int(number) for number in input_iter)
2
    print(numbers)
3
    numbers_sum = sum(numbers)

1
<generator object <genexpr> at 0x7fd5bb77e9e0>

Jeżeli powyższe rozwiązanie nie do końca Ci się nie podoba, to bardzo dobrze! Ponieważ gen expr można wstawiać bezpośrednio jako argument, pod tym względem jego zachowanie można porównać z funkcją anonimową lambda.

1
    numbers_sum = sum(int(number) for number in input_iter)

Jest pewien haczyk. Generator pozwoli nam przejść po danych wyłącznie jeden raz. Miejmy to na uwadze, ponieważ jest to ważny wyznacznik tego, czy warto z niego skorzystać w danej sytuacji.

Nested comprehension

Jeżeli powyższe przykłady nie rozwiązują do końca problemu, to może pewna wariacja to zrobi. Pomyślmy, że potrzebujemy zrobić pewną operację na danych, a następnie te dane przetworzyć. Możemy do tego skorzystać z zagnieżdżenia się jednego złożenia w drugim. W poniższym przykładzie chcemy przechowywać tuple sąsiadów liczb przekazanych przez użytkownika.

1
    number_neighbours = []
2
    for number in input_iter:
3
        number = int(number)
4
        result = (number - 1, number, number + 1)
5
        number_neighbours.append(result)
6

7
    number_neighbours = [
8
        (int(number) - 1, int(number), int(number) + 1) for number in input_iter
9
    ]
10

11
    print(number_neighbours)

1
    [
2
        (122, 123, 124), (51, 52, 53), (42, 43, 44), (-12, -11, -10),
3
        (0, 1, 2), (8, 9, 10), (3, 4, 5), (-1, 0, 1), (22, 23, 24),
4
        (0, 1, 2), (0, 1, 2), (-3, -2, -1)
5
    ]

To przekształcanie potrójne tej samej wartości odbiera trochę sens, i nasz kod wydaje się nie być wcale lepszy. Jednak jest na to rada, i jest nią zagnieżdżenie. Zagnieżdżmy gen expr, które przekształcać będzie nasze liczby w formacie str na int przed pobraniem kolejnej wartości.

1
    number_neighbours = [
2
        (number - 1, number, number + 1) for number in (int(n) for n in input_iter)
3
    ]

I teraz spójrzmy, jak byśmy zrobili tę samą konstrukcję, nie korzystając ze złożeń:

1
    def numbers(iterable):
2
        for number in iterable:
3
            yield int(number)
4

5

6
    number_neighbours = []
7
    for number in numbers(input_iter):
8
        number_neighbours.append((number - 1, number, number + 1))

Albo bez użycia generatora tak:

1
    numbers = []
2
    for number in input_iter:
3
        numbers.append(int(number))
4

5
    number_neighbours = []
6
    for number in numbers:
7
        number_neighbours.append((number - 1, number, number + 1))

Jak pokazują przykłady, poznana tutaj konstrukcja daje dużą elastyczność przy pisaniu kodu. W dodatku zdecydowanie upraszcza kod, który musielibyśmy napisać w to miejsce.

Use cases

Skoro już wiemy, jak to wszystko działa, spróbujmy z tego skorzystać w innym przypadku. Dużo ciekawiej jest, gdy operujemy na słownikach i potrzebujemy wyciągnąć na ich podstawie pewne informacje. Do przykładów skorzystamy z prostego słownika, który zawiera kilku użytkowników:

1
    from pprint import pprint
2

3
    users = [
4
        {
5
            "id": 1,
6
            "username": "jbuzzing0",
7
            "email": "[email protected]",
8
            "ip_address": "171.228.113.240",
9
            "is_active": True,
10
            "is_admin": True,
11
        },
12
        {
13
            "id": 2,
14
            "username": "afettiplace1",
15
            "email": "[email protected]",
16
            "ip_address": "239.233.89.230",
17
            "is_active": True,
18
            "is_admin": False,
19
        },
20
        {
21
            "id": 3,
22
            "username": "rsparry2",
23
            "email": "[email protected]",
24
            "ip_address": "58.241.139.102",
25
            "is_active": False,
26
            "is_admin": True,
27
        },
28
        {
29
            "id": 4,
30
            "username": "fbrinkman3",
31
            "email": "[email protected]",
32
            "ip_address": "190.242.241.245",
33
            "is_active": True,
34
            "is_admin": False,
35
        },
36
        {
37
            "id": 5,
38
            "username": "gperren4",
39
            "email": "[email protected]",
40
            "ip_address": "226.186.181.230",
41
            "is_active": True,
42
            "is_admin": False,
43
        },
44
    ]

W pierwszej kolejności chcemy sprawdzić, czy na naszej liście użytkowników jest jakiś aktywny administrator. Robimy to poprzez set comprehension i sprawdzamy, czy w kolekcji, którą otrzymaliśmy, jest wartość True.

1
    any_admin = any({user["is_admin"] and user["is_active"] for user in users})
2
    pprint(any_admin)

1
    True

Skoro już wiemy, że wśród użytkowników istnieje jakiś administrator, stwórzmy ich listę. Nowa lista będzie zawierać tylko tych użytkowników, których wartość spod klucza is_admin jest równa True.

1
    admins = [user for user in users if user["is_admin"]]
2
    print(admins)
3

4
    [{'email': '[email protected]',
5
      'id': 1,
6
      'ip_address': '171.228.113.240',
7
      'is_active': True,
8
      'is_admin': True,
9
      'username': 'jbuzzing0'},
10
     {'email': '[email protected]',
11
      'id': 3,
12
      'ip_address': '58.241.139.102',
13
      'is_active': False,
14
      'is_admin': True,
15
      'username': 'rsparry2'}]

Co jeżeli potrzebujemy obojętnie jakiego administratora? Zamiast tworzyć ich listę, tworzymy prosty generator, który zwróci nam pierwszego administratora. Dzięki temu otrzymamy do niego dostęp szybciej, niż w przypadku gdybyśmy tworzyli najpierw całą listę, aby następnie chcieli pobrać tylko jej pierwszy element.

1
    first_admin = next(user for user in users if user["is_admin"] and user["is_active"])
2
    pprint(first_admin)

1
    {'email': '[email protected]',
2
     'id': 1,
3
     'ip_address': '171.228.113.240',
4
     'is_active': True,
5
     'is_admin': True,
6
     'username': 'jbuzzing0'}

Może potrzebujemy adresy email wszystkich użytkowników, aby przesłać im wiadomości? Wystarczy dla każdego użytkownika wyciągnąć wartość spod odpowiedniego klucza

1
    emails = [user['email'] for user in users]
2
    pprint(emails)

1
    ['[email protected]',
2
     '[email protected]',
3
     '[email protected]',
4
     '[email protected]',
5
     '[email protected]']

Timing

Skoro już wiemy, z czym mamy do czynienia, i jak z tego korzystać - dowiedzmy się czy komputer też wolałby comprehension. Użyjemy do tego biblioteki timeit. Użyjemy do tego celu prostej pętli z predykatem.

1
    from timeit import timeit
2

3

4
    func = """
5
    arr = []
6
    for i in range(1000):
7
        if(i%3):
8
            arr.append(i)
9
    """
10

11
    time = timeit(func, number=10000)
12
    print(f"{'standard':>14} {time}")
13

14

15
    func = """
16
    arr = [i for i in range(1000) if i%3]
17
    """
18

19
    time = timeit(func, number=10000)
20
    print(f"{'comprehension':>14} {time}")

1
          standard 1.2301481989998138
2
     comprehension 0.8710038739955053

Wynik to ostateczny powód, dlaczego warto się tej konstrukcji nauczyć. W tym wypadku jest ona szybsza od standardowej funkcji aż o 33%!

Kiedy Nie Stosować?

Złożeń nie należy stosować w kilku sytuacjach:

Gdy robimy pętle po obiektach, które nas nie interesują. Może to być pętla, która wysyłać będzie wiadomości z listy, a metoda do wysyłania będzie zwracała None. Stosujemy to tylko wtedy gdy chcemy utworzyć nową kolekcję. Więc prostą zasadą, jeżeli ktoś zrobi taką konstrukcję, ale nie przypisze jej do żadnej wartości, to należy ją przepisać na prostą pętlę.
Gdy nasza aktualna implementacja jest skomplikowana, a zapisanie jej poprzez złożenie tylko utrudni zrozumienie kodu. Jest to fajne wyzwanie na napisanie jak najmniejszej ilości linii kodu, ale zdecydowanie antywzorzec w kodzie produkcyjnym. Pamiętajmy, że w Pythonie stawiamy na czytelność. Przykład takiego zakazanego kodu poniżej.

1
    def fibonacci(n):
2
        fib = [1, 1]
3
        return [sum(fib[-2:]) + int(str(fib.append(sum(fib[-2:]))).replace('None', '0')) for i in range(n)][-3] if n>=3 else 1

Podsumowanie

I to (chyba?) tyle. Przyjrzeliśmy się dokładnie wszystkim formą złożeń, jakie spotkać możemy w Pythonie, przyjrzeliśmy się pewnym przypadkom użycia, a na koniec udowodniliśmy, że jest to rozwiązanie szybsze!

Najlepszą praktyką, jest stosowanie złożeń wszędzie tam, gdzie potrzebujemy utworzyć kolekcje, lub gdy na podstawie jednej kolekcji chcemy utworzyć nową. Gdy chcemy przetworzyć jakieś argumenty lub je odfiltrować. Nie powinniśmy natomiast złożeń stosować także wtedy, gdy nasze rozwiązanie straci przez to na czytelności. Pamiętajmy też o generatorach, gdy kolekcja, którą tworzymy, zostanie wykorzystana tylko raz.

W ramach nauki polecam także wykorzystać plugin do flake8, który zwróci uwagę dodatkowo na złożenia w kodzie, jest to flake8-comprehensions.