[해결 전략] [Python Type Hinting] NumPy 조건부 반환 타입 효율적으로 명시하기

작성자:

Python 개발자라면 코드의 가독성과 유지보수성을 높이기 위해 타입 힌트(Type Hint)를 적극 활용하고 계실 겁니다. 특히 NumPy와 같이 동적인 타입 흐름을 가진 라이브러리를 사용할 때, 함수의 반환 타입이 입력 매개변수에 따라 조건부로 달라지는 경우가 많습니다. 이때 typing.overload를 사용하여 이러한 복잡한 동작을 정확하게 타입 힌트하는 것이 쉬운 일은 아닙니다.

이번 글에서는 Stack Overflow에서 제기된 문제 상황을 통해, NumPy 배열과 float 값을 모두 처리하는 함수가 특정 조건에서만 float을 반환하고 그 외 모든 경우에 NumPy 배열을 반환할 때, 이 동작을 효율적이고 명확하게 타입 힌트하는 방법을 알아보겠습니다.

목차 숨기기
1 1. 문제 발생 상황
2 2. 명확한 발생 원인
3 3. 해결 방법 및 코드 예시
3.1 해결 아이디어
3.2 개선된 코드 예시
4 4. 향후 예방을 위한 팁
5 Related Posts:

1. 문제 발생 상황

다음과 같은 foo 함수가 있다고 가정해 봅시다. 이 함수는 네 개의 인자를 받아 연산한 결과를 반환합니다.

import numpy as np
from typing import Union

def foo(a: Union[float, np.ndarray], b: Union[float, np.ndarray],
        c: Union[float, np.ndarray], d: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
    return a * b + c * d

이 함수의 핵심 동작은 다음과 같습니다.

  • 모든 입력 매개변수(a, b, c, d)가 float일 경우에만 반환 타입이 float입니다.
  • 단 하나라도 np.ndarray가 포함되면, 반환 타입은 np.ndarray가 됩니다.

개발자는 이러한 동작을 타입 힌트하기 위해 @overload를 사용하려고 시도했습니다. 하지만 모든 가능한 조합을 고려하면 다음과 같이 비효율적인 오버로드 정의가 필요했습니다.

from typing import overload, Union
import numpy as np

# Float 또는 np.ndarray를 나타내는 타입 별칭
FloatOrArray = Union[float, np.ndarray]

@overload
def foo(a: float, b: float, c: float, d: float) -> float: ...

@overload
def foo(a: np.ndarray, b: FloatOrArray,
        c: FloatOrArray, d: FloatOrArray) -> np.ndarray: ...

@overload
def foo(a: FloatOrArray, b: np.ndarray,
        c: FloatOrArray, d: FloatOrArray) -> np.ndarray: ...

@overload
def foo(a: FloatOrArray, b: FloatOrArray,
        c: np.ndarray, d: FloatOrArray) -> np.ndarray: ...

@overload
def foo(a: FloatOrArray, b: FloatOrArray,
        c: FloatOrArray, d: np.ndarray) -> np.ndarray: ...

# ... 물론 모든 조합을 다 지정하려면 더 많아집니다.

이 방식은 4개의 인자만으로도 5개(실제로는 더 많은)의 오버로드 정의가 필요하며, 인자 수가 늘어날수록 기하급수적으로 증가하여 비효율적이고 가독성이 떨어집니다. 개발자는 “첫 번째 오버로드(모든 float 입력)가 만족되지 않으면 항상 배열을 반환한다”는 기본 동작을 타입 검사기에 효율적으로 알릴 방법을 찾고 있었습니다.

2. 명확한 발생 원인

문제의 핵심은 typing.overload 데코레이터의 작동 방식과 개발자의 오해에서 비롯됩니다. Python의 타입 검사기(예: Mypy)는 @overload로 정의된 함수 시그니처들을 위에서 아래로 순차적으로 평가합니다. 즉, 목록의 첫 번째 오버로드부터 시작하여 현재 호출과 일치하는 시그니처를 찾습니다. 일치하는 시그니처를 찾으면 해당 오버로드의 반환 타입을 사용하고 더 이상 아래의 오버로드를 검사하지 않습니다.

따라서, 개발자는 모든 np.ndarray가 포함된 경우를 개별적으로 명시하려고 했지만, 이는 @overload의 평가 순서와 “가장 구체적인 것부터 가장 일반적인 것” 순서로 정의해야 하는 원칙을 간과한 것입니다. 특정 조건(모든 입력이 float)에서만 float을 반환하고, 그 외 *모든 다른 경우*에는 np.ndarray를 반환해야 한다는 것은, float을 반환하는 경우가 가장 *구체적*이며, np.ndarray를 반환하는 경우가 가장 *일반적*이라는 것을 의미합니다.

개발자가 제시한 방식은 각기 다른 np.ndarray 조합을 모두 나열함으로써 중복을 발생시키고, @overload의 강력한 추론 기능을 제대로 활용하지 못했습니다. 타입 검사기는 “모든 입력이 float이 아닌 경우”를 포괄하는 하나의 일반적인 오버로드로도 충분히 추론할 수 있습니다.

3. 해결 방법 및 코드 예시

이 문제의 해결책은 의외로 간단합니다. @overload의 평가 순서 원칙을 따르면, 모든 입력이 float일 때 float을 반환하는 가장 구체적인 오버로드를 먼저 정의하고, 그 뒤에 나머지 모든 경우를 포괄하는 가장 일반적인 오버로드를 정의하는 것입니다. 이렇게 하면 수많은 조합을 일일이 나열할 필요 없이 단 두 개의 오버로드만으로 정확한 타입 힌트를 제공할 수 있습니다.

해결 아이디어

  • 가장 구체적인 경우: 모든 인자가 float이고 반환 타입이 float인 오버로드를 가장 먼저 정의합니다.
  • 가장 일반적인 경우: 모든 인자가 float 또는 np.ndarray이고 반환 타입이 np.ndarray인 오버로드를 그 다음에 정의합니다. 이 오버로드는 첫 번째 오버로드에 해당하지 않는 모든 나머지 경우를 “잡아내는(catch-all)” 역할을 합니다.

개선된 코드 예시

import numpy as np
from typing import overload, Union

# Float 또는 np.ndarray를 나타내는 타입 별칭 (가독성 향상)
FloatOrArray = Union[float, np.ndarray]

# 1. 가장 구체적인 오버로드: 모든 입력이 float일 때 float 반환
@overload
def foo(a: float, b: float, c: float, d: float) -> float: ...

# 2. 모든 다른 경우를 포괄하는 일반적인 오버로드:
#    첫 번째 오버로드에 해당하지 않으면(즉, 하나라도 np.ndarray가 있으면)
#    항상 np.ndarray 반환
@overload
def foo(a: FloatOrArray, b: FloatOrArray,
        c: FloatOrArray, d: FloatOrArray) -> np.ndarray: ...

# 실제 함수 구현 (런타임에는 이 시그니처가 사용됨)
def foo(a: FloatOrArray, b: FloatOrArray,
        c: FloatOrArray, d: FloatOrArray) -> FloatOrArray:
    return a * b + c * d

# --- 타입 검증 예시 (mypy와 함께 reveal_type 사용) ---
print("--- 타입 검증 결과 (mypy reveal_type) ---")

# 모든 입력이 float인 경우
res_float = foo(1.0, 2.0, 3.0, 4.0)
# mypy: Revealed type is "builtins.float"

# 하나라도 np.ndarray가 포함된 경우
res_array_1 = foo(np.array([1.0]), 2.0, 3.0, 4.0)
# mypy: Revealed type is "numpy.ndarray"

res_array_2 = foo(1.0, np.array([2.0]), 3.0, 4.0)
# mypy: Revealed type is "numpy.ndarray"

res_array_3 = foo(1.0, 2.0, np.array([3.0]), 4.0)
# mypy: Revealed type is "numpy.ndarray"

res_array_4 = foo(1.0, 2.0, 3.0, np.array([4.0]))
# mypy: Revealed type is "numpy.ndarray"

# 모든 입력이 np.ndarray인 경우
res_array_all = foo(np.array([1.0]), np.array([2.0]), np.array([3.0]), np.array([4.0]))
# mypy: Revealed type is "numpy.ndarray"

# (실제 실행 시 print는 타입 정보를 출력하지 않습니다. Mypy를 실행해야 볼 수 있습니다.)
# 예시: 터미널에서 `mypy your_script_name.py` 실행 후 출력 확인
# your_script_name.py:27: note: Revealed type is "builtins.float"
# your_script_name.py:31: note: Revealed type is "numpy.ndarray"
# your_script_name.py:34: note: Revealed type is "numpy.ndarray"
# your_script_name.py:37: note: Revealed type is "numpy.ndarray"
# your_script_name.py:40: note: Revealed type is "numpy.ndarray"
# your_script_name.py:43: note: Revealed type is "numpy.ndarray"

위 코드에서 볼 수 있듯이, 단 두 개의 오버로드만으로 복잡한 조건부 반환 타입 로직을 완벽하게 표현할 수 있습니다. 첫 번째 @overload는 가장 구체적인 `float` 케이스를 처리하고, 두 번째 @overload는 그 외 모든 `FloatOrArray` 조합에서 `np.ndarray`가 반환되도록 타입 검사기에 지시합니다.

4. 향후 예방을 위한 팁

복잡한 함수의 타입 힌트를 다룰 때 위와 같은 문제에 직면하지 않기 위한 몇 가지 팁입니다.

  1. @overload 사용 원칙 이해: @overload를 사용할 때는 항상 “가장 구체적인(Specific) 것부터 가장 일반적인(General) 것” 순서로 정의해야 한다는 원칙을 기억하세요. 타입 검사기는 이 순서대로 일치하는 시그니처를 찾습니다.
  2. “Catch-All” 오버로드 활용: 특정 조건에서만 특별한 타입이 반환되고, 그 외 모든 경우에는 하나의 기본 타입이 반환될 때, 마지막에 가장 일반적인 “catch-all” 오버로드를 정의하여 불필요한 반복을 피하세요.
  3. 타입 별칭(Type Alias) 사용: Union[float, np.ndarray]와 같이 반복적으로 사용되는 복잡한 타입은 FloatOrArray = Union[float, np.ndarray]와 같이 타입 별칭을 정의하여 코드의 가독성을 높이고 유지보수를 용이하게 하세요.
  4. 타입 검사기 활용: Mypy, Pyright와 같은 정적 타입 검사기를 개발 워크플로우에 통합하여 타입 힌트의 정확성을 지속적으로 검증하세요. reveal_type() 같은 기능을 활용하면 특정 변수의 추론된 타입을 직접 확인할 수 있어 디버깅에 큰 도움이 됩니다.
  5. 라이브러리 문서 참고: NumPy와 같이 타입 흐름이 복잡한 라이브러리를 사용할 때는 해당 라이브러리의 타입 힌트 관련 문서를 참고하는 것이 좋습니다. 간혹 라이브러리 자체에서 제공하는 특정 TypeVar나 유틸리티 타입이 있을 수 있습니다.

이러한 원칙들을 잘 적용하면 Python 코드의 타입 안정성을 높이고, 개발 생산성을 향상시키는 동시에 동료 개발자들이 코드를 더 쉽게 이해하고 협업할 수 있게 됩니다.

Related Posts:

댓글 남기기