From da6e6c1cc94d596c511ebb31d5409b2c54610c28 Mon Sep 17 00:00:00 2001
From: Saul Shanabrook <s.shanabrook@gmail.com>
Date: Tue, 29 Oct 2024 14:36:33 -0400
Subject: [PATCH 1/3] Add better error message when using @function in class

---
 docs/changelog.md               |  1 +
 python/egglog/egraph.py         | 39 +++++++++++++++++++--------------
 python/tests/test_high_level.py | 12 ++++++++++
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/docs/changelog.md b/docs/changelog.md
index 2d1ccf1b..8b938fdb 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -6,6 +6,7 @@ _This project uses semantic versioning_
 
 - Fix pretty printing of lambda functions
 - Add support for subsuming rewrite generated by default function and method definitions
+- Add better error message when using @function in class (thanks @shinawy)
 
 ## 8.0.1 (2024-10-24)
 
diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py
index 1cc41350..c237d36e 100644
--- a/python/egglog/egraph.py
+++ b/python/egglog/egraph.py
@@ -577,23 +577,25 @@ def _generate_class_decls(  # noqa: C901,PLR0912
             decl = FunctionDecl(special_function_name, builtin=True, egg_name=egg_fn)
             decls.set_function_decl(ref, decl)
             continue
-
-        _, add_rewrite = _fn_decl(
-            decls,
-            egg_fn,
-            ref,
-            fn,
-            locals,
-            default,
-            cost,
-            merge,
-            on_merge,
-            mutates,
-            builtin,
-            ruleset=ruleset,
-            unextractable=unextractable,
-            subsume=subsume,
-        )
+        try:
+            _, add_rewrite = _fn_decl(
+                decls,
+                egg_fn,
+                ref,
+                fn,
+                locals,
+                default,
+                cost,
+                merge,
+                on_merge,
+                mutates,
+                builtin,
+                ruleset=ruleset,
+                unextractable=unextractable,
+                subsume=subsume,
+            )
+        except ValueError as e:
+            raise ValueError(f"Error processing {cls_name}.{method_name}: {e}") from e
 
         if not builtin and not isinstance(ref, InitRef) and not mutates:
             add_default_funcs.append(add_rewrite)
@@ -721,6 +723,9 @@ def _fn_decl(
     """
     Sets the function decl for the function object and returns the ref as well as a thunk that sets the default callable.
     """
+    if isinstance(fn, RuntimeFunction):
+        msg = "Inside of classes, wrap methods with the `method` decorator, not `function`"
+        raise ValueError(msg)  # noqa: TRY004
     if not isinstance(fn, FunctionType):
         raise NotImplementedError(f"Can only generate function decls for functions not {fn}  {type(fn)}")
 
diff --git a/python/tests/test_high_level.py b/python/tests/test_high_level.py
index 5edda452..c22fb997 100644
--- a/python/tests/test_high_level.py
+++ b/python/tests/test_high_level.py
@@ -762,3 +762,15 @@ def test_inserting_map(self):
 
     def test_creating_map(self):
         EGraph().simplify(Map[String, i64].empty(), 1)
+
+
+def test_helpful_error_function_class():
+    class E(Expr):
+        @function(cost=10)
+        def __init__(self) -> None: ...
+
+    with pytest.raises(
+        ValueError,
+        match="Error processing E.__init__: Inside of classes, wrap methods with the `method` decorator, not `function`",
+    ):
+        E()

From cc4349231fce4a499f7fa033eba5a9d6f85f236a Mon Sep 17 00:00:00 2001
From: Saul Shanabrook <s.shanabrook@gmail.com>
Date: Thu, 5 Dec 2024 10:08:02 -0500
Subject: [PATCH 2/3] Update python/egglog/egraph.py

Co-authored-by: Hameer Abbasi <2190658+hameerabbasi@users.noreply.github.com>
---
 python/egglog/egraph.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py
index c237d36e..7672589a 100644
--- a/python/egglog/egraph.py
+++ b/python/egglog/egraph.py
@@ -595,7 +595,7 @@ def _generate_class_decls(  # noqa: C901,PLR0912
                 subsume=subsume,
             )
         except ValueError as e:
-            raise ValueError(f"Error processing {cls_name}.{method_name}: {e}") from e
+            raise ValueError(f"Error processing {cls_name}.{method_name}") from e
 
         if not builtin and not isinstance(ref, InitRef) and not mutates:
             add_default_funcs.append(add_rewrite)

From 14a7b80bb89dc2e50ce087e8f8556f64cc11e9b6 Mon Sep 17 00:00:00 2001
From: Saul Shanabrook <s.shanabrook@gmail.com>
Date: Thu, 5 Dec 2024 10:13:55 -0500
Subject: [PATCH 3/3] Revert "Update python/egglog/egraph.py"

cc4349231fce4a499f7fa033eba5a9d6f85f236a
---
 python/egglog/egraph.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py
index 7672589a..c237d36e 100644
--- a/python/egglog/egraph.py
+++ b/python/egglog/egraph.py
@@ -595,7 +595,7 @@ def _generate_class_decls(  # noqa: C901,PLR0912
                 subsume=subsume,
             )
         except ValueError as e:
-            raise ValueError(f"Error processing {cls_name}.{method_name}") from e
+            raise ValueError(f"Error processing {cls_name}.{method_name}: {e}") from e
 
         if not builtin and not isinstance(ref, InitRef) and not mutates:
             add_default_funcs.append(add_rewrite)