[Python] 関数実行時に自身の名前と引数と戻り値を返すデコレータ

はじめに

プログラムの作成段階で動作検証のために引数などを表示する事があります。
その時に、いちいちprint()を書いていたのでは面倒です。
そのため、デコレータや関数で実行時の引数を表示する機能を作ってみようと思いました。
ただ、全ての場合に適用できるような万能なものは作れませんでした。
しかし、使い方を間違えなければ有用だと思うので公開してみます。
テストが万全ではないので、使用の際はご注意してください。

環境

実行環境

Python 3.11.0 (main, Oct 24 2022, 18:26:48) [MSC v.1933 64 bit (AMD64)]

コード

シンプルバージョン

必要最低限で難しい機能を使わない場合は、以下のようなデコレータで十分だと思います。

シンプル版

import functools

def printArgs(func):
  @functools.wraps(func)
  def wrapper(*args, **kwargs):
    result = func(*args, **kwargs)
    print(f"name: {func.__name__}, args:{args}, kwargs:{kwargs}, result: {result}")
    return result
  return wrapper

これを以下のようにすると、

テストコード

@printArgs
def t0(a, b=2, *args, **kwargs):
  return a + b + sum(args)

t0(1, 3, 5, 7, x="a", y=4)

下の結果が求まります。

実行結果

name: t0, args:(1, 3, 5, 7), kwargs:{'x': 'a', 'y': 4}, result: 16

ただ、この場合、引数の名前の情報が一部失われてしまいます。
一時的に表示させるだけなのでこれで十分なのですがもう少し高度なものが欲しいです。

inspectバージョン

inspectモジュールを使って実装したのが以下のコードです。

python

import inspect
import functools
import itertools

def encloseStr(x, cl='"', cr='"'):
  return f'{cl}{x}{cr}' if isinstance(x, str) else x

def getArgsInfo(func, context=1):
  argSpec = inspect.getfullargspec(unwrap(func))
  frame = inspect.getouterframes(inspect.currentframe())[context].frame
  argInfo = inspect.getargvalues(frame)
  info = getMergedArgsInfo(argSpec, argInfo)
  return info

def getVarargValue(varargs, argInfo):
  if varargs is not None and varargs in argInfo.locals:
    varargsValue = list(argInfo.locals[varargs])
  elif argInfo.varargs is not None and argInfo.varargs in argInfo.locals:
    varargsValue = list(argInfo.locals[argInfo.varargs])
  else:
    varargsValue = {}
  return varargsValue

def getVarkwValue(varkw, argInfo):
  if varkw is not None and varkw in argInfo.locals:
    varkwValue = argInfo.locals[varkw]
  elif argInfo.keywords is not None and argInfo.keywords in argInfo.locals:
    varkwValue = argInfo.locals[argInfo.keywords]
  else:
    varkwValue = {}
  return varkwValue

def getMergedArgsInfo(argSpec, argInfo):
  argDefaults, varargs, varkw = getArgDefaults(argSpec)
  varargsValue = getVarargValue(varargs, argInfo)
  varkwValue = getVarkwValue(varkw, argInfo)
  for k, v in argDefaults.items():
    if k == varargs:
      v["value"] = varargsValue
    elif k == varkw:
      v["value"] = varkwValue
    # elif argInfo.keywords is not None and argInfo.keywords in argInfo.locals and k in argInfo.locals[argInfo.keywords]:
    elif argInfo.keywords in argInfo.locals and k in argInfo.locals[argInfo.keywords]:
      v["value"] = argInfo.locals[argInfo.keywords][k]
    elif k in argInfo.locals:
      v["value"] = argInfo.locals[k]
    else:
      if len(varargsValue) > 0:
        v["value"] = varargsValue.pop(0)
      else:
        v["value"] = v["defaults"]
  return argDefaults

def getArgDefaults(argSpec):
  args = {}
  defaults = () if argSpec.defaults is None else argSpec.defaults
  argDefaults = dict(itertools.zip_longest(reversed(argSpec.args), reversed(defaults), fillvalue=inspect.Parameter.empty))
  for k in argSpec.args:
    args[k] = {"defaults": argDefaults[k], "kind": ""}
  if argSpec.kwonlyargs:
    for k in argSpec.kwonlyargs:
      args[k] = {"defaults": argSpec.kwonlydefaults.get(k, inspect.Parameter.empty), "kind": "*"}
  if argSpec.varargs:
    args[argSpec.varargs] = {"kind": "*args"}
  if argSpec.varkw:
    args[argSpec.varkw] = {"kind": "**kwargs"}
  return args, argSpec.varargs, argSpec.varkw

def formatMergedArgInfo(info):
  output = ""
  for k, v in info.items():
    value = encloseStr(v["value"])
    default = encloseStr(v.get("defaults", inspect.Parameter.empty))
    default = "" if default == inspect.Parameter.empty else f"({default})"
    kind = "" if v["kind"] == "" else encloseStr(v["kind"], "<", ">")
    output += f"{k}={value}{default}{kind}, "
  return output[:-2]

def printArgInfo(func, context=2):
  print(formatMergedArgInfo(getArgsInfo(func, context)))

def decoratorArgsInfo(func):
  @functools.wraps(func)
  def wrapper(*args, **kwargs):
    result = func(*args, **kwargs)
    print(func.__name__, result)
    printArgInfo(unwrap(func))
    return result
  return wrapper

def unwrap(func, depth=0):
  i = 0
  while hasattr(func, "__wrapped__"):
    func = func.__wrapped__
    i += 1
    if depth > 0 and i >= depth:
      return func
  return func

使い方

解説を後回しにして使い方は、getArgsInfo()を引数を表示したい関数内でそれ自身を引数にして呼び出します。
getArgsInfo()は、引数の情報が入ったDictonaryを返すので必要に応じてformatMergedArgInfo()などで整形します。
情報の表示させ方は、もう少し改良の余地があると思います。

print(formatMergedArgInfo(getArgsInfo()))のような長ったらしい書き方を省略したprintArgInfo()も用意してみました。

デコレータとして使いたい場合は、@decoratorArgsInfoでデコレートしてください。

test

def t1(a, b=2, *args, **kwargs):
  print(formatMergedArgInfo(getArgsInfo(t1)))
  return a + b + sum(args)

def t2(a, b=2, *args, **kwargs):
  printArgInfo(t2)
  return a + b + sum(args)

@decoratorArgsInfo
def t3(a, b=2, *args, **kwargs):
  return a + b + sum(args)

t1(1, 3, 5, 7, x="a", y="b")
t2(1, 3, 5, 7, x="a", y="b")
t3(1, 3, 5, 7, x="a", y="b")

出力は、以下のようになります。

a=1, b=3, args=[5, 7]<*args>, kwargs={'x': 'a', 'y': 'b'}<**kwargs>
a=1, b=3, args=[5, 7]<*args>, kwargs={'x': 'a', 'y': 'b'}<**kwargs>
t3 16
a=1, b=3, args=[5, 7]<*args>, kwargs={'x': 'a', 'y': 'b'}<**kwargs>

注意点

ここで以下のようなデコレータを用意します。
注目していただきたいのは、@functools.wrapsを使っていないデコレータである事です。

def meow(func):
  def wrapper(*args, **kwargs):
    result = func(*args, **kwargs)
    print(f"meow, meow")
    return result
  return wrapper

@meow
def t4(a, b=2, *args, **kwargs):
  printArgInfo(t4)
  return a + b + sum(args)

@decoratorArgsInfo
@meow
def t5(a, b=2, *args, **kwargs):
  return a + b + sum(args)

t4(1, 3, 5, 7, x="a", y=4)
print("----------")
t5(1, 3, 5, 7, x="a", y=4)

これを実行すると、以下のような情報が欠落した結果が得られます。

args=[5, 7]<*args>, kwargs={'x': 'a', 'y': 4}<**kwargs>
meow, meow
----------
meow, meow
wrapper 16
args=[1, 3, 5, 7]<*args>, kwargs={'x': 'a', 'y': 4}<**kwargs>

これは、デコレートされていないオリジナルの関数をgetArgsInfo()が必要としているからです。
一応、unwrap()でデコレータを外すようにはしています。
しかし、この関数が対応しているのは、@functools.wrapsを使ったデコレータのみです。
そのため、欠落した情報のみが得られます。

pipでundecoratedあたりを導入すればこの問題は解決できそうですが試していません。
自分で作るにしてもそれだけで記事が出来そうです。
今回は、これで妥協しています。

また、ほかの複雑な構造をしたデコレータには対応していないかもしれません。

解説

このコードの核となる関数は、以下のgetArgsInfo()です。

def getArgsInfo(func, context=1):
  argSpec = inspect.getfullargspec(unwrap(func))
  frame = inspect.getouterframes(inspect.currentframe())[context].frame
  argInfo = inspect.getargvalues(frame)
  info = getMergedArgsInfo(argSpec, argInfo)
  return info

始めに、inspect.getfullargspec()に関数オブジェクトを渡し、その引数情報を取得します。
上記のように、これに渡す関数は、デコレートされていないものが必要です。
そのため、unwrap()でデコレータを外しています。

inspect.getfullargspec()を実行すると以下のようなFullArgSpecというnamedtupleが返ってきます。

FullArgSpec(args=['a', 'b'], varargs='args', varkw='kwargs', defaults=(2,), kwonlyargs=[], kwonlydefaults=None, annotations={})

これには、引数名、デフォルト値、位置引数かキーワード引数かの情報が含まれています。
現在の値は分かりません。

そこで、現在の引数の値を得るために、inspect.getargvalues()を実行します。
これには、関数が動いているフレームを渡す必要があります。
フレームは、正確かわかりませんが、スコープと言い直しても良いかもしれません。
inspect.currentframe()でgetArgsInfo()が動いているフレームが得られます。
しかし、getArgsInfo()を呼び出している関数のフレームが必要なので、inspect.getouterframes(inspect.currentframe())[context].frameで1つ外側のフレームを取得しています。
必ず1つ外側と決まっているならば、inspect.currentframe().f_backで良いのですが、ほかの条件にも対応できるようにこうしています。

こうして得られたフレームでinspect.getargvalues()を実行すると以下のようなArgInfoというnamedtupleが返ってきます。

ArgInfo(args=['a', 'b'], varargs='args', keywords='kwargs', locals={'a': 1, 'b': 3, 'args': (5, 7), 'kwargs': {'x': 'a', 'y': 'b'}})

これには、引数名、現在の値、位置引数かキーワード引数かの情報が含まれています。
一見すると、これだけで目的が果たせそうです。
しかし、デコレータとして実装しようとすると、以下のようなラッパー関数の情報が得られてしまいます。

ArgInfo(args=[], varargs='args', keywords='kwargs', locals={'args': (1, 3, 5, 7), 'kwargs': {'x': 'a', 'y': 'b'}, 'result': 16, 'func': <function t3 at 0x000001E8540CE160>})

そのため、関数の情報である事が確実なinspect.getfullargspec()を基準として、inspect.getargvalues()と組み合わせる事で目的の情報を得る事にしています。

組み合わせる処理をgetMergedArgsInfo()で行っています。
getMergedArgsInfo()では、名前を照らし合わせて泥臭くDictonaryを作っているだけなので簡単な解説のみで済ませます。
注意すべき点が2つあります。
1つ目は、FullArgSpecに含まれているデフォルト値は、最後のn個だけで位置引数やキーワード引数(args)の個数とあっていない点です。
そのため、zip()で結合する時にreversed()してから結合しています。
2つ目は、デコレータ内から呼び出された場合は、引数が全てvarargsかkwargsに渡されてしまっている点です。
そのため、そのため、varargsの始めのn個を引数に渡して残りを本来のvarargsの値としています。
もっとも、デコレータの構造によっては、このような条件では上手くいかない場合があるかもしれません。

テストコード

def test1():
  def t0():
    printArgInfo(t0)
    return
  def t1(a, b, c):
    printArgInfo(t1)
    return a + b + c
  def t2(a, b=1, c=2):
    printArgInfo(t2)
    return a + b + c
  def t3(a, *, b, c=2):
    printArgInfo(t3)
    return a + b + c
  def t4(a, b, c, *args, **kgargs):
    printArgInfo(t4)
    return a + b + c + sum(args)
  t0()
  t1(1, 2, 3)
  t2(1)
  t2(2, 4, 6)
  t3(1, b=2)
  t3(1, b=2, c=4)
  t4(1, 2, 3, 4, 5, d=2, e=4)

def test2():
  @decoratorArgsInfo
  def t0():
    return
  @decoratorArgsInfo
  def t1(a, b, c):
    return a + b + c
  @decoratorArgsInfo
  def t2(a, b=1, c=2):
    return a + b + c
  @decoratorArgsInfo
  def t3(a, *, b, c=2):
    return a + b + c
  @decoratorArgsInfo
  def t4(a, b, c, *args, **kgargs):
    return a + b + c + sum(args)
  t0()
  t1(1, 2, 3)
  t2(1)
  t2(2, 4, 6)
  t3(1, b=2)
  t3(1, b=2, c=4)
  t4(1, 2, 3, 4, 5, d=2, e=4)
test1()
test2()