Add get_origin

Author: msullivan

gel/_internal/_qb/__init__.py

@@ -75,6 +75,7 @@
     SortAlias,
     VarAlias,
     is_pointer_descriptor,
+    get_origin,
 )
 
 from ._protocols import (
@@ -170,6 +171,7 @@
     "empty_set_if_none",
     "exprmethod",
     "get_object_type_splat",
+    "get_origin",
     "is_expr_compatible",
     "is_pointer_descriptor",
     "toplevel_edgeql",

gel/_internal/_qb/_generics.py

@@ -10,6 +10,7 @@
     TYPE_CHECKING,
     Any,
     TypedDict,
+    TypeVar,
     NoReturn,
     get_args,
 )
@@ -375,6 +376,17 @@ def __edgeql__(self) -> tuple[type, tuple[str, dict[str, object]]]:
         return type_, toplevel_edgeql(self, splat_cb=splat_cb)
 
 
+_T = TypeVar("_T")
+
+
+def get_origin(x: type[_T] | BaseAlias) -> type[_T]:
+    return (
+        x.__gel_origin__
+        if isinstance(x, BaseAlias)  # type: ignore[return-value]
+        else x
+    )
+
+
 class PathAlias(BaseAlias):
     pass
 

gel/_internal/_qbmodel/_abstract/_syntax.py

@@ -24,9 +24,10 @@ def for_(iterator: type[_T], body: Callable[[type[_T]], type[_X]]) -> type[_X]:
     iter_expr = _qb.edgeql_qb_expr(iterator)
     scope = _qb.Scope()
     var = _qb.Variable(type_=iter_expr.type, scope=scope)
-    body_ = body(_qb.AnnotatedVar(iterator.__gel_origin__, var))  # type: ignore [arg-type, attr-defined]
+    t = _qb.get_origin(iterator)
+    body_ = body(_qb.AnnotatedVar(t, var))  # type: ignore [arg-type]
     return _qb.AnnotatedExpr(  # type: ignore [return-value]
-        body_.__gel_origin__,  # type: ignore [attr-defined]
+        _qb.get_origin(body_),
         _qb.ForStmt(
             iter_expr=iter_expr,
             body=_qb.edgeql_qb_expr(body_),

pyproject.toml

@@ -126,7 +126,6 @@ module = [
     "gel._internal._qbmodel.*",
     "gel._internal._qbmodel._abstract.*",
     "gel._internal._qbmodel._pydantic.*",
-    "gel.qb",
     "tests.test_dataclass_extras",
     "tests.test_dis_bool",
     "tests.test_is_overload",

tests/test_qb.py

@@ -595,7 +595,7 @@ def test_qb_enum_01(self):
         self.assertEqual(e.name, "red")
 
     def test_qb_for_01(self):
-        from models.orm import std
+        from models.orm import default, std
 
         res = self.client.query(
             std.for_(
@@ -617,6 +617,17 @@ def test_qb_for_01(self):
 
         self.assertEqual(set(res), {11, 12, 21, 22})
 
+        res2 = self.client.query(
+            std.for_(
+                default.User,
+                lambda x: x.name,
+            )
+        )
+        self.assertEqual(
+            set(res2),
+            {'Alice', 'Zoe', 'Billie', 'Dana', 'Cameron', 'Elsa'},
+        )
+
 
 class TestQueryBuilderModify(tb.ModelTestCase):
     """This test suite is for data manipulation using QB."""