שיתוף קוד בין מחלקות Python באמצעות Decorators

24/01/2021

פוסט זה כולל טיפ קצר לעבודה יעילה עם Python. אם אתם רוצים ללמוד פייתון יותר לעומק אני ממליץ על קורס Python כאן באתר.
הקורס כולל עשרות שיעורי וידאו והמון תרגול מעשי וילמד אתכם Python בצורה מקצועית מההתחלה ועד הנושאים המתקדמים.

כולם היום יודעים לדקלם איך כשיש לנו מחלקות עם קוד משותף אפשר להשתמש בירושה כדי לכתוב את הקוד הכפול רק פעם אחת, אבל מה קורה כשהדבר שאנחנו רוצים לשתף הוא קצת יותר מסובך מפונקציות ומידע? מה קורה אם כל מחלקה משתפת את הקוד הזה בצורה קצת אחרת? בפוסט זה נראה דוגמה בה ירושה לא עובדת כל כך טוב, ובמקומה נראה פיתרון אחר לשיתוף קוד בין מחלקות באמצעות Decorators.

1. מה אנחנו בונים

נניח שאנחנו רוצים לבנות משחק שיש בו כדור קופץ ומלבן שצריך לתפוס את הכדור הזה. אפשר לדמיין שכל אחת מהצורות במערכת תיוצג באמצעות מחלקה ולכל צורה יש מיקום ומהירות. בקוד פייתון (עם pyxel בשביל הציורים) הקוד נראה כך:

@dataclass
class Ball:
    x: int = 0
    vx: int = 2
    y: int = 20

    color: int = 2
    radius: int = 5

    def draw(self):
        pyxel.circ(self.x, self.y, self.radius, self.color)

    def update(self):
        self.x += self.vx


@dataclass
class Rect:
    x: int = 0
    vx: int = 2
    y: int = 0
    color: int = 5

    def draw(self):
        pyxel.rect(self.x, self.y, 60, 20, self.color)

    def update(self):
        self.x += self.vx

בתוכנית עצמה אפשר ליצור כדורים ומקלות ואפילו להזיז אותם (כרגע פשוט כולם זזים ימינה) עם קוד שנראה כך:

shapes = [Ball(x=x, y=10) for x in range(0, 100, 30)] + [Rect(x=0, y=100)]


def draw():
    pyxel.cls(0)
    for shape in shapes:
        shape.draw()


def update():
    for shape in shapes:
        shape.update()

pyxel.init(160, 120)
pyxel.run(update, draw)

הקוד עובד אבל קל לראות שיש לנו כפל קוד בין שתי המחלקות:

  1. בשתי המחלקות יש משתנים לייצוג המיקום על ציר x והמהירות על ציר x

  2. בשתי המחלקות יש פונקציית update שמעדכנת את המיקום על ציר x לפי המהירות

אפשר גם לדמיין שבעתיד נצטרך להזיז את הכדורים גם על ציר y, ואולי נרצה משתנה נפרד עבור המהירות בציר y. אולי אפשר לדמיין משחק תלת מימדי עם מיקום ומהירות על ציר z. וזה מביא אותנו לאתגר - איך אפשר לשתף בצורה קלה את הקוד שמטפל במיקום ובמהירות על ציר מסוים, כאשר כל צורה תוכל להגדיר לעצמה על איזה צירים היא יכולה לזוז?

2. הפיתרון: Decorators

כיוון אחד לפיתרון יהיה להגדיר את כל הצירים האפשריים במחלקת-אב בשם Shape. אני משאיר לכם את כתיבת הפיתרון הזה בתור תרגיל, אבל אגיד רק שאני לא חושב שזה פיתרון מספיק טוב במקרה הכללי. כי לא כל הצורות צריכות את המשתנים לכל הצירים.

כיוון יותר מעניין הוא שימוש ב Decorators: אם אני יכול להוסיף בצורה דינמית את המשתנים המתאימים לציר מסוים, ולעדכן את הפונקציה update ולהוסיף לה את התנועה באותו כיוון, אוכל לייצר Decorator שיוסיף את המשתנים ויעדכן את update על הצירים שאני צריך.

קוד? ברור. שימו לב לפונקציה הבאה:

def add_axis(axis, initial_value=0, initial_speed=2):
    def class_wrapper(cls):
        cls.__annotations__[axis] = int
        cls.__annotations__['v' + axis] = int
        setattr(cls, axis, initial_value)
        setattr(cls, 'v' + axis, initial_speed)
        original_update = getattr(cls, 'update')

        def update(self):
            original_update(self)
            setattr(self, axis, getattr(self, axis) + getattr(self, 'v' + axis))

        setattr(cls, 'update', update)

        return cls

    return class_wrapper

הפונקציה מקבלת מחלקה אחרת ומעדכנת את שדה __annotations__ שלה עם המשתנים המתאימים לציר החדש שאני יוצר, וגם יוצרת ערכי ברירת מחדל למשתנים בציר זה. אם axis יהיה x נקבל את המשתנים x ו-vx. אם axis יהיה y נקבל את המשתנים y ו vy, וכך לכל ציר שאפשר לדמיין.

בשביל להפעיל את הפונקציה אני משתמש בכתיב ה Decorators על המחלקה, ועל הדרך גם מוסיף לכדורים תנועה על ציר y:

@dataclass
@add_axis('x', 0, 2)
@add_axis('y', 20, 1)
class Ball:
    color: int = 2
    radius: int = 5

    def draw(self):
        pyxel.circ(self.x, self.y, self.radius, self.color)

    def update(self):
        pass


@dataclass
@add_axis('x', 0, 2)
@add_axis('y', 20, 0)
class Rect:
    color: int = 5

    def draw(self):
        pyxel.rect(self.x, self.y, 60, 20, self.color)

    def update(self):
        pass

בניגוד לירושה, שיתוף קוד באמצעות Decorators באופן זה מעודד אותנו להדביק לתוך המחלקה התנהגויות מסוגים שונים, מה שקצת מזכיר תבנית של Mixins. בצד החסרונות מדובר במימוש מסובך יותר ולכן קשה יותר לתחזוקה.