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

Answer 1

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