restore yank-last-arg and yank-nth-arg in ipython 5

01-esc-dot.py

# ~/.ipython/profile_default/startup/01-esc-dot.py
from IPython import get_ipython
from prompt_toolkit.enums import DEFAULT_BUFFER
from prompt_toolkit.keys import Keys
from prompt_toolkit.filters import HasFocus, HasSelection, ViInsertMode, EmacsInsertMode

ip = get_ipython()
insert_mode = ViInsertMode() | EmacsInsertMode()


class State(object):
    def __init__(self):
        self.reset_state()

    def reset_state(self):
        self.depth = 0
        self.arg = None

state = State()


def yank_last_arg(event):
    """'readline'-style yank-last-arg:
        Insert last argument to the previous command (the last word of the previous history entry).
        Successive calls to yank-last-arg move back through the history list, inserting the last
        word of each line in turn.
        see: https://www.gnu.org/software/bash/manual/bashref.html#Commands-For-History
        note: this doesn't support the numeric argument option of readline's yank-last-arg
    """
    b = event.current_buffer
    hist = ip.history_manager.input_hist_raw
    if state.depth == 0:
        state.depth = 1
    else:
        state.depth += 1

    if not event._arg is None:
        state.arg = event._arg

    arg = state.arg is None and -1 or state.arg
    if len(hist) and len(hist) >= state.depth:
        lastline = hist[-state.depth]

        if state.depth > 1:
            b.undo()

        b.save_to_undo_stack()
        if len(lastline.split()):
            b.insert_text(lastline.split()[arg])


if getattr(ip, 'pt_cli'):
    registry = ip.pt_cli.application.key_bindings_registry
    registry.add_binding(Keys.Escape, u'.',
                     filter=(HasFocus(DEFAULT_BUFFER)
                             & ~HasSelection()
                             & insert_mode))(yank_last_arg)
    registry.add_binding(Keys.Escape, u'_',
                     filter=(HasFocus(DEFAULT_BUFFER)
                             & ~HasSelection()
                             & insert_mode))(yank_last_arg)
    ip.events.register('post_execute', state.reset_state)