From 965cad76bb8519926d0389f703250b3d66ec763c Mon Sep 17 00:00:00 2001
From: Zoe Roux <zoe.roux@zoriya.dev>
Date: Sat, 10 May 2025 15:37:04 +0200
Subject: [PATCH] Map `Language` for pydantic

---
 scanner/scanner/models/collection.py |  4 +--
 scanner/scanner/models/entry.py      |  4 +--
 scanner/scanner/models/movie.py      |  4 +--
 scanner/scanner/models/season.py     |  4 +--
 scanner/scanner/models/serie.py      |  4 +--
 scanner/scanner/utils.py             | 51 ++++++++++++++++++++++++----
 6 files changed, 50 insertions(+), 21 deletions(-)

diff --git a/scanner/scanner/models/collection.py b/scanner/scanner/models/collection.py
index c3dcd132..1229357b 100644
--- a/scanner/scanner/models/collection.py
+++ b/scanner/scanner/models/collection.py
@@ -1,8 +1,6 @@
 from __future__ import annotations
 
-from langcodes import Language
-
-from ..utils import Model
+from ..utils import Language, Model
 from .genre import Genre
 from .metadataid import MetadataId
 
diff --git a/scanner/scanner/models/entry.py b/scanner/scanner/models/entry.py
index eb408f19..ae5cfe56 100644
--- a/scanner/scanner/models/entry.py
+++ b/scanner/scanner/models/entry.py
@@ -3,9 +3,7 @@ from __future__ import annotations
 from datetime import date
 from typing import Literal
 
-from langcodes import Language
-
-from ..utils import Model
+from ..utils import Language, Model
 from .metadataid import EpisodeId, MetadataId
 
 
diff --git a/scanner/scanner/models/movie.py b/scanner/scanner/models/movie.py
index 25776d36..756881bc 100644
--- a/scanner/scanner/models/movie.py
+++ b/scanner/scanner/models/movie.py
@@ -3,9 +3,7 @@ from __future__ import annotations
 from datetime import date
 from enum import StrEnum
 
-from langcodes import Language
-
-from ..utils import Model
+from ..utils import Language, Model
 from .collection import Collection
 from .genre import Genre
 from .metadataid import MetadataId
diff --git a/scanner/scanner/models/season.py b/scanner/scanner/models/season.py
index 72da60b9..be857a0b 100644
--- a/scanner/scanner/models/season.py
+++ b/scanner/scanner/models/season.py
@@ -2,9 +2,7 @@ from __future__ import annotations
 
 from datetime import date
 
-from langcodes import Language
-
-from ..utils import Model
+from ..utils import Language, Model
 from .metadataid import SeasonId
 
 
diff --git a/scanner/scanner/models/serie.py b/scanner/scanner/models/serie.py
index 05058a3f..4e50f895 100644
--- a/scanner/scanner/models/serie.py
+++ b/scanner/scanner/models/serie.py
@@ -3,9 +3,7 @@ from __future__ import annotations
 from datetime import date
 from enum import StrEnum
 
-from langcodes import Language
-
-from ..utils import Model
+from ..utils import Language, Model
 from .collection import Collection
 from .entry import Entry
 from .extra import Extra
diff --git a/scanner/scanner/utils.py b/scanner/scanner/utils.py
index 7c0127f5..0f015497 100644
--- a/scanner/scanner/utils.py
+++ b/scanner/scanner/utils.py
@@ -1,15 +1,16 @@
-from langcodes import Language
-from pydantic import AliasGenerator, BaseModel, ConfigDict
+from typing import Annotated, Any, Callable
+
+from langcodes import Language as BaseLanguage
+from pydantic import AliasGenerator, BaseModel, ConfigDict, GetJsonSchemaHandler
 from pydantic.alias_generators import to_camel
-
-
-def normalize_lang(lang: str) -> str:
-	return str(Language.get(lang))
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import core_schema
 
 
 def to_slug(title: str) -> str:
 	return title
 
+
 def clean(val: str) -> str | None:
 	return val or None
 
@@ -21,3 +22,41 @@ class Model(BaseModel):
 			serialization_alias=lambda x: to_camel(x[:-1] if x[-1] == "_" else x),
 		),
 	)
+
+
+class _LanguagePydanticAnnotation:
+	@classmethod
+	def __get_pydantic_core_schema__(
+		cls,
+		_source_type: Any,
+		_handler: Callable[[Any], core_schema.CoreSchema],
+	) -> core_schema.CoreSchema:
+		def validate_from_str(value: str) -> BaseLanguage:
+			return BaseLanguage.get(value)
+
+		from_str_schema = core_schema.chain_schema(
+			[
+				core_schema.str_schema(),
+				core_schema.no_info_plain_validator_function(validate_from_str),
+			]
+		)
+
+		return core_schema.json_or_python_schema(
+			json_schema=from_str_schema,
+			python_schema=core_schema.union_schema(
+				[
+					core_schema.is_instance_schema(BaseLanguage),
+					from_str_schema,
+				]
+			),
+			serialization=core_schema.to_string_ser_schema(),
+		)
+
+	@classmethod
+	def __get_pydantic_json_schema__(
+		cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+	) -> JsonSchemaValue:
+		return handler(core_schema.str_schema())
+
+
+Language = Annotated[BaseLanguage, _LanguagePydanticAnnotation]