Skip to content

Instantly share code, notes, and snippets.

@jdahm
Last active May 21, 2020 19:52
Show Gist options
  • Save jdahm/dd3d4772eba79540d6d983c608012247 to your computer and use it in GitHub Desktop.
Save jdahm/dd3d4772eba79540d6d983c608012247 to your computer and use it in GitHub Desktop.
Exact value from function via AST visitor
import ast
from collections import namedtuple
import inspect
class ExtractValue(ast.NodeVisitor):
def __init__(self, func, *, call=None, keyword=None):
self.source = inspect.getsource(func)
self.call = call
self.keyword = keyword
def extract(self):
return self.visit(ast.parse(self.source))
def visit_Call(self, node):
def get_key(n):
return n.arg
def get_value(n):
if type(n.value) == ast.Num:
return str(n.value.n)
else:
raise
if node.func.id != self.call:
return
# Check if a call or object
if type(node.args[0]) == ast.Call:
call = node.args[0]
split_source = self.source.split('\n')
line_source = split_source[call.lineno-1][call.col_offset-1+1:]
# find index after first paren
index = line_source.find('(') + 1
# find matching paren
depth = 1
while depth > 0:
if line_source[index] == '(':
depth += 1
elif line_source[index] == ')':
depth -= 1
index += 1
self.value = eval("{}.{}".format(line_source[:index], self.keyword))
def extract_value(func, *, call=None, keyword=None):
v = ExtractValue(func, call=call, keyword=keyword)
return v.extract()
Test = namedtuple('Test', field_names='max')
def test():
return get_max(Test(max=2))
print(extract_value(test, call='get_max', keyword='max'))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment