Support passing python types to scalars in casts. (#953)

Instead of generating cast functions for everything and letting the server figure it out, generate cast functions for scalars based on the reflected explicit casts.

Add support to convert python types into gel scalars if possible. For example:
```py
std.str.cast(1)
std.str.cast("abc")
```

This works for basic scalars (like `int64` and `str`) as well as enums.

Author: dnwpark

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,

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",

gel/_internal/_qb/_protocols.py

@@ -121,6 +121,12 @@ def edgeql(
     return value
 
 
+ExprPrimitive = TypeAliasType(
+    "ExprPrimitive",
+    int,
+)
+
+
 def edgeql_qb_expr(
     x: ExprCompatible | ExprClosure,
     *,

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",

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,

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",

tests/test_qb.py

@@ -1356,6 +1356,14 @@ def test_qb_cast_scalar_01(self):
         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_qb import default, std
@@ -1370,6 +1378,10 @@ def test_qb_cast_scalar_03(self):
         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