Support exact max of select inputs

piercefreeman · Dec 5, 2024 · d32891f · d32891f
1 parent 5d73592
commit d32891f
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/iceaxe/queries.py b/iceaxe/queries.py
@@ -131,6 +131,13 @@ def select(
         self, fields: tuple[T | Type[T], T2 | Type[T2]]
     ) -> QueryBuilder[tuple[T, T2], Literal["SELECT"]]: ...
 
+    @overload
+    def select(
+        # TypeVarTuples only match the typing of one-or-more elements, so we also
+        # need a overloaded signature for 3 elements.
+        self, fields: tuple[T | Type[T], T2 | Type[T2], T3 | Type[T3]]
+    ) -> QueryBuilder[tuple[T, T2, T3], Literal["SELECT"]]: ...
+
     @overload
     def select(
         self, fields: tuple[T | Type[T], T2 | Type[T2], T3 | Type[T3], *Ts]
@@ -144,12 +151,14 @@ def select(
             | Type[T]
             | tuple[T | Type[T]]
             | tuple[T | Type[T], T2 | Type[T2]]
+            | tuple[T | Type[T], T2 | Type[T2], T3 | Type[T3]]
             | tuple[T | Type[T], T2 | Type[T2], T3 | Type[T3], *Ts]
         ),
     ) -> (
         QueryBuilder[T, Literal["SELECT"]]
         | QueryBuilder[tuple[T], Literal["SELECT"]]
         | QueryBuilder[tuple[T, T2], Literal["SELECT"]]
+        | QueryBuilder[tuple[T, T2, T3], Literal["SELECT"]]
         | QueryBuilder[tuple[T, T2, T3, *Ts], Literal["SELECT"]]
     ):
         """