В принципе, вы можете напрямую изменить ast («абстрактное синтаксическое дерево») функции, хотя на практике это может быть довольно проблематично. В любом случае, вот как это сделать для вашего простого примера:
Это создает из источника ast и происходит от класса NodeTransformer
для модификации ast на месте. Преобразователь узла имеет общий метод посещения, который пересекает узел и его поддерево, делегирующее узлам конкретных посетителей в производных классах. Здесь мы меняем все имена np
на sp
, а затем меняем эти атрибуты на прежние np
, теперь sp
, которые пишутся по-разному. Вы должны добавить все такие различия к translate
dict.
Наконец, мы компилируем обратно из ast в объект кода и выполняем его, чтобы сделать измененную функцию доступной.
import ast, inspect
import numpy as np
import sympy as sp
def f(a):
return np.array([np.sin(a), np.cos(a)])
z = ast.parse(inspect.getsource(f))
translate = {'array': 'Array'}
class np_to_sp(ast.NodeTransformer):
def visit_Name(self, node):
if node.id=='np':
node = ast.copy_location(ast.Name(id='sp', ctx=node.ctx), node)
return node
def visit_Attribute(self, node):
self.generic_visit(node)
if node.value.id=='sp' and node.attr in translate:
fields = {k: getattr(node, k) for k in node._fields}
fields['attr'] = translate[node.attr]
node = ast.copy_location(ast.Attribute(**fields), node)
return node
np_to_sp().visit(z)
exec(compile(z, '', 'exec'))
x = sp.Symbol('x')
print(f(x))
Выход:
[sin(x), cos(x)]
ОБНОВЛЕНИЕ простое усовершенствование: изменение функций, вызываемых функцией:
import ast, inspect
import numpy as np
import sympy as sp
def f(a):
return np.array([np.sin(a), np.cos(a)])
def f2(a):
return np.array([1, np.sin(a)])
def f3(a):
return f(a) + f2(a)
translate = {'array': 'Array'}
class np_to_sp(ast.NodeTransformer):
def visit_Name(self, node):
if node.id=='np':
node = ast.copy_location(ast.Name(id='sp', ctx=node.ctx), node)
return node
def visit_Attribute(self, node):
self.generic_visit(node)
if node.value.id=='sp' and node.attr in translate:
fields = {k: getattr(node, k) for k in node._fields}
fields['attr'] = translate[node.attr]
node = ast.copy_location(ast.Attribute(**fields), node)
return node
from types import FunctionType
for fn in f3.__code__.co_names:
fo = globals()[fn]
if not isinstance(fo, FunctionType):
continue
z = ast.parse(inspect.getsource(fo))
np_to_sp().visit(z)
exec(compile(z, '', 'exec'))
x = sp.Symbol('x')
print(f3(x))
Печать:
[sin(x) + 1, sin(x) + cos(x)]