Add a for_ operator (#887)

Following Elvis's original branch, I added it to a `gel.qb` module.
I also, on Elvis's suggestion, modified `std` to re-export everything
from `gel.qb`.

I renamed his `foreach` to `for_` to better match EdgeQL.

Fixes #729.

---------

Co-authored-by: Elvis Pranskevichus <[email protected]>

Author: msullivan

gel/_internal/_codegen/_models/_pydantic.py

@@ -48,6 +48,8 @@
 from gel._internal._collections_extras import ImmutableChainMap
 from gel._internal._namespace import ident, dunder
 from gel._internal._qbmodel import _abstract as _qbmodel
+from gel._internal._qbmodel._abstract import _syntax as _qbsyntax
+
 from gel._internal._reflection._enums import SchemaPart, TypeModifier
 from gel._internal._schemapath import SchemaPath
 from gel._internal._polyfills import _strenum
@@ -2273,6 +2275,15 @@ def prepare_namespace(self, mod: IntrospectedModule) -> None:
                     self.py_file.update_globals(
                         gt.name for gt in GENERIC_TYPES
                     )
+                    names = _qbsyntax.__all__
+                    self.py_file.export(*names)
+                    self.py_file.update_globals(names)
+
+                    self.write("# Re-export top-level query-builder functions")
+                    self.write(
+                        "from gel._internal._qbmodel._abstract._syntax "
+                        "import *"
+                    )
 
                 self.py_file.update_globals(
                     ident(t.schemapath.name)

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/_expressions.py

@@ -457,7 +457,7 @@ def __edgeql_expr__(self, *, ctx: ScopeContext) -> str:
                 rexpr=item,
                 type_=SchemaPath("std", "bool"),
             )
-        return f"FILTER {edgeql(fexpr, ctx=ctx)}"
+        return f"FILTER {edgeql_exprstmt(fexpr, ctx=ctx)}"
 
 
 class OrderDirection(_strenum.StrEnum):
@@ -529,7 +529,7 @@ def __edgeql_expr__(self, *, ctx: ScopeContext) -> str:
                 type_=SchemaPath("std", "bool"),
             )
 
-        return f"ORDER BY {edgeql(dexpr, ctx=ctx)}"
+        return f"ORDER BY {edgeql_exprstmt(dexpr, ctx=ctx)}"
 
 
 @dataclass(kw_only=True, frozen=True)
@@ -548,7 +548,7 @@ def __edgeql_expr__(self, *, ctx: ScopeContext) -> str:
         if isinstance(self.limit, IntLiteral) and self.limit.val == 1:
             return "LIMIT 1"
         else:
-            clause = edgeql(self.limit, ctx=ctx)
+            clause = edgeql_exprstmt(self.limit, ctx=ctx)
             return f"LIMIT {clause}"
 
 
@@ -564,7 +564,7 @@ def precedence(self) -> _edgeql.Precedence:
         return _edgeql.PRECEDENCE[_edgeql.Token.OFFSET]
 
     def __edgeql_expr__(self, *, ctx: ScopeContext) -> str:
-        return f"OFFSET {edgeql(self.offset, ctx=ctx)}"
+        return f"OFFSET {edgeql_exprstmt(self.offset, ctx=ctx)}"
 
 
 @dataclass(kw_only=True, frozen=True)
@@ -726,7 +726,7 @@ def _body_edgeql(self, ctx: ScopeContext) -> str:
 
 
 @dataclass(kw_only=True, frozen=True)
-class ForStmt(IteratorExpr):
+class ForStmt(IteratorExpr, Stmt):
     stmt: _edgeql.Token = _edgeql.Token.FOR
     iter_expr: Expr
     body: Expr
@@ -745,12 +745,19 @@ def subnodes(self) -> Iterable[Node]:
         return (self.iter_expr, self.body)
 
     def _edgeql(self, ctx: ScopeContext) -> str:
+        ctx.bind(self.var)
         return (
             f"{self.stmt} {edgeql(self.var, ctx=ctx)} IN "
             f"({edgeql(self.iter_expr, ctx=ctx)})\n"
             f"UNION ({edgeql(self.body, ctx=ctx)})"
         )
 
+    def _iteration_edgeql(self, ctx: ScopeContext) -> str:
+        raise AssertionError('...')
+
+    def _body_edgeql(self, ctx: ScopeContext) -> str:
+        raise AssertionError('...')
+
 
 class Splat(_strenum.StrEnum):
     STAR = "*"
@@ -908,6 +915,17 @@ def get_object_type_splat(cls: type[GelTypeMetadata]) -> Shape:
     return shape
 
 
+def edgeql_exprstmt(
+    source: ExprCompatible,
+    *,
+    ctx: ScopeContext,
+) -> str:
+    res = edgeql(source, ctx=ctx)
+    if isinstance(source, Stmt):
+        res = f'({res})'
+    return res
+
+
 def toplevel_edgeql(
     x: ExprCompatible,
     *,

gel/_internal/_qb/_generics.py

@@ -10,6 +10,7 @@
     TYPE_CHECKING,
     Any,
     TypedDict,
+    TypeVar,
     NoReturn,
     get_args,
 )
@@ -315,9 +316,6 @@ def __infix_op__(
         *,
         swapped: bool = False,
     ) -> Any:
-        if op == "__eq__" and operand is self:
-            return True
-
         # Check for None comparison and raise appropriate error
         if operand is None and op in {"__eq__", "__ne__"}:
             _raise_op_error(_Op.IS_NONE if op == "__eq__" else _Op.IS_NOT_NONE)
@@ -375,6 +373,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

@@ -0,0 +1,39 @@
+# SPDX-PackageName: gel-python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright Gel Data Inc. and the contributors.
+
+
+"""Query building constructs"""
+
+from typing import TypeVar
+from collections.abc import Callable
+
+from gel._internal import _qb
+from gel._internal._qbmodel import _abstract
+
+
+_T = TypeVar("_T", bound=_abstract.GelType)
+_X = TypeVar("_X", bound=_abstract.GelType)
+
+
+def for_(iterator: type[_T], body: Callable[[type[_T]], type[_X]]) -> type[_X]:
+    """Evaluate the expression returned by *body* for each element in *iter*.
+
+    This is the Pythonic representation of the EdgeQL FOR expression."""
+
+    iter_expr = _qb.edgeql_qb_expr(iterator)
+    scope = _qb.Scope()
+    var = _qb.Variable(type_=iter_expr.type, scope=scope)
+    t = _qb.get_origin(iterator)
+    body_ = body(_qb.AnnotatedVar(t, var))  # type: ignore [arg-type]
+    return _qb.AnnotatedExpr(  # type: ignore [return-value]
+        _qb.get_origin(body_),
+        _qb.ForStmt(
+            iter_expr=iter_expr,
+            body=_qb.edgeql_qb_expr(body_),
+            scope=scope,
+        ),
+    )
+
+
+__all__ = ("for_",)

gel/_internal/ruff.toml

@@ -24,7 +24,6 @@ extend-select = [
     "PIE", # flake8-pie
     "PL",  # pylint
     "PYI", # flake8-pyi
-    "Q",   # flake8-quotes
     "RUF", # ruff specific
     "S",   # flake8-bandit
     "SIM", # flake8-simplify

gel/models/ruff.toml

@@ -24,7 +24,6 @@ extend-select = [
     "PIE", # flake8-pie
     "PL",  # pylint
     "PYI", # flake8-pyi
-    "Q",   # flake8-quotes
     "RUF", # ruff specific
     "S",   # flake8-bandit
     "SIM", # flake8-simplify

pyproject.toml

@@ -221,6 +221,7 @@ extend-ignore = [
     "E402", # module-import-not-at-top-of-file
     "E252", # missing-whitespace-around-parameter-equals
     "F541", # f-string-missing-placeholders
+    "Q000", # prefer double quotes
 ]
 
 [tool.ruff.format]

tests/test_qb.py

@@ -632,6 +632,61 @@ def test_qb_enum_01(self):
         self.assertEqual(e.color, default.Color.Red)
         self.assertEqual(e.name, "red")
 
+    def test_qb_for_01(self):
+        from models.orm import default, std
+
+        res = self.client.query(
+            std.for_(
+                std.range_unpack(std.range(std.int64(1), std.int64(10))),
+                lambda x: x * 2,
+            )
+        )
+        self.assertEqual(set(res), {i * 2 for i in range(1, 10)})
+
+        res = self.client.query(
+            std.for_(
+                std.range_unpack(std.range(std.int64(1), std.int64(3))),
+                lambda x: std.for_(
+                    std.range_unpack(std.range(std.int64(1), std.int64(3))),
+                    lambda y: x * 10 + y,
+                )
+            )
+        )
+
+        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'},
+        )
+
+        res3 = self.client.query(
+            default.User.filter(
+                lambda u: std.for_(
+                    std.assert_exists(std.int64(0)),
+                    # HMMMMM
+                    lambda x: x == x,
+                )
+            )
+        )
+        self.assertEqual(len(res3), 6)
+
+        res4 = self.client.query(
+            default.User.filter(
+                lambda u: std.for_(
+                    u.name,
+                    lambda x: x == "Alice",
+                )
+            )
+        )
+        self.assertEqual(len(res4), 1)
+
 
 class TestQueryBuilderModify(tb.ModelTestCase):
     """This test suite is for data manipulation using QB."""