iTerm 2 unzoom on switch pane script

unzoom.py

#!/usr/bin/env python3.7
# Put this script in ~/.config/iterm/Scripts or whatever, then enable in the Scripts menu
# Map switch keys to "Invoke Script Function...": `unzoom_and_switch(direction: "above/below/left/right")` 
import iterm2

async def async_turn_off_menu_option(conn: iterm2.Connection, option: str):
    state = await iterm2.MainMenu.async_get_menu_item_state(conn, option)
    if state.checked:
        await iterm2.MainMenu.async_select_menu_item(conn, option)

async def main(conn: iterm2.Connection):
    app = await iterm2.async_get_app(conn)
    @iterm2.RPC
    async def unzoom_and_switch(direction: str):
        window = app.current_window
        tab = window.current_tab
        await async_turn_off_menu_option(conn, 'Maximize Active Pane')
        await tab.async_select_pane_in_direction(iterm2.NavigationDirection(direction))
    
    await unzoom_and_switch.async_register(conn)


iterm2.run_forever(main)