Temporarily replacing list elements in Python

Purpose

I suddenly felt like doing something unusual: iterating through a loop while temporarily replacing elements in a Python list. Since I found a way to achieve this, I'm writing it down as a memo.

Implementation

Referring to contextlib — Utilities for with-statement contexts, the implementation could look like the following. It allows for temporarily replacing a single element or multiple elements at once.

from __future__ import annotations
import contextlib
from typing import Any
from collections.abc import Sequence


@contextlib.contextmanager
def temporarily_replaced_with(
    arr: list[Any],
    index: int | Sequence[int],
    by: Any | Sequence[Any]
):
    if isinstance(index, Sequence):
        backups = [arr[i] for i in index]
    else:
        backups = [arr[index]]
        index = [index]
        by = [by]

    try:
        for i, v in zip(index, by):
            arr[i] = v
        yield arr
    finally:
        for i, v in zip(index, backups):
            arr[i] = v

Experiments

• I want to temporarily triple the value at a specific position:

import numpy as np

a = np.arange(10)

for i in range(1, len(a)):
    with temporarily_replaced_with(a, i, by=a[i]*3):
        print(a)
print('original:', a)

[0 3 2 3 4 5 6 7 8 9]
[0 1 6 3 4 5 6 7 8 9]
[0 1 2 9 4 5 6 7 8 9]
[ 0 1 2 3 12 5 6 7 8 9]
[ 0 1 2 3 4 15 6 7 8 9]
[ 0 1 2 3 4 5 18 7 8 9]
[ 0 1 2 3 4 5 6 21 8 9]
[ 0 1 2 3 4 5 6 7 24 9]
[ 0 1 2 3 4 5 6 7 8 27]
original: [0 1 2 3 4 5 6 7 8 9]

• I want to temporarily swap the value at a specific position with its neighbor:

a = np.arange(10)

for i in range(1, len(a)-1):
    with temporarily_replaced_with(a, [i, i+1], [a[i+1], a[i]]):
        print(a)
print('original:', a)

[0 2 1 3 4 5 6 7 8 9]
[0 1 3 2 4 5 6 7 8 9]
[0 1 2 4 3 5 6 7 8 9]
[0 1 2 3 5 4 6 7 8 9]
[0 1 2 3 4 6 5 7 8 9]
[0 1 2 3 4 5 7 6 8 9]
[0 1 2 3 4 5 6 8 7 9]
[0 1 2 3 4 5 6 7 9 8]
original: [0 1 2 3 4 5 6 7 8 9]

It seems to be working as expected.

Summary

I don't usually use @contextlib.contextmanager very often, so this might be my second time, or maybe even my first. However, if you keep the idea that "such a feature exists" in the back of your mind, it might just click when you need it, even if that opportunity only comes once in a decade.