From 04b7cb08ad90c39d643840f2ad1dadbf41dac4d9 Mon Sep 17 00:00:00 2001
From: Jay Qi <jayqi@users.noreply.github.com>
Date: Tue, 5 Mar 2024 22:58:10 -0500
Subject: [PATCH] Add example.py and typecheck it

---
 example.py                   | 32 ++++++++++++++++++++++++++++++++
 pyproject.toml               |  5 +++--
 sortedcontainers_pydantic.py | 15 ++++++++++-----
 3 files changed, 45 insertions(+), 7 deletions(-)
 create mode 100644 example.py

diff --git a/example.py b/example.py
new file mode 100644
index 0000000..6d2895c
--- /dev/null
+++ b/example.py
@@ -0,0 +1,32 @@
+from typing import TYPE_CHECKING, Tuple
+
+from pydantic import BaseModel
+
+from sortedcontainers_pydantic import SortedDict, SortedList, SortedSet
+
+
+class MyModel(BaseModel):
+    sorted_dict: SortedDict[str, int]
+    sorted_list: SortedList[float]
+    sorted_set: SortedSet[Tuple[int, int]]
+
+
+m = MyModel(
+    sorted_dict={"charlie": 2, "alpha": 3, "bravo": 1},
+    sorted_list=[3.0, 2.0, 1.0],
+    sorted_set={(2, 9), (3, 5), (1, 3)},
+)
+print(repr(m))
+#> MyModel(sorted_dict=SortedDict({'alpha': 3, 'bravo': 1, 'charlie': 2}), sorted_list=SortedList([1.0, 2.0, 3.0]), sorted_set=SortedSet([(1, 3), (2, 9), (3, 5)]))
+print(m.model_dump_json())
+#> {"sorted_dict":{"alpha":3,"bravo":1,"charlie":2},"sorted_list":[1.0,2.0,3.0],"sorted_set":[[1,3],[2,9],[3,5]]}
+
+m2 = MyModel.model_validate_json(m.model_dump_json())
+assert m == m2
+print(repr(m2))
+#> MyModel(sorted_dict=SortedDict({'alpha': 3, 'bravo': 1, 'charlie': 2}), sorted_list=SortedList([1.0, 2.0, 3.0]), sorted_set=SortedSet([(1, 3), (2, 9), (3, 5)]))
+
+if TYPE_CHECKING:
+    reveal_type(m.sorted_dict)  # noqa: F821
+    reveal_type(m.sorted_list)  # noqa: F821
+    reveal_type(m.sorted_set)  # noqa: F821
diff --git a/pyproject.toml b/pyproject.toml
index 196bb4e..073b3fe 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,7 @@ path = "sortedcontainers_pydantic.py"
 
 [tool.ruff]
 line-length = 99
-src = ["*.py"]
+exclude = ["example.py"]
 
 [tool.ruff.lint]
 select = [
@@ -54,8 +54,9 @@ known-first-party = ["sortedcontainers_pydantic"]
 force-sort-within-sections = true
 
 [tool.mypy]
-files = ["sortedcontainers_pydantic.py"]
+files = ["sortedcontainers_pydantic.py", "example.py"]
 ignore_missing_imports = true
+plugins = ["pydantic.mypy"]
 
 [tool.pytest.ini_options]
 minversion = "6.0"
diff --git a/sortedcontainers_pydantic.py b/sortedcontainers_pydantic.py
index 9a565b3..0beebb6 100644
--- a/sortedcontainers_pydantic.py
+++ b/sortedcontainers_pydantic.py
@@ -1,5 +1,5 @@
-from collections.abc import Mapping
-from typing import Any, Iterable, Set, Tuple, get_args
+from collections.abc import Hashable, Mapping
+from typing import Any, Iterable, Set, Tuple, TypeVar, get_args
 
 from pydantic import (
     GetCoreSchemaHandler,
@@ -9,8 +9,13 @@
 
 __version__ = "1.0.0"
 
+_KT = TypeVar("_KT", bound=Hashable)  # Key type.
+_VT = TypeVar("_VT")  # Value type.
+_T = TypeVar("_T")
+_HashableT = TypeVar("_HashableT", bound=Hashable)
 
-class SortedDict(sortedcontainers.SortedDict):
+
+class SortedDict(sortedcontainers.SortedDict[_KT, _VT]):
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Any, handler: GetCoreSchemaHandler
@@ -67,7 +72,7 @@ def __get_pydantic_core_schema__(
         )
 
 
-class SortedList(sortedcontainers.SortedList):
+class SortedList(sortedcontainers.SortedList[_T]):
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Any, handler: GetCoreSchemaHandler
@@ -116,7 +121,7 @@ def __get_pydantic_core_schema__(
         )
 
 
-class SortedSet(sortedcontainers.SortedSet):
+class SortedSet(sortedcontainers.SortedSet[_HashableT]):
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Any, handler: GetCoreSchemaHandler