一日一技：限定Python函数只能被特定函数调用

摄影：产品经理

产品经理做的朝鲜冷面

有时候，我们的某些函数可能要限制调用。例如函数 A 只能被函数 B、函数 C 调用，不能被其他函数调用。

这并不是一个假想的场景，而是实实在在的场景。比如说，某些函数的传入条件非常苛刻，必须经过前置函数做周密的边界条件检查才能调用。不能让其他人随意调用。但当一个项目活久了以后，开发者或者新的接手者就会忘记这些限制，于是直接调用这些函数，导致出现问题。

但是我们知道，作为一个动态语言，Python 原生是没有这样的功能的。即使是加了双下划线的函数或者方法，也可以强行在另外的模块在调用。

但是我们可以通过分析函数的调用栈来通过代码解决这个问题。查询调用栈，可以使用inspect模块的stack()函数。我们来看看运行效果：

我们写一段测试代码：

import inspect


def foot_up():
    print('双脚离地')
    stack = inspect.stack()
    stop = 1


def jump():
    print('既然是跳舞，肯定要跳起来')
    foot_up()


def dance():
    print('开始跳舞')
    jump()


dance()

代码运行效果如下图所示：

可以看到，inspect.stack()返回了一个列表，列表里面有很多的FrameInfo对象。其中，第0项的.function属性对应的是当前下断点的这个函数自身；第1项的.function属性对应的是调用当前函数的函数名。如下图所示：

显然，如果我们判断 stack列表第1项的 function是不是我们允许的函数名，就可以决定要不要继续运行函数了。我们加一个判断试一试：

是 jump 函数调用，允许执行

不是 jump 函数调用，拒绝执行

如果你有很多个函数需要限制调用，显然每次都这样写会很麻烦，于是我们可以改写成装饰器：

def call_stack_check(valid_function_list=None):
    def decorate(func):
        def wrap(*args, **kwargs):
            if valid_function_list:
                stack = inspect.stack()
                upper_function_name = stack[1].function
                if upper_function_name not in valid_function_list:
                    raise Exception('不是 jump 函数调用，拒绝执行后续代码！')
                else:
                    print('是 jump 函数调用，允许执行后续代码')
            result = func(*args, **kwargs)
            return result
        return wrap
    return decorate

这个装饰器允许接收一个列表参数，如果这个列表不为空，那么只有列表中的函数名可以调用被装饰的函数。如果被列表外的函数调用了，程序就自动抛出一个异常。运行效果如下图所示：

正常执行

拒绝执行，抛出异常

以后，当你有一个函数需要限定调用者的时候，只需要用这个装饰器装饰它就可以了。装饰器的参数就是允许发起调用的函数名列表。