Search code examples
pythonpython-3.xiterationaccessor

Is there a way to disable iteration for classes that define a __getitem__ method without putting constraints on the key?

The class that acts as my starting point is as follows:

class Test:
    def __getitem__(self, key):
        global frame
        frame = inspect.currentframe()
        if key > 9:
            raise KeyError
        return key

My thought was to use frame.f_back to discover that an iterator was automatically created for an instance as in the following example:

for x in Test():
    x

After running both of these and looking at frame.f_back, it was not obvious if __getitem__ can get enough information to detect if it is being called from whatever "iterator" that is interacting with it. The easiest solution would be to make the container start at one instead of zero to access its contents or maybe even to force wrapping the key in a list before passing it to the function as shown here:

>>> class Test:
    def __getitem__(self, key):
        if not isinstance(key, list):
            raise TypeError
        if len(key) != 1:
            raise ValueError
        key = key.pop()
        if not isinstance(key, int):
            raise TypeError
        if not 0 <= key < 10:
            raise KeyError
        return key


>>> for x in Test():
    x


Traceback (most recent call last):
  File "<pyshell#42>", line 1, in <module>
    for x in Test():
  File "<pyshell#39>", line 4, in __getitem__
    raise TypeError
TypeError
>>> Test()[[5]]
5
>>>

Is there a way that __getitem__ can know that it is being used automatically as an iterator and raise an exception to prevent such usage?

Related: Why does defining __getitem__ on a class make it iterable in python?

Solution

  • If the goal is to stop the object from being an iterable, you can just force an error on __iter__ method:

    class Test:
    def __getitem__(self, key):
        if key > 9:
            raise KeyError
        return key
    def __iter__(self):
        raise TypeError('Object is not iterable')

    Test run:

    >>> t = Test()
>>> for x in t:
    print(x)

Traceback (most recent call last):
  File "<pyshell#126>", line 1, in <module>
    for x in t:
  File "<pyshell#122>", line 7, in __iter__
    raise TypeError('Object is not iterable')
TypeError: Object is not iterable

    But __getitem__ will still work:

    >>> t[0]
0
>>> t[1]
1
>>> t[10]
Traceback (most recent call last):
  File "<pyshell#129>", line 1, in <module>
    t[10]
  File "<pyshell#122>", line 4, in __getitem__
    raise KeyError
KeyError