בואו נתקן את cycle בפייתון

29/02/2024

הפונקציה cycle מתוך המודול itertools בפייתון לוקחת אוסף של דברים ומחזירה איטרטור אינסופי שכל פעם מחזיר את הדבר הבא מתוך הרצף במעגל. לדוגמה אפשר להשתמש בה עם המחרוזת abc באופן הבא:

print(list(itertools.islice(itertools.cycle("abc"), 10)))

ולקבל את התוצאה:

['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c', 'a']

אבל לסייקל יש בעיית ממשק - היא יכולה לקבל רק פרמטר אחד. לכן אם במקום רשימת האותיות הייתי רוצה להעביר רשימה של מספרים הייתי צריך לכתוב:

print(list(itertools.islice(itertools.cycle([1, 2, 3]), 10)))

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

def easy_cycle(*args):
    match args:
        case [col] if hasattr(col, '__iter__'):
            return itertools.cycle(col)

        case _:
            return itertools.cycle(args)

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

התוצאה פשוטה ועכשיו הממשק מסתדר גם כשאני שוכח סוגריים מרובעים:

print(list(itertools.islice(easy_cycle("abc"), 10)))
print(list(itertools.islice(easy_cycle("a", "b", "c"), 10)))
print(list(itertools.islice(easy_cycle(["a", "b", "c"]), 10)))

נ.ב. היה נחמד אם מנגנון ה Type Hints של פייתון היה מאפשר לי להגדיר שערך ההחזר של הפונקציה שכתבתי זהה לערך ההחזר של הפונקציה itertools.cycle. בינתיים אפשר להעתיק חתימות או לוותר על ה Type Hint במקרה כזה.