Skip to content

Commit 06d8a21

Browse files
authored
Support using enums in the query builder (#932)
Work on #635
1 parent cf37cee commit 06d8a21

File tree

4 files changed

+20
-21
lines changed

4 files changed

+20
-21
lines changed

gel/_internal/_codegen/_models/_pydantic.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2508,11 +2508,15 @@ def _write_enum_scalar_type(
25082508
type_name = stype.schemapath
25092509
tname = type_name.name
25102510
assert stype.enum_values
2511-
anyenum = self.import_name(BASE_IMPL, "AnyEnum")
2511+
anyenum = self.get_object(
2512+
SchemaPath('std', 'anyenum'),
2513+
aspect=ModuleAspect.SHAPES,
2514+
)
25122515
with self._class_def(tname, [anyenum]):
25132516
self.write_description(stype)
25142517
for value in stype.enum_values:
25152518
self.write(f"{ident(value)} = {value!r}")
2519+
self.write_type_reflection(stype)
25162520
self.write_section_break()
25172521

25182522
def _write_scalar_type(
@@ -2601,6 +2605,8 @@ def _write_regular_scalar_type(
26012605
typecheck_meta_parents.append(anyenummeta_)
26022606
typecheck_meta_ignores = ["misc", "unused-ignore"]
26032607

2608+
self.export("anyenum")
2609+
26042610
if not runtime_parents:
26052611
typecheck_parents = [self.get_type(self._types_by_name["anytype"])]
26062612
runtime_parents = typecheck_parents
@@ -4176,7 +4182,7 @@ def _write_object_type_qb_methods(
41764182
)
41774183
union = []
41784184
select_union = [builtin_bool, expr_closure, expr_proto]
4179-
if reflection.is_non_enum_scalar_type(target_t):
4185+
if reflection.is_scalar_type(target_t):
41804186
broad_ptr_t = self._get_pytype_for_scalar(target_t)
41814187
union.extend(broad_ptr_t)
41824188
select_union.extend(broad_ptr_t)

gel/_internal/_qbmodel/_abstract/_expressions.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
from ._primitive import (
2525
GelPrimitiveType,
2626
PyConstType,
27-
PyTypeScalar,
27+
GelScalarType,
2828
get_literal_for_scalar,
2929
get_literal_for_value,
3030
)
@@ -106,7 +106,7 @@ def _const_to_expr(
106106
if isinstance(attr, ModelFieldDescriptor) and (
107107
tp := attr.get_resolved_type()
108108
):
109-
if issubclass(tp, PyTypeScalar):
109+
if issubclass(tp, GelScalarType):
110110
return get_literal_for_scalar(tp, val)
111111

112112
raise TypeError(f"unsupported const variant for {cls!r}.{pname!r}")

gel/_internal/_qbmodel/_abstract/_primitive.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,11 @@ def __setitem__(self, key: str, value: Any) -> None:
153153

154154

155155
class AnyEnum(GelScalarType, StrEnum, metaclass=AnyEnumMeta):
156-
pass
156+
def __edgeql_literal__(self) -> _qb.Literal | _qb.CastOp:
157+
return _qb.Literal(
158+
val=str(self),
159+
type_=type(self).__gel_reflection__.name,
160+
)
157161

158162

159163
if TYPE_CHECKING:
@@ -776,12 +780,14 @@ def get_literal_for_value(
776780

777781

778782
def get_literal_for_scalar(
779-
t: type[PyTypeScalar[_PT_co]],
783+
t: type[GelScalarType],
780784
v: Any,
781785
) -> _qb.Literal | _qb.CastOp:
782786
if not isinstance(v, t):
783-
v = t(v)
784-
ltype = _py_type_to_literal.get(t.__gel_py_type__) # type: ignore [arg-type]
787+
v = t(v) # type: ignore [call-arg]
788+
ltype = None
789+
if issubclass(t, PyTypeScalar):
790+
ltype = _py_type_to_literal.get(t.__gel_py_type__)
785791
if ltype is not None:
786792
return ltype(val=v) # type: ignore [call-arg]
787793
else:

tests/test_qb.py

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1091,11 +1091,6 @@ def test_qb_boolean_operator_error_01(self):
10911091
with self.assertRaisesRegex(TypeError, "use std.exists"):
10921092
default.User.filter(lambda u: u.name is not None) # type: ignore
10931093

1094-
@tb.xfail('''
1095-
Filtering on enums broken. And Enums need __gel_type_class__?
1096-
See #635
1097-
''')
1098-
@tb.skip_typecheck
10991094
def test_qb_enum_01(self):
11001095
from models.orm import default
11011096

@@ -1338,10 +1333,6 @@ def test_qb_delete_03(self):
13381333
self.assertEqual(res[1].body, "I'm Alice")
13391334
self.assertEqual(res[1].author.name, "Alice")
13401335

1341-
@tb.xfail('''
1342-
Filtering on enums broken. And Enums need __gel_type_class__?
1343-
See #635
1344-
''')
13451336
def test_qb_enum_edit_01(self):
13461337
from models.orm import default
13471338

@@ -1356,10 +1347,6 @@ def test_qb_enum_edit_01(self):
13561347
self.assertEqual(e.color, default.Color.Orange)
13571348
self.assertEqual(e.name, "red")
13581349

1359-
@tb.xfail('''
1360-
Filtering on enums broken. And Enums need __gel_type_class__?
1361-
See #635
1362-
''')
13631350
def test_qb_enum_edit_02(self):
13641351
from models.orm import default
13651352

0 commit comments

Comments
 (0)