with_unlocked.py

""" 
Examples:

>>> a = lock_array(np.zeros((5, 5)))   # locked array
>>> b = np.zeros((5, 5))               # default array

# >>> a[:, 3] = 1.   # You can't update locked arrays
>>> b[:, 3] = 1.

>>> with unlocked(a, b):
>>>    a[:, 2] = 1   # You can update unlocked arrays here
>>>    b[:, 2] = 1
      
"""

def lock_array(values):
    values.setflags(write=False)
    return values

def unlock_array(values):
    values.setflags(write=True)
    return values

class unlocked:
    def __init__(self, *values):
        for v in values:
            assert isinstance(v, np.ndarray)
        self._values = values

    def __enter__(self, *values):
        self._org_flags = [v.flags.writeable for v in self._values]
        for v in self._values:
            v.setflags(write=True)

    def __exit__(self, *args):
        for v, flag in zip(self._values, self._org_flags):
            v.setflags(write=flag)