diff --git a/gel/_internal/_codegen/_models/_pydantic.py b/gel/_internal/_codegen/_models/_pydantic.py index 2e299a2e..de417203 100644 --- a/gel/_internal/_codegen/_models/_pydantic.py +++ b/gel/_internal/_codegen/_models/_pydantic.py @@ -2575,17 +2575,38 @@ def _write_enum_scalar_cast( aexpr = self.import_name(BASE_IMPL, "AnnotatedExpr") cast_op = self.import_name(BASE_IMPL, "CastOp") + py_to_gel_casts = self._get_scalar_py_to_gel_casts(stype) + if not py_to_gel_casts: + return + + arg_name: str = "expr" + arg_types: list[str] = [expr_compat, *py_to_gel_casts.keys()] + with self._classmethod_def( "cast", - [f"expr: {expr_compat}"], + [f"{arg_name}: {' | '.join(arg_types)}"], type_self, ): + self.write() + self.write(f"match {arg_name}:") + with self.indented(): + for py_type, gel_cast in py_to_gel_casts.items(): + self.write(f"case {py_type}():") + with self.indented(): + cast_text: str + if isinstance(gel_cast, str): + cast_text = f"{gel_cast}({arg_name})" + else: + cast_text = gel_cast(arg_name) + self.write(f"{arg_name} = {cast_text}") + + self.write() self.write(f"return {aexpr}( # type: ignore [return-value]") with self.indented(): self.write("cls,") self.write(f"{cast_op}(") with self.indented(): - self.write("expr=expr,") + self.write(f"expr={arg_name},") self.write("type_=cls.__gel_reflection__.type_name,") self.write(")") self.write(")") @@ -2818,25 +2839,38 @@ def _write_regular_scalar_cast( self_ = self.import_name("typing_extensions", "Self") type_self = f"{type_}[{self_}]" - if signature_only: - self.write() - with self._classmethod_def( - "cast", - [f"expr: {expr_compat}"], - type_self, - ): + py_to_gel_casts = self._get_scalar_py_to_gel_casts(stype) + if not py_to_gel_casts: + return + + arg_name: str = "expr" + arg_types: list[str] = [expr_compat, *py_to_gel_casts.keys()] + + with self._classmethod_def( + "cast", + [f"{arg_name}: {' | '.join(arg_types)}"], + type_self, + ): + if signature_only: self.write("...") - else: - aexpr = self.import_name(BASE_IMPL, "AnnotatedExpr") - cast_op = self.import_name(BASE_IMPL, "CastOp") + else: + aexpr = self.import_name(BASE_IMPL, "AnnotatedExpr") + cast_op = self.import_name(BASE_IMPL, "CastOp") - self.write() - with self._classmethod_def( - "cast", - [f"expr: {expr_compat}"], - type_self, - ): + self.write(f"match {arg_name}:") + with self.indented(): + for py_type, gel_cast in py_to_gel_casts.items(): + self.write(f"case {py_type}():") + with self.indented(): + cast_text: str + if isinstance(gel_cast, str): + cast_text = f"{gel_cast}({arg_name})" + else: + cast_text = gel_cast(arg_name) + self.write(f"{arg_name} = {cast_text}") + + self.write() self.write(f"return {aexpr}( # type: ignore [return-value]") with self.indented(): self.write("cls,") @@ -2847,6 +2881,98 @@ def _write_regular_scalar_cast( self.write(")") self.write(")") + def _get_scalar_py_to_gel_casts( + self, + stype: reflection.ScalarType, + ) -> dict[str, str | Callable[[str], str]] | None: + if not (explicit_casts := self._casts.explicit_casts_to.get(stype.id)): + return None + + py_to_gel_casts: dict[str, str | Callable[[str], str]] = {} + + # Determine if the result type can be directly cast from a literal + direct_py_type_name: tuple[str, str] | None = None + + if py_type_names := _qbmodel.get_py_type_for_scalar( + stype.name, + consider_generic=False, + ): + # with consider_generic=False, there should be 1 value + direct_py_type_name = py_type_names[0] + if literal_name := _qbmodel.get_literal_name_for_py_type( + direct_py_type_name + ): + py_type = self.import_name(*direct_py_type_name) + literal = self.import_name(BASE_IMPL, literal_name) + + py_to_gel_casts[py_type] = lambda x: ( + f"{literal}(" + f"val={x}," + f"type_=cls.__gel_reflection__.type_name," + f")" + ) + + # Determine what python types can converted to a gel type before cast + + # Get the gel types that can be cast to result type + scalar_arg_types = [ + arg_type + for arg_type_id in explicit_casts + if (arg_type := self._types.get(arg_type_id)) + if reflection.is_scalar_type(arg_type) + if arg_type.schemapath not in GENERIC_TYPES + ] + + # Find the python types associated with the gel types + py_to_scalar_types: dict[ + tuple[str, str], list[reflection.ScalarType] + ] = {} + for scalar_arg_type in scalar_arg_types: + if py_type_names := _qbmodel.get_py_type_for_scalar( + scalar_arg_type.name, + consider_generic=False, + ): + # with consider_generic=False, there should be 1 value + py_type_name = py_type_names[0] + + if py_type_name == direct_py_type_name: + # Skip the directly converted type + continue + + if py_type_name not in py_to_scalar_types: + py_to_scalar_types[py_type_name] = [] + + py_to_scalar_types[py_type_name].append(scalar_arg_type) + + # Pick the best gel type to convert to + for py_type_name, scalar_types in py_to_scalar_types.items(): + py_type = self.import_name(*py_type_name) + + scalars_with_rank: list[tuple[reflection.ScalarType, int]] = [] + for scalar_type in scalar_types: + rank = _qbmodel.get_py_type_scalar_match_rank( + py_type_name, scalar_type.name + ) + if rank is None: + continue + scalars_with_rank.append((scalar_type, rank)) + + if not scalars_with_rank: + # This can happen for scalars which don't convert to simple + # python primitives. eg. ext::pgvector::halfvec + continue + + best_scalar_type = min( + scalars_with_rank, key=operator.itemgetter(1) + )[0] + gel_type = self.get_type( + best_scalar_type, import_time=ImportTime.typecheck_runtime + ) + + py_to_gel_casts[py_type] = gel_type + + return py_to_gel_casts + def render_callable_return_type( self, tp: reflection.Type, diff --git a/gel/_internal/_qb/__init__.py b/gel/_internal/_qb/__init__.py index e4881317..cc5074dd 100644 --- a/gel/_internal/_qb/__init__.py +++ b/gel/_internal/_qb/__init__.py @@ -18,6 +18,7 @@ ) from ._expressions import ( + BigIntLiteral, BinaryOp, BoolLiteral, BytesLiteral, @@ -112,6 +113,7 @@ "AnnotatedPath", "AnnotatedVar", "BaseAlias", + "BigIntLiteral", "BinaryOp", "BoolLiteral", "BytesLiteral", diff --git a/gel/_internal/_qb/_protocols.py b/gel/_internal/_qb/_protocols.py index 00d2f944..0183af96 100644 --- a/gel/_internal/_qb/_protocols.py +++ b/gel/_internal/_qb/_protocols.py @@ -121,6 +121,12 @@ def edgeql( return value +ExprPrimitive = TypeAliasType( + "ExprPrimitive", + int, +) + + def edgeql_qb_expr( x: ExprCompatible | ExprClosure, *, diff --git a/gel/_internal/_qbmodel/_abstract/__init__.py b/gel/_internal/_qbmodel/_abstract/__init__.py index 30079e6e..6d54bc7d 100644 --- a/gel/_internal/_qbmodel/_abstract/__init__.py +++ b/gel/_internal/_qbmodel/_abstract/__init__.py @@ -95,6 +95,7 @@ Tuple, TupleMeta, UUIDImpl, + get_literal_name_for_py_type, get_py_type_from_gel_type, get_base_scalars_backed_by_py_type, get_overlapping_py_types, @@ -177,6 +178,7 @@ "empty_set_if_none", "field_descriptor", "get_base_scalars_backed_by_py_type", + "get_literal_name_for_py_type", "get_overlapping_py_types", "get_proxy_linkprops", "get_py_base_for_scalar", diff --git a/gel/_internal/_qbmodel/_abstract/_primitive.py b/gel/_internal/_qbmodel/_abstract/_primitive.py index 5b436efe..c363191d 100644 --- a/gel/_internal/_qbmodel/_abstract/_primitive.py +++ b/gel/_internal/_qbmodel/_abstract/_primitive.py @@ -817,6 +817,15 @@ def get_py_type_scalar_match_rank( decimal.Decimal: _qb.DecimalLiteral, } +_py_type_name_to_literal_name: dict[tuple[str, str], str] = { + ("builtins", "bool"): "BoolLiteral", + ("builtins", "int"): "IntLiteral", + ("builtins", "float"): "FloatLiteral", + ("builtins", "str"): "StringLiteral", + ("builtins", "bytes"): "BytesLiteral", + ("decimal", "Decimal"): "DecimalLiteral", +} + _PT_co = TypeVar("_PT_co", bound=PyConstType, covariant=True) _ST = TypeVar("_ST", bound=GelScalarType, default=GelScalarType) @@ -850,6 +859,12 @@ def get_literal_for_scalar( ) +def get_literal_name_for_py_type( + py_type_name: tuple[str, str], +) -> str | None: + return _py_type_name_to_literal_name.get(py_type_name) + + class PyTypeScalar( _typing_parametric.ParametricType, GelScalarType, diff --git a/gel/models/pydantic.py b/gel/models/pydantic.py index b46701d3..0b36e424 100644 --- a/gel/models/pydantic.py +++ b/gel/models/pydantic.py @@ -22,9 +22,14 @@ OPERAND_IS_ALIAS, AnnotatedExpr, BaseAlias, + BigIntLiteral, + BoolLiteral, + BytesLiteral, CastOp, + DecimalLiteral, EmptyDirection, Direction, + FloatLiteral, GelLinkMetadata, GelObjectTypeMetadata, GelPointerReflection, @@ -34,10 +39,12 @@ ExprCompatible, IndexOp, InfixOp, + IntLiteral, FuncCall, ObjectWhenType, PathAlias, SetLiteral, + StringLiteral, UnaryOp, construct_infix_op_chain, ) @@ -117,6 +124,9 @@ "Array", "ArrayMeta", "BaseAlias", + "BigIntLiteral", + "BoolLiteral", + "BytesLiteral", "Cardinality", "CastOp", "ComputedLink", @@ -128,12 +138,14 @@ "DateImpl", "DateTimeImpl", "DateTimeLike", + "DecimalLiteral", "DefaultValue", "DeferredImport", "Direction", "EmptyDirection", "ExprClosure", "ExprCompatible", + "FloatLiteral", "FuncCall", "GelLinkMetadata", "GelLinkModel", @@ -151,6 +163,7 @@ "IdProperty", "IndexOp", "InfixOp", + "IntLiteral", "JSONImpl", "LazyClassProperty", "LinkClassNamespace", @@ -183,6 +196,7 @@ "RequiredMultiLinkWithProps", "SchemaPath", "SetLiteral", + "StringLiteral", "TimeDeltaImpl", "TimeImpl", "Tuple", diff --git a/tests/test_qb.py b/tests/test_qb.py index 0fd4fd2a..1d61b3cd 100644 --- a/tests/test_qb.py +++ b/tests/test_qb.py @@ -1351,25 +1351,37 @@ def test_qb_array_agg_01(self): def test_qb_cast_scalar_01(self): # scalar to scalar - from models.orm import std + from models.orm_qb import std result = self.client.get(std.str.cast(std.int64(1))) self.assertEqual(result, "1") + # python scalar to scalar + from models.orm_qb import std + + result = self.client.get(std.str.cast(1)) + self.assertEqual(result, "1") + result = self.client.get(std.str.cast("1")) + self.assertEqual(result, "1") + def test_qb_cast_scalar_02(self): # enum to scalar - from models.orm import default, std + from models.orm_qb import default, std result = self.client.get(std.str.cast(default.Color.Red)) self.assertEqual(result, "Red") def test_qb_cast_scalar_03(self): # scalar to enum - from models.orm import default, std + from models.orm_qb import default, std result = self.client.get(default.Color.cast(std.str("Red"))) self.assertEqual(result, default.Color.Red) + # python scalar to enum + result = self.client.get(default.Color.cast("Red")) + self.assertEqual(result, default.Color.Red) + def test_qb_cast_array_01(self): # array[scalar] to array[scalar] from models.orm_qb import std @@ -1696,9 +1708,7 @@ def test_qb_is_type_basic_07(self): # Link TypeIntersection from models.orm_qb import default - result = self.client.query( - default.Link_Inh_A.l.is_(default.Inh_B) - ) + result = self.client.query(default.Link_Inh_A.l.is_(default.Inh_B)) self._assertObjectsWithFields( result, @@ -1901,9 +1911,9 @@ def test_qb_is_type_for_01(self): from models.orm_qb import default, std result = self.client.query( - std.for_( - default.Inh_A.is_(default.Inh_B), lambda x: x - ).select(a=True) + std.for_(default.Inh_A.is_(default.Inh_B), lambda x: x).select( + a=True + ) ) self._assertObjectsWithFields(