Skip to content
Merged
8 changes: 4 additions & 4 deletions gel/_internal/_codegen/_models/_pydantic.py
Original file line number Diff line number Diff line change
Expand Up @@ -2909,15 +2909,15 @@ def _write_prefix_op_method_node_ctor(
) -> None:
"""Generate the query node constructor for a prefix operator method.

Creates the code that builds a PrefixOp query node for unary operator
Creates the code that builds a UnaryOp query node for unary operator
methods like __neg__. The operator is applied to 'self' as the
operand.

Args:
op: The operator reflection object containing metadata
"""
op_name = op.schemapath.name
node_cls = self.import_name(BASE_IMPL, "PrefixOp")
node_cls = self.import_name(BASE_IMPL, "UnaryOp")

args = [
"expr=self", # The operand is always 'self' for method calls
Expand Down Expand Up @@ -2967,15 +2967,15 @@ def _write_prefix_op_func_node_ctor(
) -> None:
"""Generate the query node constructor for a prefix operator function.

Creates the code that builds a PrefixOp query node for unary operator
Creates the code that builds a UnaryOp query node for unary operator
functions. Unlike method versions, this takes the operand from function
arguments and applies special type casting for tuple parameters.

Args:
op: The operator reflection object containing metadata
"""
op_name = op.schemapath.name
node_cls = self.import_name(BASE_IMPL, "PrefixOp")
node_cls = self.import_name(BASE_IMPL, "UnaryOp")
expr_compat = self.import_name(BASE_IMPL, "ExprCompatible")
cast_ = self.import_name("typing", "cast")

Expand Down
12 changes: 10 additions & 2 deletions gel/_internal/_qb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
AbstractFieldDescriptor,
Expr,
PathPrefix,
PathTypeIntersectionPrefix,
Scope,
ScopeContext,
Stmt,
Expand All @@ -36,14 +37,14 @@
IntLiteral,
Limit,
Literal,
ObjectWhenType,
Offset,
OrderBy,
OrderByElem,
OrderByExpr,
OrderDirection,
OrderEmptyDirection,
Path,
PrefixOp,
SchemaSet,
SelectStmt,
SetLiteral,
Expand All @@ -52,11 +53,13 @@
ShapeOp,
Splat,
StringLiteral,
UnaryOp,
UpdateStmt,
Variable,
construct_infix_op_chain,
empty_set,
empty_set_if_none,
expr_uses_auto_splat,
get_object_type_splat,
toplevel_edgeql,
)
Expand Down Expand Up @@ -95,6 +98,7 @@
GelSchemaMetadata,
GelSourceMetadata,
GelTypeMetadata,
GelObjectTypeExprMetadata,
GelObjectTypeMetadata,
)

Expand Down Expand Up @@ -128,6 +132,7 @@
"ForStmt",
"FuncCall",
"GelLinkMetadata",
"GelObjectTypeExprMetadata",
"GelObjectTypeMetadata",
"GelPointerReflection",
"GelReflectionProto",
Expand All @@ -141,6 +146,7 @@
"IntLiteral",
"Limit",
"Literal",
"ObjectWhenType",
"Offset",
"OrderBy",
"OrderByElem",
Expand All @@ -150,7 +156,7 @@
"Path",
"PathAlias",
"PathPrefix",
"PrefixOp",
"PathTypeIntersectionPrefix",
"SchemaSet",
"Scope",
"ScopeContext",
Expand All @@ -163,6 +169,7 @@
"Splat",
"Stmt",
"StringLiteral",
"UnaryOp",
"UpdateStmt",
"VarAlias",
"Variable",
Expand All @@ -171,6 +178,7 @@
"edgeql_qb_expr",
"empty_set",
"empty_set_if_none",
"expr_uses_auto_splat",
"exprmethod",
"get_object_type_splat",
"get_origin",
Expand Down
18 changes: 13 additions & 5 deletions gel/_internal/_qb/_abstract.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@

if TYPE_CHECKING:
from collections.abc import Iterable, Iterator, Mapping
from gel._internal._schemapath import TypeName
from gel._internal._schemapath import TypeNameExpr


@dataclass(kw_only=True, frozen=True)
Expand Down Expand Up @@ -89,7 +89,7 @@ class Expr(Node):
def precedence(self) -> _edgeql.Precedence: ...

@abc.abstractproperty
def type(self) -> TypeName: ...
def type(self) -> TypeNameExpr: ...

@abc.abstractmethod
def __edgeql_expr__(self, *, ctx: ScopeContext) -> str: ...
Expand All @@ -100,10 +100,10 @@ def __edgeql_qb_expr__(self) -> Self:

@dataclass(kw_only=True, frozen=True)
class TypedExpr(Expr):
type_: TypeName
type_: TypeNameExpr

@property
def type(self) -> TypeName:
def type(self) -> TypeNameExpr:
return self.type_


Expand Down Expand Up @@ -494,12 +494,20 @@ def compute_must_bind_refs(
return (self,)


@dataclass(frozen=True, kw_only=True)
class PathTypeIntersectionPrefix(IdentLikeExpr):
type_filter: TypeNameExpr

def __edgeql_expr__(self, *, ctx: ScopeContext) -> str:
return f"[is {self.type_filter.as_quoted_schema_name()}]"


@dataclass(kw_only=True, frozen=True)
class ImplicitIteratorStmt(IteratorExpr, Stmt):
"""Base class for statements that are implicit iterators"""

@property
def type(self) -> TypeName:
def type(self) -> TypeNameExpr:
return self.iter_expr.type

@property
Expand Down
74 changes: 59 additions & 15 deletions gel/_internal/_qb/_expressions.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@

from gel._internal import _edgeql
from gel._internal._polyfills import _strenum
from gel._internal._schemapath import SchemaPath, TypeName
from gel._internal._schemapath import SchemaPath, TypeName, TypeNameExpr

from ._abstract import (
AtomicExpr,
Expand Down Expand Up @@ -239,7 +239,7 @@ def __init__(
/,
*,
op: _edgeql.Token | str,
type_: TypeName,
type_: TypeNameExpr,
) -> None:
super().__init__(type_=type_)
if not isinstance(op, _edgeql.Token):
Expand All @@ -252,15 +252,15 @@ def precedence(self) -> _edgeql.Precedence:


@dataclass(kw_only=True, frozen=True)
class PrefixOp(Op):
class UnaryOp(Op):
expr: Expr

def __init__(
self,
*,
expr: ExprCompatible,
op: _edgeql.Token | str,
type_: TypeName,
type_: TypeNameExpr,
) -> None:
object.__setattr__(self, "expr", edgeql_qb_expr(expr))
super().__init__(op=op, type_=type_)
Expand All @@ -276,7 +276,7 @@ def __edgeql_expr__(self, *, ctx: ScopeContext) -> str:


@dataclass(kw_only=True, frozen=True)
class CastOp(PrefixOp):
class CastOp(UnaryOp):
def __init__(
self,
*,
Expand All @@ -300,6 +300,35 @@ def __edgeql_expr__(self, *, ctx: ScopeContext) -> str:
return f"<{self.type.as_quoted_schema_name()}>{expr}"


@dataclass(kw_only=True, frozen=True)
class ObjectWhenType(UnaryOp):
type_filter: TypeNameExpr

def __init__(
self,
*,
expr: ExprCompatible,
type_filter: TypeNameExpr,
type_: TypeNameExpr,
) -> None:
op = _edgeql.Token.RANGBRACKET
object.__setattr__(self, "type_filter", type_filter)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is an issue for another PR (or not bothering), but all these __setattr__s and explicit __init__s in a dataclass are dodgy; are they all just so that op doesn't show up in the dataclass string like it would if it was declared with a default

super().__init__(expr=expr, op=op, type_=type_)

@property
def precedence(self) -> _edgeql.Precedence:
return _edgeql.PRECEDENCE[_edgeql.Token.LBRACKET]

def subnodes(self) -> Iterable[Node]:
return (self.expr,)

def __edgeql_expr__(self, *, ctx: ScopeContext) -> str:
expr = edgeql(self.expr, ctx=ctx)
if _need_left_parens(self.precedence, self.expr):
expr = f"({expr})"
return f"{expr} [is {self.type_filter.as_quoted_schema_name()}]"


def empty_set(type_: TypeName) -> CastOp:
return CastOp(expr=SetLiteral(items=(), type_=type_), type_=type_)

Expand All @@ -319,7 +348,7 @@ def __init__(
lexpr: ExprCompatible,
rexpr: ExprCompatible,
op: _edgeql.Token | str,
type_: TypeName,
type_: TypeNameExpr,
) -> None:
object.__setattr__(self, "lexpr", edgeql_qb_expr(lexpr))
object.__setattr__(self, "rexpr", edgeql_qb_expr(rexpr))
Expand All @@ -337,7 +366,7 @@ def __init__(
lexpr: ExprCompatible,
rexpr: ExprCompatible,
op: _edgeql.Token | str,
type_: TypeName,
type_: TypeNameExpr,
) -> None:
super().__init__(lexpr=lexpr, rexpr=rexpr, op=op, type_=type_)

Expand Down Expand Up @@ -406,7 +435,7 @@ def __init__(
fname: str,
args: list[ExprCompatible] | None = None,
kwargs: dict[str, ExprCompatible] | None = None,
type_: TypeName,
type_: TypeNameExpr,
) -> None:
object.__setattr__(self, "fname", fname)
if args is not None:
Expand Down Expand Up @@ -508,7 +537,7 @@ def subnodes(self) -> Iterable[Node]:
return (self.expr,)

@property
def type(self) -> TypeName:
def type(self) -> TypeNameExpr:
return self.expr.type

@property
Expand Down Expand Up @@ -588,9 +617,15 @@ def __edgeql_expr__(self, *, ctx: ScopeContext) -> str:

@dataclass(kw_only=True, frozen=True)
class InsertStmt(Stmt, TypedExpr):
type_: TypeName # insert can only deal with simple type names

stmt: _edgeql.Token = _edgeql.Token.INSERT
shape: Shape | None = None

@property
def type(self) -> TypeName:
return self.type_

def subnodes(self) -> Iterable[Node | None]:
return (self.shape,)

Expand Down Expand Up @@ -661,7 +696,7 @@ def wrap(
kwargs = {}
if isinstance(expr, ShapeOp):
kwargs["body_scope"] = expr.scope
elif isinstance(expr, (SchemaSet, Path)):
elif expr_uses_auto_splat(expr):
if splat_cb is not None:
shape = splat_cb()
else:
Expand Down Expand Up @@ -757,7 +792,7 @@ def __post_init__(self) -> None:
object.__setattr__(self, "var", var)

@property
def type(self) -> TypeName:
def type(self) -> TypeNameExpr:
return self.body.type

def subnodes(self) -> Iterable[Node]:
Expand Down Expand Up @@ -786,7 +821,7 @@ class Splat(_strenum.StrEnum):
@dataclass(kw_only=True, frozen=True)
class ShapeElement(Node):
name: str | Splat
origin: TypeName
origin: TypeNameExpr
expr: Expr | None = None

def subnodes(self) -> Iterable[Node]:
Expand All @@ -798,7 +833,7 @@ def subnodes(self) -> Iterable[Node]:
@classmethod
def splat(
cls,
source: TypeName,
source: TypeNameExpr,
*,
kind: Splat = Splat.STAR,
) -> Self:
Expand All @@ -815,14 +850,23 @@ def subnodes(self) -> Iterable[Node]:
@classmethod
def splat(
cls,
source: TypeName,
source: TypeNameExpr,
*,
kind: Splat = Splat.STAR,
) -> Self:
elements = [ShapeElement.splat(source=source, kind=kind)]
return cls(elements=elements)


def expr_uses_auto_splat(expr: Expr) -> bool:
if isinstance(expr, (SchemaSet, Path)):
return True
elif isinstance(expr, ObjectWhenType):
return expr_uses_auto_splat(expr.expr)
else:
return False


def _render_shape(
shape: Shape,
source: Expr,
Expand Down Expand Up @@ -873,7 +917,7 @@ def subnodes(self) -> Iterable[Node]:
return (self.iter_expr, self.shape)

@property
def type(self) -> TypeName:
def type(self) -> TypeNameExpr:
return self.iter_expr.type

@property
Expand Down
Loading