Pydantic 모델에서 Union(|) 연산자로 상속 문제 해결하기
Pydantic을 사용하여 컨텐츠 요소(Content Element) 클래스를 설계할 때, 상속 구조에서 발생할 수 있는 한 가지 문제는 요소들이 부모 클래스(Base Class)로만 인식되는 경우다. 이를 해결하기 위해 | 연산자를 활용하는 방법을 소개한다.
문제 상황: 부모 클래스로만 인식되는 컨텐츠 요소
기본적으로 컨텐츠의 구성 요소는 공통된 속성을 가지며, 이를 기반으로 개별적인 요소(예: 설명, 대화 등)를 정의할 수 있다. 처음 구현한 코드에서 ContentElement를 부모 클래스로 두고, 이를 상속하는 Description과 Conversation 클래스를 만들었다. 그리고 Content 클래스의 elements 속성에 여러 요소를 포함하도록 설계했다.
from pydantic import BaseModel, Field
from typing import List
class ContentElement(BaseModel):
category: str = Field(title="요소 카테고리", description="컨텐츠 요소의 카테고리")
element_id: int = Field(title="요소 ID", description="컨텐츠 요소의 고유 ID")
class Description(ContentElement):
category: str = "description"
text: str = Field(title="설명 내용", description="설명의 텍스트 내용")
mood: str = Field(title="분위기", description="설명의 분위기")
class Conversation(ContentElement):
category: str = "conversation"
speaker: str = Field(title="화자", description="대화를 하는 인물의 이름")
text: str = Field(title="대화 내용", description="대화의 텍스트 내용")
mood: str = Field(title="분위기", description="대화의 분위기")
Python
복사
이후 Content 클래스에서 elements 필드의 타입을 List[ContentElement]로 지정했지만, 실제로 데이터를 변환할 때 모든 요소가 ContentElement로만 인식되는 문제가 발생했다.
class Content(BaseModel):
title: str = Field(title="컨텐츠 제목", description="컨텐츠의 제목")
scene_number: int = Field(title="씬 번호", description="컨텐츠의 씬 번호")
elements: List[ContentElement] = Field(title="컨텐츠 요소들", description="컨텐츠를 구성하는 요소들의 리스트")
Python
복사
이 코드에서는 elements 내부에 Description과 Conversation 객체가 들어가더라도 Pydantic이 이를 부모 클래스 ContentElement로만 처리했다. 즉, 데이터 변환 시 객체가 원래의 서브클래스가 아닌 부모 클래스의 속성만 유지하는 문제가 발생했다.
해결 방법: | 연산자로 다형성 처리
이 문제를 해결하기 위해 | 연산자를 사용하여 여러 개의 서브클래스를 명확하게 정의할 수 있다. 즉, elements의 타입을 List[ContentElement | Description | Conversation]으로 명시적으로 지정하면, Pydantic이 요소들을 개별적인 서브클래스로 인식할 수 있다.
from typing import Literal
class Content(BaseModel):
title: str = Field(title="컨텐츠 제목", description="컨텐츠의 제목")
scene_number: int = Field(title="씬 번호", description="컨텐츠의 씬 번호")
elements: List[ContentElement | Description | Conversation] = Field(title="컨텐츠 요소들", description="컨텐츠를 구성하는 요소들의 리스트")
Python
복사
이렇게 변경한 후 테스트를 진행하면, elements 리스트 내부의 객체들이 원래의 서브클래스(Description, Conversation) 타입을 유지하는 것을 확인할 수 있다.
content = Content(
title="예제 컨텐츠",
scene_number=1,
elements=[
Description(element_id=1, text="어둠 속에서 한 사람이 걸어온다.", mood="긴장감"),
Conversation(element_id=2, speaker="주인공", text="여긴 어디지?", mood="당황")
]
)
print(content.elements[0]) # Description 객체로 유지됨
print(content.elements[1]) # Conversation 객체로 유지됨
Python
복사
이제 elements 리스트 안에 들어가는 객체들이 더 이상 ContentElement로만 인식되지 않고, 각각의 서브클래스로 정확히 유지되는 것을 확인할 수 있다.
정리
Pydantic에서 다형성을 활용할 때, 기본적으로 리스트 내의 요소들이 부모 클래스 타입으로만 인식될 수 있는 문제가 있다. 이를 해결하기 위해 | 연산자를 사용하면, Pydantic이 올바른 서브클래스 타입을 인식하여 변환할 수 있도록 도와준다. 특히, 여러 타입의 객체를 포함하는 리스트를 다룰 때 | 연산자를 적극적으로 활용하면 더욱 직관적인 데이터 모델을 만들 수 있다.