From 853d787923b74b9eb4f1c74a156d2c8652910fce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebasti=C3=A1n=20Ram=C3=ADrez?= <tiangolo@gmail.com>
Date: Mon, 4 Dec 2023 10:46:59 +0100
Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20Refactor=20type=20generati?=
 =?UTF-8?q?on=20of=20selects=20re-order=20to=20prioritize=20models=20to=20?=
 =?UTF-8?q?optimize=20editor=20support=20(#718)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 scripts/generate_select.py        |   2 +-
 sqlmodel/sql/expression.py        | 346 +++++++++++++++---------------
 sqlmodel/sql/expression.py.jinja2 |   4 +-
 3 files changed, 176 insertions(+), 176 deletions(-)

diff --git a/scripts/generate_select.py b/scripts/generate_select.py
index 88e0e0a..8f01cb8 100644
--- a/scripts/generate_select.py
+++ b/scripts/generate_select.py
@@ -24,7 +24,7 @@ arg_groups: List[Arg] = []
 signatures: List[Tuple[List[Arg], List[str]]] = []
 
 for total_args in range(2, number_of_types + 1):
-    arg_types_tuples = product(["scalar", "model"], repeat=total_args)
+    arg_types_tuples = product(["model", "scalar"], repeat=total_args)
     for arg_type_tuple in arg_types_tuples:
         args: List[Arg] = []
         return_types: List[str] = []
diff --git a/sqlmodel/sql/expression.py b/sqlmodel/sql/expression.py
index a8a5725..2c931a1 100644
--- a/sqlmodel/sql/expression.py
+++ b/sqlmodel/sql/expression.py
@@ -328,12 +328,12 @@ _T3 = TypeVar("_T3")
 
 
 @overload
-def select(__ent0: _TScalar_0) -> SelectOfScalar[_TScalar_0]:  # type: ignore
+def select(__ent0: _TCCA[_T0]) -> SelectOfScalar[_T0]:
     ...
 
 
 @overload
-def select(__ent0: _TCCA[_T0]) -> SelectOfScalar[_T0]:
+def select(__ent0: _TScalar_0) -> SelectOfScalar[_TScalar_0]:  # type: ignore
     ...
 
 
@@ -342,17 +342,9 @@ def select(__ent0: _TCCA[_T0]) -> SelectOfScalar[_T0]:
 
 @overload
 def select(  # type: ignore
-    entity_0: _TScalar_0,
-    entity_1: _TScalar_1,
-) -> Select[Tuple[_TScalar_0, _TScalar_1]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
+    __ent0: _TCCA[_T0],
     __ent1: _TCCA[_T1],
-) -> Select[Tuple[_TScalar_0, _T1]]:
+) -> Select[Tuple[_T0, _T1]]:
     ...
 
 
@@ -366,9 +358,9 @@ def select(  # type: ignore
 
 @overload
 def select(  # type: ignore
-    __ent0: _TCCA[_T0],
+    entity_0: _TScalar_0,
     __ent1: _TCCA[_T1],
-) -> Select[Tuple[_T0, _T1]]:
+) -> Select[Tuple[_TScalar_0, _T1]]:
     ...
 
 
@@ -376,53 +368,16 @@ def select(  # type: ignore
 def select(  # type: ignore
     entity_0: _TScalar_0,
     entity_1: _TScalar_1,
-    entity_2: _TScalar_2,
-) -> Select[Tuple[_TScalar_0, _TScalar_1, _TScalar_2]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
-    entity_1: _TScalar_1,
-    __ent2: _TCCA[_T2],
-) -> Select[Tuple[_TScalar_0, _TScalar_1, _T2]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
-    __ent1: _TCCA[_T1],
-    entity_2: _TScalar_2,
-) -> Select[Tuple[_TScalar_0, _T1, _TScalar_2]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
-    __ent1: _TCCA[_T1],
-    __ent2: _TCCA[_T2],
-) -> Select[Tuple[_TScalar_0, _T1, _T2]]:
+) -> Select[Tuple[_TScalar_0, _TScalar_1]]:
     ...
 
 
 @overload
 def select(  # type: ignore
     __ent0: _TCCA[_T0],
-    entity_1: _TScalar_1,
-    entity_2: _TScalar_2,
-) -> Select[Tuple[_T0, _TScalar_1, _TScalar_2]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    __ent0: _TCCA[_T0],
-    entity_1: _TScalar_1,
+    __ent1: _TCCA[_T1],
     __ent2: _TCCA[_T2],
-) -> Select[Tuple[_T0, _TScalar_1, _T2]]:
+) -> Select[Tuple[_T0, _T1, _T2]]:
     ...
 
 
@@ -438,89 +393,9 @@ def select(  # type: ignore
 @overload
 def select(  # type: ignore
     __ent0: _TCCA[_T0],
-    __ent1: _TCCA[_T1],
-    __ent2: _TCCA[_T2],
-) -> Select[Tuple[_T0, _T1, _T2]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
-    entity_1: _TScalar_1,
-    entity_2: _TScalar_2,
-    entity_3: _TScalar_3,
-) -> Select[Tuple[_TScalar_0, _TScalar_1, _TScalar_2, _TScalar_3]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
-    entity_1: _TScalar_1,
-    entity_2: _TScalar_2,
-    __ent3: _TCCA[_T3],
-) -> Select[Tuple[_TScalar_0, _TScalar_1, _TScalar_2, _T3]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
     entity_1: _TScalar_1,
     __ent2: _TCCA[_T2],
-    entity_3: _TScalar_3,
-) -> Select[Tuple[_TScalar_0, _TScalar_1, _T2, _TScalar_3]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
-    entity_1: _TScalar_1,
-    __ent2: _TCCA[_T2],
-    __ent3: _TCCA[_T3],
-) -> Select[Tuple[_TScalar_0, _TScalar_1, _T2, _T3]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
-    __ent1: _TCCA[_T1],
-    entity_2: _TScalar_2,
-    entity_3: _TScalar_3,
-) -> Select[Tuple[_TScalar_0, _T1, _TScalar_2, _TScalar_3]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
-    __ent1: _TCCA[_T1],
-    entity_2: _TScalar_2,
-    __ent3: _TCCA[_T3],
-) -> Select[Tuple[_TScalar_0, _T1, _TScalar_2, _T3]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
-    __ent1: _TCCA[_T1],
-    __ent2: _TCCA[_T2],
-    entity_3: _TScalar_3,
-) -> Select[Tuple[_TScalar_0, _T1, _T2, _TScalar_3]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    entity_0: _TScalar_0,
-    __ent1: _TCCA[_T1],
-    __ent2: _TCCA[_T2],
-    __ent3: _TCCA[_T3],
-) -> Select[Tuple[_TScalar_0, _T1, _T2, _T3]]:
+) -> Select[Tuple[_T0, _TScalar_1, _T2]]:
     ...
 
 
@@ -529,38 +404,43 @@ def select(  # type: ignore
     __ent0: _TCCA[_T0],
     entity_1: _TScalar_1,
     entity_2: _TScalar_2,
-    entity_3: _TScalar_3,
-) -> Select[Tuple[_T0, _TScalar_1, _TScalar_2, _TScalar_3]]:
+) -> Select[Tuple[_T0, _TScalar_1, _TScalar_2]]:
     ...
 
 
 @overload
 def select(  # type: ignore
-    __ent0: _TCCA[_T0],
+    entity_0: _TScalar_0,
+    __ent1: _TCCA[_T1],
+    __ent2: _TCCA[_T2],
+) -> Select[Tuple[_TScalar_0, _T1, _T2]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
+    __ent1: _TCCA[_T1],
+    entity_2: _TScalar_2,
+) -> Select[Tuple[_TScalar_0, _T1, _TScalar_2]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
+    entity_1: _TScalar_1,
+    __ent2: _TCCA[_T2],
+) -> Select[Tuple[_TScalar_0, _TScalar_1, _T2]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
     entity_1: _TScalar_1,
     entity_2: _TScalar_2,
-    __ent3: _TCCA[_T3],
-) -> Select[Tuple[_T0, _TScalar_1, _TScalar_2, _T3]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    __ent0: _TCCA[_T0],
-    entity_1: _TScalar_1,
-    __ent2: _TCCA[_T2],
-    entity_3: _TScalar_3,
-) -> Select[Tuple[_T0, _TScalar_1, _T2, _TScalar_3]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    __ent0: _TCCA[_T0],
-    entity_1: _TScalar_1,
-    __ent2: _TCCA[_T2],
-    __ent3: _TCCA[_T3],
-) -> Select[Tuple[_T0, _TScalar_1, _T2, _T3]]:
+) -> Select[Tuple[_TScalar_0, _TScalar_1, _TScalar_2]]:
     ...
 
 
@@ -568,19 +448,9 @@ def select(  # type: ignore
 def select(  # type: ignore
     __ent0: _TCCA[_T0],
     __ent1: _TCCA[_T1],
-    entity_2: _TScalar_2,
-    entity_3: _TScalar_3,
-) -> Select[Tuple[_T0, _T1, _TScalar_2, _TScalar_3]]:
-    ...
-
-
-@overload
-def select(  # type: ignore
-    __ent0: _TCCA[_T0],
-    __ent1: _TCCA[_T1],
-    entity_2: _TScalar_2,
+    __ent2: _TCCA[_T2],
     __ent3: _TCCA[_T3],
-) -> Select[Tuple[_T0, _T1, _TScalar_2, _T3]]:
+) -> Select[Tuple[_T0, _T1, _T2, _T3]]:
     ...
 
 
@@ -598,9 +468,139 @@ def select(  # type: ignore
 def select(  # type: ignore
     __ent0: _TCCA[_T0],
     __ent1: _TCCA[_T1],
+    entity_2: _TScalar_2,
+    __ent3: _TCCA[_T3],
+) -> Select[Tuple[_T0, _T1, _TScalar_2, _T3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    __ent0: _TCCA[_T0],
+    __ent1: _TCCA[_T1],
+    entity_2: _TScalar_2,
+    entity_3: _TScalar_3,
+) -> Select[Tuple[_T0, _T1, _TScalar_2, _TScalar_3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    __ent0: _TCCA[_T0],
+    entity_1: _TScalar_1,
     __ent2: _TCCA[_T2],
     __ent3: _TCCA[_T3],
-) -> Select[Tuple[_T0, _T1, _T2, _T3]]:
+) -> Select[Tuple[_T0, _TScalar_1, _T2, _T3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    __ent0: _TCCA[_T0],
+    entity_1: _TScalar_1,
+    __ent2: _TCCA[_T2],
+    entity_3: _TScalar_3,
+) -> Select[Tuple[_T0, _TScalar_1, _T2, _TScalar_3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    __ent0: _TCCA[_T0],
+    entity_1: _TScalar_1,
+    entity_2: _TScalar_2,
+    __ent3: _TCCA[_T3],
+) -> Select[Tuple[_T0, _TScalar_1, _TScalar_2, _T3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    __ent0: _TCCA[_T0],
+    entity_1: _TScalar_1,
+    entity_2: _TScalar_2,
+    entity_3: _TScalar_3,
+) -> Select[Tuple[_T0, _TScalar_1, _TScalar_2, _TScalar_3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
+    __ent1: _TCCA[_T1],
+    __ent2: _TCCA[_T2],
+    __ent3: _TCCA[_T3],
+) -> Select[Tuple[_TScalar_0, _T1, _T2, _T3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
+    __ent1: _TCCA[_T1],
+    __ent2: _TCCA[_T2],
+    entity_3: _TScalar_3,
+) -> Select[Tuple[_TScalar_0, _T1, _T2, _TScalar_3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
+    __ent1: _TCCA[_T1],
+    entity_2: _TScalar_2,
+    __ent3: _TCCA[_T3],
+) -> Select[Tuple[_TScalar_0, _T1, _TScalar_2, _T3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
+    __ent1: _TCCA[_T1],
+    entity_2: _TScalar_2,
+    entity_3: _TScalar_3,
+) -> Select[Tuple[_TScalar_0, _T1, _TScalar_2, _TScalar_3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
+    entity_1: _TScalar_1,
+    __ent2: _TCCA[_T2],
+    __ent3: _TCCA[_T3],
+) -> Select[Tuple[_TScalar_0, _TScalar_1, _T2, _T3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
+    entity_1: _TScalar_1,
+    __ent2: _TCCA[_T2],
+    entity_3: _TScalar_3,
+) -> Select[Tuple[_TScalar_0, _TScalar_1, _T2, _TScalar_3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
+    entity_1: _TScalar_1,
+    entity_2: _TScalar_2,
+    __ent3: _TCCA[_T3],
+) -> Select[Tuple[_TScalar_0, _TScalar_1, _TScalar_2, _T3]]:
+    ...
+
+
+@overload
+def select(  # type: ignore
+    entity_0: _TScalar_0,
+    entity_1: _TScalar_1,
+    entity_2: _TScalar_2,
+    entity_3: _TScalar_3,
+) -> Select[Tuple[_TScalar_0, _TScalar_1, _TScalar_2, _TScalar_3]]:
     ...
 
 
diff --git a/sqlmodel/sql/expression.py.jinja2 b/sqlmodel/sql/expression.py.jinja2
index f1a2541..9bf93e1 100644
--- a/sqlmodel/sql/expression.py.jinja2
+++ b/sqlmodel/sql/expression.py.jinja2
@@ -273,12 +273,12 @@ _T{{ i }} = TypeVar("_T{{ i }}")
 # Generated TypeVars end
 
 @overload
-def select(__ent0: _TScalar_0) -> SelectOfScalar[_TScalar_0]:  # type: ignore
+def select(__ent0: _TCCA[_T0]) -> SelectOfScalar[_T0]:
     ...
 
 
 @overload
-def select(__ent0: _TCCA[_T0]) -> SelectOfScalar[_T0]:
+def select(__ent0: _TScalar_0) -> SelectOfScalar[_TScalar_0]:  # type: ignore
     ...