pythonfunctional-programmingpython-3.8python-assignment-expression

Expression that returns mutated list


I'm looking for a single expression, mutates an element and returns the modified list

The following is a bit verbose

# key=0; value=3; rng=[1,2]
[(v if i != key else value) for i, v in enumerate(rng)]

Edit:

I'm looking for a way to inline the following function in a single expression

def replace(rng: List, key: int, value):
    a = list(rng) 
    a[key] = value
    return a

Edit 2: the code that actually motivated this question

class TextDecoder(nn.Module):
    def forward(self, x: Tensor, kv_cache: Tensor):
        kv_cache_write = torch.zeros((_:=list(kv_cache.shape))).__setitem__(-2, x.shape[-1]) or _)
        ...

Solution

  • Maybe better than list concatenation:

    [*rng[:key], value, *rng[key+1:]]
    

    Another:

    [a for a, a[key] in [(rng[:], value)]][0]
    

    Or if you are just assigning the result to a variable (as discussed) you can do:

    a, a[key] = rng[:], value