In python, Two functions calls some same workflows, how to distinguish the origin function at bottom?

Question

I have two python functions A() and B(), they both call f1(), and f1() calls f2()... In f4(), I want to get the origin function name(A or B), Is there any smart way? I thought I can add a param for each functions, but it makes my codes quite ugly. These functions are in different files and some start in new Thread, I also read some doc about context but I'm not sure if it works here.

def A():
    f1()

def B():
    f1()


def f1():
    f2()

def f2():
    f3()

def f3():
    f4()

def f4():

    if call_from_A:
        print(123)
    else:
        print(456)

Answer 1

You can use the traceback module to do this. It will give you the whole call stack, which you can parse to get the information you need.

Example:

import traceback


def A():
    f1()

def B():
    f1()


def f1():
    f2()

def f2():
    f3()

def f3():
    f4()

def f4():
    tb = traceback.extract_stack()
    call_chain = [f.name for f in tb]
    if 'A' in call_chain:
        print('Called from A')
    elif 'B' in call_chain:
        print('Called from B')
    
        
def main():
    A()
    B()
    
        
if __name__ == '__main__':
    main()

Answer 2

We can use inspect from Python to identify this.

import inspect


def f4():
    curframe = inspect.currentframe()
    calframe = inspect.getouterframes(curframe, 2)
    call_from_A = calframe[4].function == "A"

    if call_from_A:
        print(123)
    else:
        print(456)

Note that calframe[4] in the line call_from_A = calframe[4].function == "A" has to be modified appropriately if the nesting of your functions is changed. In other words, [4] is hard-coded to "go back in the frame stack 4 times" (i.e. f4 <- f3 <- f2 <- f1 <-source).

Answer 3

Your question boils down to "what kind of workflow am I in?" and the traditional way to solve such problems is using object-oriented patterns. Using object-oriented design properly often results in the elimination of conditionals.

Put the methods on a class. You can either pass the "kind of thing it is" at instantiation, or you can model that using inheritance, depending on your needs.

Using a single class and passing the "kind" at instantiation:

class Workflow:

    def __init__(self, kind):
        self.kind = kind

    def f1(self):
        self.f2()

    def f2(self):
        self.f3()

    def f3(self):
        self.f4()

    def f4(self):
        print(self.kind)

 # now use it

 a = Workflow(123)
 b = Workflow(456)

 a.f1()     # prints 123
 b.f1()     # prints 456

Notice there's no if involved. There doesn't need to be. Something like this would be missing the point:

def f4(self):
    if self.kind == "A":
        print(123)
    else if self.kind == "B":
        print(456)

The distinction is made earlier, when the object is instantiated. Thereafter each instance has its own behavior which is not predicated on anything.

If the two workflows have more complex behavior (e.g. f4() actually does some significant work) then you probably want inheritance. Define a base class, then subclass it for each kind of thing you want to be able to create.

 class WorkflowBase:

    def f1(self):
        self.f2()

    def f2(self):
        self.f3()

    def f3(self):
        self.f4()

    def f4(self):
        return

class WorkflowA(WorkflowBase):

    def f4(self):
         print(123)

class WorkflowB(WorkflowBase):

    def f4(self):
         print(456)

# using them

a = WorkflowA()
b = WorkflowB()

a.f1()     # prints 123
b.f1()     # prints 456

Here the behavior is baked into either WorkflowA or WorkflowB and the fact that you have an instance of one class or the other is what activates the differing result.

It also does not require you to inspect the call stack, which is quite a complicated condition. That is fun to do, I understand the allure well, but it is slow-ish, and it makes your code difficult for others to understand.

Answer 4

Although inspecting the Frames in the stack will allow one to find ot where the call came from, this is a somewhat expensive solution -these are objects created lazily from Python 3.11.

This is one of the problems context vars can help solve: one can set a value in a context var, that will be unique for all functions called on that "branch" of calls, as they will share the same contextvar context.

The only requisite is that the outermost function you want to track have to be called with contextvars.copy_context().run()

This can be achieved with a decorator, so that each functions can be called as usual, but internally it works as the entry point for a context. ContextVars, on the other hand, are a bit bureaucratic to deal with, requiring a set call.

All in all, this is how it would like with contextvars:

import contextvars
from functools import wraps

Value = contextvars.ContextVar("Value")



def isolate_context(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return contextvars.copy_context().run(func, *args, **kwargs)
    return wrapper


@isolate_context
def A():
    Value.set(123)
    f1()
    B()

@isolate_context
def B():
    Value.set(456)
    f1()


def f1():
    f2()

def f2():
    f3()

def f3():
    f4()

def f4():
    print(Value.get())


A()
B()