How to intercept a variable in a function from another function

Question

Suppose I had a function like this in test.py:

from math import sqrt

def a():
    intermediate_val = sqrt(4)
    return 5

def b():
    another_val = sqrt(9)
    return 8

I want to write a function that looks at both a() and b() and returns the result of any call made to sqrt() without modifying the original code (decorators would be fine). Something like this:

import test

def intercept_value(fnc, intercept_fnc):
    # What goes here?

intercept_value('a', 'sqrt') == 2  # True
intercept_value('b', 'sqrt') == 3  # True

You could replace `sqrt` with a version which logs all values passed through it, and then returns the result of `sqrt`, then you'd have a list of all arguments used in the call to sqrt. — will, Apr 15 '15 at 22:34
I want to be able to do it without modifying any of the original code. In particular, I have a project that has unittest test cases. I'm trying to write a script that goes through each test and intercept the results of any http requests made (which are always made by self.client.get(), self.client.post(), etc). — Ryan Norton, Apr 15 '15 at 22:37
If you want it to be specifically inside function calls, you could use the `inspect` module to get the name of the function in the calling frame. — will, Apr 15 '15 at 22:37
I don't think you can do it without modifying the code, as you have to change the way the function calls work. — will, Apr 15 '15 at 22:38
Could I use ast to parse the source code and find the node that is a function call to self.client.get(), then see the name of the variable that it is being assigned to, then have a decorator that runs the actual function and returns the value of the variable name I found before the function returns? — Ryan Norton, Apr 15 '15 at 22:41
I mean, you could sure, but parsing source code is always an ugly option. I've done it in the past though. You probably don't need to use `ast` to achieve it though. It sounds like you know how to achieve what you want, and are just hoping there's a cleaner way - maybe put your attempt in your question? — will, Apr 15 '15 at 22:43
Instead of trying to modify the functions that call self.client.get(), etc, why not modify (i.e. decorate) those functions themselves? — 1.618, Apr 15 '15 at 22:56
i'm assuming `test1` is `a` and `test2` is `b`, but can you fix that? — abcd, Apr 16 '15 at 02:35

Ajax1234 · Answer 1 · 2021-11-20T00:00:45.540

As this is both an interesting question and tagged abstract-syntax-tree, here is a solution that will locate all function definitions and any routines called in the body of those functions and then evaluate the sub function:

import ast, collections
class Parse:
    def __init__(self):
       self.f_subs = collections.defaultdict(dict)
    def memoize(self, sig):
       if all(isinstance(i, ast.Constant) for i in sig.args) and all(isinstance(i.value, ast.Constant) for i in sig.keywords):
           return {'args':[i.value for i in sig.args], 'keywords':{i.arg:i.value.value for i in sig.keywords}}
    def walk(self, tree, f = None):
       if isinstance(tree, ast.FunctionDef):
           f = tree.name
       elif isinstance(tree, ast.Call) and f is not None:
           if (m:=self.memoize(tree)) is not None:
              self.f_subs[f][tree.func.id] = m
       for i in getattr(tree, '_fields', []):
           v = getattr(tree, i)
           for j in ([v] if not isinstance(v, list) else v):
              self.walk(j, f)

import test
def intercept_value(fnc, intercept_fnc):
    p = Parse()
    with open(test.__file__) as f:
       p.walk(ast.parse(f.read()))
       return getattr(test, intercept_fnc)(*(f:=p.f_subs[fnc][intercept_fnc])['args'], **f['keywords'])
       
print(intercept_value('a', 'sqrt'))
print(intercept_value('b', 'sqrt'))

Output:

2
3

score 0 · Answer 2 · answered Apr 16 '15 at 02:44

Here is a simple way to do what you want. It has the downside of needing an adjustment to the source code. Another downside is that it only works for the outputs of functions that are assigned to variables -- e.g., it wouldn't work for sqrt in x = sum(sqrt(some_array)).

I'm sure it's not the best answer to your question, but it's a start.

Source code:

from math import sqrt

def a(intercept=None):
    intermediate_val = sqrt(4)
    if intercept is not None:
        return locals()[intercept]
    return 5

def b(intercept=None):
    another_val = sqrt(9)
    if intercept is not None:
        return locals()[intercept]
    return 8

Test code:

import test

def intercept_value(fnc, intercept_var):
    return fnc(intercept=intercept_var)

intercept_value(test.a, 'intermediate_val') == 2  # True
intercept_value(test.b, 'another_val') == 3  # True

How to intercept a variable in a function from another function

2 Answers2