Merge pull request #286 from egraphs-good/fix-import-bug

Fix bug resolving types

Author: saulshanabrook

Cargo.lock

@@ -4,6 +4,10 @@ _This project uses semantic versioning_
 
 ## UNRELEASED
 
+- Fix bug on resolving types if not all imported to your module [#286](https://github.com/egraphs-good/egglog-python/pull/286)
+  - Also stops special casing including `Callable` as a global. So if you previously included this in a `TYPE_CHECKING` block so it wasn
+    available at runtime you will have to move this to a runtime import if used in a type alias.
+
 ## 10.0.0 (2025-03-28)
 
 - Change builtins to not evaluate values in egraph and changes facts to compare structural equality instead of using an egraph when converting to a boolean, removing magic context (`EGraph.current` and `Schedule.current`) that was added in release 9.0.0.

docs/changelog.md

@@ -198,6 +198,10 @@ ignore = [
     "SLF001",
     # allow blind exception to add context
     "BLE001",
+    # Don't move type checking around so that can be accessed at runtime
+    "TCH001",
+    "TCH002",
+    "TCH003",
 ]
 select = ["ALL"]
 
@@ -217,6 +221,15 @@ preview = true
 # Don't require annotations for tests
 "python/tests/**" = ["ANN001", "ANN201", "INP001"]
 
+# Disable these tests instead for now since ruff doesn't support including all method annotations of decorated class
+# [tool.ruff.lint.flake8-type-checking]
+# runtime-evaluated-decorators = [
+#     "egglog.function",
+#     "egglog.method",
+#     "egglog.ruleset",
+# ]
+# runtime-evaluated-base-classes = ["egglog.Expr"]
+
 [tool.mypy]
 ignore_missing_imports = true
 warn_redundant_casts = true

pyproject.toml

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from fractions import Fraction
 from functools import partial, reduce
 from types import FunctionType, MethodType
@@ -20,7 +21,7 @@
 from .thunk import Thunk
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterator
+    from collections.abc import Iterator
 
 
 __all__ = [

python/egglog/builtins.py

@@ -1817,12 +1817,13 @@ def _rewrite_or_rule_generator(gen: RewriteOrRuleGenerator, frame: FrameType) ->
     """
     Returns a thunk which will call the function with variables of the type and name of the arguments.
     """
-    # Get the local scope from where the function is defined, so that we can get any type hints that are in the scope
-    # but not in the globals
-    globals = gen.__globals__.copy()
-    if "Callable" not in globals:
-        globals["Callable"] = Callable
-    hints = get_type_hints(gen, globals, frame.f_locals)
+    # Need to manually pass in the frame locals from the generator, because otherwise classes defined within function
+    # will not be available in the annotations
+    # combine locals and globals so that they are the same dict. Otherwise get_type_hints will go through the wrong
+    # path and give an error for the test
+    # python/tests/test_no_import_star.py::test_no_import_star_rulesset
+    combined = {**gen.__globals__, **frame.f_locals}
+    hints = get_type_hints(gen, combined, combined)
     args = [_var(p.name, hints[p.name]) for p in signature(gen).parameters.values()]
     return list(gen(*args))  # type: ignore[misc]
 

python/egglog/egraph.py

@@ -6,13 +6,10 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from collections.abc import Callable
 
 from egglog import *
 
-if TYPE_CHECKING:
-    from collections.abc import Callable
-
 
 class Math(Expr):
     def __init__(self, i: i64Like) -> None: ...

python/egglog/examples/higher_order_functions.py

@@ -7,12 +7,11 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, ClassVar
+from collections.abc import Callable
+from typing import ClassVar
 
 from egglog import *
-
-if TYPE_CHECKING:
-    from collections.abc import Callable
+from egglog import Expr
 
 
 class Val(Expr):

python/egglog/examples/lambda_.py

@@ -60,6 +60,7 @@
 import numbers
 import os
 import sys
+from collections.abc import Callable
 from copy import copy
 from types import EllipsisType
 from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, cast
@@ -72,9 +73,10 @@
 from .program_gen import *
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterator
+    from collections.abc import Iterator
     from types import ModuleType
 
+
 # Pretend that exprs are numbers b/c sklearn does isinstance checks
 numbers.Integral.register(RuntimeExpr)
 

python/egglog/exp/array_api.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import egglog as el
 
 
@@ -10,3 +12,15 @@ class Num(el.Expr):
         def __init__(self, value: el.i64Like) -> None: ...
 
     Num(1)  # gets an error "NameError: name 'i64' is not defined"
+
+
+def test_no_import_star_rulesset():
+    """
+    https://github.com/egraphs-good/egglog-python/issues/283
+    """
+
+    @el.ruleset
+    def _rules(_: el.i64Like):
+        return []
+
+    el.EGraph().run(_rules)

python/tests/test_no_import_star.py

@@ -1,6 +1,7 @@
 # mypy: disable-error-code="empty-body"
 from __future__ import annotations
 
+from collections.abc import Callable
 from copy import copy
 from functools import partial
 from typing import TYPE_CHECKING, ClassVar
@@ -10,8 +11,6 @@
 from egglog import *
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
-
     from egglog.runtime import RuntimeExpr