[Python] 함수 정의시 화살표(->)와 변수 설정시 콜론(:)

최근 transformers github에 있는 BART모델 코드를 볼 일이 있었는데,
이때 함수를 정의하는데 화살표(->)와 콜론(:)이 사용되는 것을 보고 어떤 기능을 하는지 몰라서 알아보게 되었다.

아래는 BartEncoderLayer 클래스의 forward 함수에서 갖고온 코드이다.

def forward(
    self,
    hidden_states: torch.FloatTensor,
    attention_mask: torch.FloatTensor,
    layer_head_mask: torch.FloatTensor,
    output_attentions: Optional[bool] = False,
) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]:

화살표(->)의 의미

python3에서는 함수 정의시 화살표(->)를 사용하여 함수의 리턴 값에 대한 type을 주석 형태로 제공할 수 있다고 한다.

콜론(:)의 의미

콜론 역시 python3에서 함수의 변수 설정시 변수의 type을 주석 형태로 제공할 수 있고,
output_attentions: Optional[bool] = False,을 보면 만약 함수 실행시 output_attentions의 값을 넣어주지 않았을 때 초기값을 False로 설정한다는 것을 의미한다.

---

코드로 좀 더 알아보자

테스트 코드

def test(x: int=1)->int:
    print(type(x))
    return f"x is {x}"

print(test())           # test1
print(test(2))          # test2
print(test(1.1))        # test3

실행 결과

# in test1
>>> <class 'int'>
>>> x is 1

# in test2
>>> <class 'int'>
>>> x is 2

# in test3
>>> <class 'float'>
>>> x is 1.1

test1

위에서 설명 했듯이 함수를 호출 할 때 변수를 넣어주지 않아도 초기값을 자동적으로 설정해주는 것을 볼 수 있다.

test2

만약 입력값이 있다면 들어온 입력값에 대해 작동한다.

test3

하지만 결국 콜론과 화살표 모두 주석이기 때문에 다른 type의 값이 들어와도 그대로 수행하는 것을 볼 수 있다.