Order_by and group_by accept labels

Improve typing to accept labels in ordey_by mand group_by.

Change-Id: I33e5d6f64633d39a220108d412ef84d6478b25e6

1 files changed, 36 insertions, 0 deletions

diff --git a/test/ext/mypy/plain_files/common_sql_element.py b/test/ext/mypy/plain_files/common_sql_element.py
index af36c85ee..586a130d2 100644
--- a/test/ext/mypy/plain_files/common_sql_element.py
+++ b/test/ext/mypy/plain_files/common_sql_element.py
@@ -9,7 +9,9 @@ unions.
 
 from __future__ import annotations
 
+from sqlalchemy import asc
 from sqlalchemy import Column
+from sqlalchemy import desc
 from sqlalchemy import Integer
 from sqlalchemy import MetaData
 from sqlalchemy import select
@@ -19,6 +21,7 @@ from sqlalchemy import Table
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.orm import Mapped
 from sqlalchemy.orm import mapped_column
+from sqlalchemy.orm import Session
 
 
 class Base(DeclarativeBase):
@@ -79,6 +82,21 @@ reveal_type(stmt)
 
 stmt = stmt.where(e2)
 
+stmt2 = select(User.id).order_by("email").group_by("email")
+stmt2 = select(User.id).order_by("id", "email").group_by("email", "id")
+stmt2 = (
+    select(User.id).order_by(asc("id"), desc("email")).group_by("email", "id")
+)
+# EXPECTED_TYPE: Select[Tuple[int]]
+reveal_type(stmt2)
+
+stmt2 = select(User.id).order_by(User.id).group_by(User.email)
+stmt2 = (
+    select(User.id).order_by(User.id, User.email).group_by(User.email, User.id)
+)
+# EXPECTED_TYPE: Select[Tuple[int]]
+reveal_type(stmt2)
+
 
 receives_str_col_expr(User.email)
 receives_str_col_expr(User.email + "some expr")
@@ -92,3 +110,21 @@ receives_bool_col_expr(User.email == "x")
 receives_bool_col_expr(e2)
 receives_bool_col_expr(e2.label("x"))
 receives_bool_col_expr(user_table.c.email == "x")
+
+
+# query
+
+q1 = Session().query(User.id).order_by("email").group_by("email")
+q1 = Session().query(User.id).order_by("id", "email").group_by("email", "id")
+# EXPECTED_TYPE: RowReturningQuery[Tuple[int]]
+reveal_type(q1)
+
+q1 = Session().query(User.id).order_by(User.id).group_by(User.email)
+q1 = (
+    Session()
+    .query(User.id)
+    .order_by(User.id, User.email)
+    .group_by(User.email, User.id)
+)
+# EXPECTED_TYPE: RowReturningQuery[Tuple[int]]
+reveal_type(q1)