Adam Venis Adam Venis - 1 month ago 7
Python Question

Can a python function know when it's being called by a list comprehension?

I want to make a python function that behaves differently when it's being called from a list comprehension:

def f():
# this function returns False when called normally,
# and True when called from a list comprehension
pass

>>> f()
False
>>> [f() for _ in range(3)]
[True, True, True]


I tried looking at the inspect module, the dis module, and lib2to3's parser for something to make this trick work, but haven't found anything. There also might be a simple reason why this cannot exist, that I haven't thought of.

Answer

You can determine this by inspecting the stack frame in the following sort of way:

def f():
    try:
        raise ValueError
    except Exception as e:
        if e.__traceback__.tb_frame.f_back.f_code.co_name == '<listcomp>':
            return True

Then:

>>> print(f())
None
>>> print([f() for x in range(10)])
[True, True, True, True, True, True, True, True, True, True]

Its not to be recommended though. Really, its not.

NOTE

As it stands this only detects list comprehensions as requested. It will not detect the use of a generator. For example:

>>> print(list(f() for x in range(10)))
[None, None, None, None, None, None, None, None, None, None]
Comments