既存クラスのメソッドにエラーハンドリングを追加する

この記事の内容

ライブラリなどに定義されているクラスの全ての関数に対し try...except でエラーを捕捉し、適切なエラーを返す方法を紹介します。

背景

PySparkを用いて開発する際、spark.sql.DataFrame のメソッド関係のエラーが頻発していました。最初はいちいち try...except で囲うことも考えたのですが、可読性の問題を考えて何か良い方法はないかと考えてこの方法に辿り着きました。

関数にエラーハンドリングを付与する

本題に入る前に、関数定義の際にエラーハンドリングを付加するための方法としてデコレータを用いた方法が取られます。

python

class MyException(Exception):
    def __init__(self, message = "An error occurred"):
        super().__init__(message)

def error_handling_wrapper(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            raise MyException(f"An error occurred in {func.__name__}: {e}") from e

    return wrapper

もちろん except Exception as e の部分は全てのエラーを補足してしまうので適切なエラーのみに書き換えるべきですが、今はこのまま話を進めます。この関数をデコレータとして使用すると例えば以下のようになります。

python

@error_handling_wrapper
def test():
    1/0

test()

# 前略
# MyException: An error occurred in test: division by zero

この関数をクラスのメソッドに適用していきます。

クラスメソッドへの適用

以下の decorate_all_methods_in_cls を使用することでエラーハンドリングをクラスの全ての関数に適用することができます。
例として pyspark.sql.DataFrame に適用しています。

python

from pyspark.sql import DataFrame

def decorate_all_methods_in_cls(decorator):
    def decorate(cls):
        for attr in dir(cls):
            if callable(getattr(cls, attr)) and not attr.startswith("_"):
                setattr(cls, attr, decorator(getattr(cls, attr)))
        return cls
    return decorate

DataFrame = decorate_all_methods_in_cls(error_handling_wrapper)(DataFrame)

全ての関数に適用、と言いましたが実際は _ から始まらない、public な使用を想定された関数に対して error_handling_wrapper を適用するような形になっています。これによってクラスに含まれる関数がエラーを出した時に例外処理が行われます。

使い道

背景の部分にも書いたように、PySpark を使用する際（厳密には Spark と Sedona で地理的操作をする際）にエラーが頻発しており、関連する操作すべてに try...except を使用するのは憚られたのでこのような対策を思いつきました。ただこんなことが起こるのはそもそもコードが悪い可能性もあるのでまずはコードの見直しが必要であると思います。
ライブラリ内で発生するエラーをカスタムエラーで統一的に処理したいなんて場合にもこの方法は使えそうです。
また今回紹介した例は DataFrame クラスに対し適用する形となっていますが、PySpark では DataFrame を直接生成することよりも csv やデータベースからの読み込みによって生成することのほうが多く、今回の手法は使えません。多分ラッパークラスとかを使うんだと思いますがまだよく理解していないのでまた理解したら記事にしたいです。

おわりに

今回紹介した方法とは別のベストプラクティスがあればぜひコメントでお教えいただければ幸いです。最後までお読みいただきありがとうございました。