💬

【Python】関数内のローカル変数名を取得する

2022/11/05に公開

はじめに

どうやらpythonにおいてローカル変数の情報は関数オブジェクトが所持していないらしく、
スマートに取得する方法は見つかりませんでした。

2022.11.23追記
ありました
https://qiita.com/chankane/items/3909e9f2d1c5910cc60b
上の記事は関数の引数を取るものなので、以下のように書くとローカル変数だけを取得できます

関数名.__code__.co_varnames[関数名.__code__.co_argcount:]

上のコードでやりたかったことはできたので、これ以降の内容は見る必要ないです。

自分はどうしてもローカル変数名を取得したかったので、astモジュールを使って取得を行いましたが、今回のコードだと漏れが生じる可能性もあるので完璧ではないです。

環境

python 3.9.12
windows

astモジュールとは

抽象構文木を作成するモジュールです。詳しいことは私も理解していませんが、pythonプログラムの文字列からpythonで実行できるようなデータ形式への変換を行っているっぽいです。
https://docs.python.org/ja/3/library/ast.html

抽象構文木

pythonにおける1つ1つの命令をノードの木構造として表現したもの?
たとえば、以下のコードを実行する

import ast
import inspect


def hogefunc():
    x=1
    y=2
    return x

print(ast.dump(ast.parse(inspect.getsource(hogefunc)),indent=4))

出力

Module(
    body=[
        FunctionDef(
            name='hogefunc',
            args=arguments(
                posonlyargs=[],
                args=[],
                kwonlyargs=[],
                kw_defaults=[],
                defaults=[]),
            body=[
                Assign(
                    targets=[
                        Name(id='x', ctx=Store())],
                    value=Constant(value=1)),
                Assign(
                    targets=[
                        Name(id='y', ctx=Store())],
                    value=Constant(value=2)),
                Return(
                    value=Name(id='x', ctx=Load()))],
            decorator_list=[])],
    type_ignores=[])

関数hogefuncが構文木に変換されました。ここで、inspect.getsource()は関数を一度文字列に変換し、ast.parse()は文字列から構文木への変換、ast.dump()は構文木から文字列への変換を行っています。

ast.NodeVisitor

astモジュールで作成した構文木を(おそらく上から準に)たどることができます。
ast.NodeVisitorを継承したクラス内でvisit_{構文木ノードのクラス名}という関数を定義すると、特定のノードに到達したときに処理を実行することができます。
以下のコードを実行すると、

import ast
import inspect


def hogefunc():
    x=1
    y=2
    return x



class MyNodeVisitor(ast.NodeVisitor):

    def visit_Name(self, node: ast.Name) -> None:
        print(node.id)  
        self.generic_visit(node) #より深いノードの探索を行う場合はこれをかく

tree=ast.parse(inspect.getsource(hogefunc))  
visitor=MyNodeVisitor()
visitor.visit(tree)

出力

x
y
x

抽象構文木内の'Name'ノードのidを出力できました。

実装の方針

抽象構文木内の'Assign'(代入)ノード、NameEpr(代入式)ノードを洗い出し、ノードの中身を確認してローカル変数だろうと判断したものをリストに入れています。

また、global、nonlocal宣言があればその変数は別のリストに入れてローカル変数リストには入れないようにします。

関数内関数、関数内クラスについては新たなNodeVisitorを作成してローカル変数を取得します。

コード

"""
You can get local variables name of a function by using this.
"""
import inspect
import ast
from typing import Callable
import textwrap
def get_localvariables_name(f:Callable):
    """
    関数内のローカル変数名をリストで返す
    """

    class LocalvarNodeVisitor(ast.NodeVisitor):
        """抽象構文木を探索してローカル変数を見つけたらリストに入れてくれるクラス"""

        def __init__(self) -> None:
            self.local_vars_list=[] #ローカル変数のリスト
            self.nonlocal_vars_list=[] #global,nonlocal宣言された変数
            super().__init__()


        def append_var(self,name):
            """global,nonlocal宣言された変数でなければリストに追加"""
            if name not in self.nonlocal_vars_list:
                self.local_vars_list.append(name)
        
        def visit_Assign(self, node: ast.Assign) -> None:
            """代入(=)"""
            for target in node.targets: # x=y=19などの書き方だと node.targetsの要素が2つ以上になる
                 # Name,Tuple以外にもAttribute(Class.attr=value)やSubscript(list[num]=value)があるが
                 # それらはローカル変数を作成しないので排除
                if type(target) is ast.Name:
                    self.append_var(target.id)
                elif type(target) is ast.Tuple:
                    for element in target.elts:
                        self.append_var(element.id)

            self.generic_visit(node)

        def visit_NamedExpr(self, node: ast.NamedExpr) -> None:
            """代入式(:=)"""
            if type(tar:=node.target) is ast.Name:# python3.9.12では代入式にTuple,Attribute,Subscriptは使えないっぽいが一応書いておく
                self.append_var(tar.id)
            self.generic_visit(node)

        def visit_arguments(self, node: ast.arguments) ->None:
            """関数の引数もローカル変数に含める"""
            for arg in node.args: #普通の引数
                self.append_var(arg.arg)
            for arg in node.posonlyargs: #位置専用パラメータ
                self.append_var(arg.arg)
            for arg in node.kwonlyargs: #キーワード専用パラメータ
                self.append_var(arg.arg)
            if node.vararg is not None: # キーワード引数のtuple化
                self.append_var(node.vararg.arg)
            if node.kwarg is not None: # キーワード引数の辞書化
                self.append_var(node.kwarg.arg)
            self.generic_visit(node)
        

        def visit_Global(self, node: ast.Global) -> None:
            """global宣言された変数を取得"""
            for name in node.names:
                self.nonlocal_vars_list.append(name)
            self.generic_visit(node)
        def visit_Nonlocal(self, node: ast.Nonlocal) -> None:
            """nonlocal宣言された変数を取得(これはなくてもいいかも)"""
            for name in node.names:
                self.nonlocal_vars_list.append(name)
            self.generic_visit(node)

        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            self.append_var(node.name) #関数名やクラス名もローカル変数としてリストに入れる
            self.get_childlocal(node)
        def visit_ClassDef(self, node: ast.ClassDef) -> None:
            self.append_var(node.name)
            self.get_childlocal(node)
        def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
            self.append_var(node.name)
            self.get_childlocal(node)

        def get_childlocal(self,node:ast.AST):
            """
            スコープが変わればglobal,nonlocalの宣言が及ばなくなるので
            スコープが変わるノードに到達したら新しいvisitorを作成する
            """
            visitor= LocalvarNodeVisitor()
            visitor.generic_visit(node)
            self.local_vars_list+=visitor.local_vars_list

    
    class InitNodeVisitor(ast.NodeVisitor):
        """
        抽象構文木から最初のFunctionDefノードを取り出すクラス
        これとvisitor.generic_visitを使って引数の関数名はリストに含めずに関数内関数はリストに入れるようにする
        """
        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            self.funcdef=node
        def get_funcnode(self,tree:ast.AST)->ast.FunctionDef:
            self.visit(tree)
            return self.funcdef
    
    text=inspect.getsource(f) #関数のソースコードを文字列に変換
    text=textwrap.dedent(text) # 文字列を左寄せ
    tree=ast.parse(text) #抽象構文木へ変換
    tree=InitNodeVisitor().get_funcnode(tree) # 一番上の階層のFunctionDefノードをみつける
    visitor=LocalvarNodeVisitor()
    visitor.generic_visit(tree) # 構文木を探索させる(generic_visitは引数の下の階層から探索を始めているっぽい)
    return list(set(visitor.local_vars_list)) #リストからsetで重複を削除してリストで返す setのせいで順番はぐちゃぐちゃになる
 
    
    

テスト

g=0
a1=10
def test(a,b1="aaa",**b4):
    a0=100
    d=100
    d=9
    global g
    g=11


    (o,k)=(u,i)=j=(d,100)

    if (e:=100)>1000:
        h=100

    f=12
    def test2(s):
        nonlocal f
        f=0
        g=0
        s=s*s
        a=1
        return s
    
    x=test2(1)

    test3=test2


print(get_localvariables_name(test))

結果

['u', 'j', 'b4', 'g', 'test2', 'test3', 'o', 's', 'x', 'a0', 'd', 'a', 'b1', 'i', 'e', 'k', 'f', 'h']

set()を使っているため出力される変数名の順番は毎回変わります。
gはtest2()のなかで代入を行っているのでローカル変数として判定されています。

Discussion