From d5033e658cdc2ca1ed89942ac8dd10893c627f91 Mon Sep 17 00:00:00 2001 From: Seth Howes Date: Tue, 1 Jul 2025 12:15:28 +0100 Subject: [PATCH] refactor: Replace Literal with Enum in sources.py --- shared/types/models/sources.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/shared/types/models/sources.py b/shared/types/models/sources.py index e00e8650..927a6ee2 100644 --- a/shared/types/models/sources.py +++ b/shared/types/models/sources.py @@ -1,10 +1,15 @@ from typing import Annotated, Any, Generic, Literal, TypeVar, Union, final +from enum import Enum from pydantic import AnyHttpUrl, BaseModel, Field, TypeAdapter from shared.types.models.common import ModelId -SourceType = Literal["HuggingFace", "GitHub"] + +class SourceType(str, Enum): + HuggingFace = "HuggingFace" + GitHub = "GitHub" + T = TypeVar("T", bound=SourceType) @@ -28,14 +33,14 @@ class GitHubModelSourceData(BaseModel): @final -class HuggingFaceModelSource(BaseModelSource[Literal["HuggingFace"]]): - source_type: Literal["HuggingFace"] = "HuggingFace" +class HuggingFaceModelSource(BaseModelSource[SourceType.HuggingFace]): + source_type: Literal[SourceType.HuggingFace] = SourceType.HuggingFace source_data: HuggingFaceModelSourceData @final -class GitHubModelSource(BaseModelSource[Literal["GitHub"]]): - source_type: Literal["GitHub"] = "GitHub" +class GitHubModelSource(BaseModelSource[SourceType.GitHub]): + source_type: Literal[SourceType.GitHub] = SourceType.GitHub source_data: GitHubModelSourceData