Repository: edgedb/edgedb-python (language: python)

Issue:
Generated Enum types are not valid as parameters
**Describe the bug**

When my schema contains a reference to [an Enum type](https://www.edgedb.com/docs/stdlib/enum#ref-std-enum), then the code generator creates an `enum.Enum` subclass. However, I can pass an instance of this Enum type into a parameter function, and mypy is happy. However, when executing the code, I see the error:

```python
>   ???
E   TypeError: a str or edgedb.EnumValue(__tid__=59ec73b8-4261-11ed-a3eb-b1b44f04f696) is expected as a valid enum argument, got Color

edgedb/protocol/./codecs/./enum.pyx:27: TypeError
```

Changing the passed argument to the Enum value (e.g `Color.RED.value`) results in mypy complaining that `str` is not a valid type.

<details>
  <summary>Full trace back - click to expand!</summary>

```python
test.py:62: in insert_shirt
    return await executor.query_single(
../../Library/Caches/pypoetry/virtualenvs/test-edgedb-bug/lib/python3.10/site-packages/edgedb/abstract.py:238: in query_single
    return await self._query(QueryContext(
../../Library/Caches/pypoetry/virtualenvs/test-edgedb-bug/lib/python3.10/site-packages/edgedb/base_client.py:739: in _query
    return await con.raw_query(query_context)
../../Library/Caches/pypoetry/virtualenvs/test-edgedb-bug/lib/python3.10/site-packages/edgedb/base_client.py:239: in raw_query
    raise e
../../Library/Caches/pypoetry/virtualenvs/test-edgedb-bug/lib/python3.10/site-packages/edgedb/base_client.py:234: in raw_query
    return await self._protocol.query(**args)
edgedb/protocol/protocol.pyx:594: in query
    ???
edgedb/protocol/protocol.pyx:559: in execute
    ???
edgedb/protocol/protocol.pyx:379: in _execute
    ???
edgedb/protocol/protocol.pyx:1125: in edgedb.protocol.protocol.SansIOProtocol.encode_args
    ???
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

>   ???
E   edgedb.errors.InvalidArgumentError: invalid input for query argument $color: <Color.RED: 'Red'> (a str or edgedb.EnumValue(__tid__=59ec73b8-4261-11ed-a3eb-b1b44f04f696) is expected as a valid enum argument, got Color)

edgedb/protocol/./codecs/./object.pyx:114: InvalidArgumentError
```

</details> 

**Reproduction**

Given the schema:
```esdl
scalar type Color extending enum<Red, Green, Blue>;

type Shirt {
  property color -> Color;
}
```

When we have a query:
```edgeql
select (
    insert Shirt {
        color := <Color>$color
    }
) { id, color };
```

Run `edgedb-py`, and execute:

```python
insert_shirt(color=Color.RED)
```

**Expected behavior**

I expect to be able to pass the Enum value directly as a parameter, as it's the type declared by the code generator.

**Versions (please complete the following information):**

<!--
For EdgeDB Version: run `edgedb query 'select sys::get_version_as_str()'` from your project directory (or run `select sys::get_version_as_str();` in the EdgeDB interactive shell).
For EdgeDB CLI version: Run `edgedb --version` from anywhere
For `edgedb-python` version: Run `python -m pip show edgedb`.
For Python version: Run `python --version`
-->

- OS: macOS Ventura 13.2.1 (22D68)
- EdgeDB version: 2.12+a5dfd9f
- EdgeDB CLI version:  2.3.1+5d93f42
- `edgedb-python` version: 1.3.0
- Python version: 3.10.6

**Additional context**

Note: I raised another issue/question in discord regarding optional parameters: https://discord.com/channels/841451783728529451/849374935209607199/1081881229676265592



Discussion / hints:

Fix patch:
diff --git a/edgedb/datatypes/enum.pyx b/edgedb/datatypes/enum.pyx
index a3fa6705..9ca4b155 100644
--- a/edgedb/datatypes/enum.pyx
+++ b/edgedb/datatypes/enum.pyx
@@ -27,33 +27,42 @@ class EnumValue(enum.Enum):
     def __repr__(self):
         return f'<edgedb.EnumValue {self._value_!r}>'
 
+    @classmethod
+    def _try_from(cls, value):
+        if isinstance(value, EnumValue):
+            return value
+        elif isinstance(value, enum.Enum):
+            return cls(value.value)
+        else:
+            raise TypeError
+
     def __lt__(self, other):
-        if not isinstance(other, EnumValue):
-            return NotImplemented
+        other = self._try_from(other)
         if self.__tid__ != other.__tid__:
             return NotImplemented
         return self._index_ < other._index_
 
     def __gt__(self, other):
-        if not isinstance(other, EnumValue):
-            return NotImplemented
+        other = self._try_from(other)
         if self.__tid__ != other.__tid__:
             return NotImplemented
         return self._index_ > other._index_
 
     def __le__(self, other):
-        if not isinstance(other, EnumValue):
-            return NotImplemented
+        other = self._try_from(other)
         if self.__tid__ != other.__tid__:
             return NotImplemented
         return self._index_ <= other._index_
 
     def __ge__(self, other):
-        if not isinstance(other, EnumValue):
-            return NotImplemented
+        other = self._try_from(other)
         if self.__tid__ != other.__tid__:
             return NotImplemented
         return self._index_ >= other._index_
 
+    def __eq__(self, other):
+        other = self._try_from(other)
+        return self is other
+
     def __hash__(self):
-        return hash((self.__tid__, self._value_))
+        return hash(self._value_)
diff --git a/edgedb/protocol/codecs/enum.pyx b/edgedb/protocol/codecs/enum.pyx
index 78a30f6b..d7fe5620 100644
--- a/edgedb/protocol/codecs/enum.pyx
+++ b/edgedb/protocol/codecs/enum.pyx
@@ -24,9 +24,13 @@ cdef class EnumCodec(BaseCodec):
 
     cdef encode(self, WriteBuffer buf, object obj):
         if not isinstance(obj, (self.cls, str)):
-            raise TypeError(
-                f'a str or edgedb.EnumValue(__tid__={self.cls.__tid__}) is '
-                f'expected as a valid enum argument, got {type(obj).__name__}')
+            try:
+                obj = self.cls._try_from(obj)
+            except (TypeError, ValueError):
+                raise TypeError(
+                    f'a str or edgedb.EnumValue(__tid__={self.cls.__tid__}) '
+                    f'is expected as a valid enum argument, '
+                    f'got {type(obj).__name__}') from None
         pgproto.text_encode(DEFAULT_CODEC_CONTEXT, buf, str(obj))
 
     cdef decode(self, FRBuffer *buf):


Test patch:
diff --git a/tests/test_async_query.py b/tests/test_async_query.py
index e51ac197..103e6393 100644
--- a/tests/test_async_query.py
+++ b/tests/test_async_query.py
@@ -19,6 +19,7 @@
 
 import datetime
 import decimal
+import enum
 import json
 import random
 import unittest
@@ -912,6 +913,31 @@ async def test_enum_argument_01(self):
                 edgedb.InvalidArgumentError, 'a str or edgedb.EnumValue'):
             await self.client.query_single('SELECT <MyEnum>$0', 123)
 
+    async def test_enum_argument_02(self):
+        class MyEnum(enum.Enum):
+            A = "A"
+            B = "B"
+            C = "C"
+
+        A = await self.client.query_single('SELECT <MyEnum>$0', MyEnum.A)
+        self.assertEqual(str(A), 'A')
+        self.assertEqual(A, MyEnum.A)
+        self.assertEqual(MyEnum.A, A)
+        self.assertLess(A, MyEnum.B)
+        self.assertGreater(MyEnum.B, A)
+
+        mapping = {MyEnum.A: 1, MyEnum.B: 2}
+        self.assertEqual(mapping[A], 1)
+
+        with self.assertRaises(ValueError):
+            _ = A > MyEnum.C
+        with self.assertRaises(ValueError):
+            _ = A < MyEnum.C
+        with self.assertRaises(ValueError):
+            _ = A == MyEnum.C
+        with self.assertRaises(edgedb.InvalidArgumentError):
+            await self.client.query_single('SELECT <MyEnum>$0', MyEnum.C)
+
     async def test_json(self):
         self.assertEqual(
             await self.client.query_json('SELECT {"aaa", "bbb"}'),
diff --git a/tests/test_enum.py b/tests/test_enum.py
index a5e99b15..af005f84 100644
--- a/tests/test_enum.py
+++ b/tests/test_enum.py
@@ -43,8 +43,10 @@ async def test_enum_01(self):
         self.assertEqual(repr(ct_red), "<edgedb.EnumValue 'red'>")
 
         self.assertEqual(str(ct_red), 'red')
-        self.assertNotEqual(ct_red, 'red')
-        self.assertFalse(ct_red == 'red')
+        with self.assertRaises(TypeError):
+            _ = ct_red != 'red'
+        with self.assertRaises(TypeError):
+            _ = ct_red == 'red'
         self.assertFalse(ct_red == c_red)
 
         self.assertEqual(ct_red, ct_red)
@@ -74,8 +76,8 @@ async def test_enum_01(self):
         with self.assertRaises(TypeError):
             _ = ct_red >= c_red
 
-        self.assertNotEqual(hash(ct_red), hash(c_red))
-        self.assertNotEqual(hash(ct_red), hash('red'))
+        self.assertEqual(hash(ct_red), hash(c_red))
+        self.assertEqual(hash(ct_red), hash('red'))
 
     async def test_enum_02(self):
         c_red = await self.client.query_single('SELECT <Color>"red"')