• בלוג
  • היום למדתי: הגדרת פרוטוקול ב Python

היום למדתי: הגדרת פרוטוקול ב Python

04/02/2023

מנגנון Type Hints של פייתון מציע דרך להוסיף הגדרת טיפוסים למשתנים ופונקציות שלנו. כלים כמו mypy יכולים אז לעבור על הקוד ולזהות מתי ניסינו לעשות משהו שלא יעבוד בגלל בעיית טיפוסים ולהציל אותנו מבאגים טפשיים.

1. הבעיה

כמובן שבשביל לקבל את ההגנה צריך לדעת לכתוב את הטיפוס לכל "דבר" בקוד, וזה לא תמיד פשוט. נתבונן בתוכנית הבאה:

def paint_circle(*, x: int, y: int):
    print(f"Circle: x={x}, y={y}")


def paint_triangle(*, x: int, y: int):
    print(f"Triangle: x={x}, y={y}")


def bad_painter(x, y, z):
    print("Bad painter: x={x}, y={y}, z={z}")


def paint(shapes: List[Callable]):
    for shape in shapes:
        shape(x=10, y=20)


paint([
    paint_circle,
    paint_circle,
    bad_painter,
    paint_triangle,
])

התוכנית מגדירה פונקציה בשם paint שמקבלת רשימה של פונקציות ציור, ומפעילה כל פונקציית ציור עם keyword arguments בשמות x ו y. בעזרת Callable הצלחתי להגדיר בקלות שאני מחפש רשימה של פונקציות, אבל Callable לא נותן לי דרך להגיד שהפונקציות צריכות לקבל פרמטרים בשמות מסוימים (כן אפשר להוסיף סוגריים מרובעים אחרי Callable כדי לקבוע את הטיפוסים הרגילים שפונקציה מקבלת, אבל לא את ה Keyword Arguments).

לכן כשמפעילים mypy על הקוד שלמעלה מקבלים:

Success: no issues found in 1 source file

2. פרוטוקולים להצלה

הפיצ'ר של mypy שיכול לעזור כאן נקרא Protocol. לפי התיעוד אנחנו מגדירים אותו באופן הבא:

from typing import Optional, Iterable
from typing_extensions import Protocol

class Combiner(Protocol):
    def __call__(self, *vals: bytes, maxlen: Optional[int] = None) -> list[bytes]: ...

וכך יכולים לקבוע מה יהיו הפרמטרים של הפונקציה כשקוראים לה כולל מה יהיו השמות של ה Keyword Arguments. בתוכנית שלנו פרוטוקול מציל את המצב בקלות:

from typing import Callable, TypeVar, List
from typing_extensions import Protocol


class Painter(Protocol):
    def __call__(self, *, x: int, y: int) -> None: ...


def paint_circle(*, x: int, y: int):
    print(f"Circle: x={x}, y={y}")


def paint_triangle(*, x: int, y: int):
    print(f"Triangle: x={x}, y={y}")


def bad_painter(*, x, y, z):
    print("Bad painter: x={x}, y={y}, z={z}")


def paint(shapes: List[Painter]):
    for shape in shapes:
        shape(x=10, y=20)


paint([
    paint_circle,
    paint_circle,
    bad_painter,
    paint_triangle,
])

ועכשיו כשמפעילים mypy מקבלים את השגיאה:

named_arguments.py:29: error: List item 2 has incompatible type "Callable[[NamedArg(Any, 'x'), NamedArg(Any, 'y'), NamedArg(Any, 'z')], Any]"; expected "Painter"
Found 1 error in 1 file (checked 1 source file)

שווה לשים לב שלפחות אצלי PyCharm בלי mypy לא הצליח לזהות את הטעות, ואת השגיאה ראיתי רק בהפעלה ישירה של mypy משורת הפקודה או דרך הפלאגין בפייצ'ארם.