Merge pull request #301 from egraphs-good/fix-import-f64

Fix f64Like import

Author: saulshanabrook

Cargo.lock

@@ -4,10 +4,12 @@ _This project uses semantic versioning_
 
 ## UNRELEASED
 
+- fix using `f64Like` when not importing star (also properly includes removal of `Callable` special case from previous release).
+
 ## 10.0.1 (2025-04-06)
 
 - 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
+  - Also stops special casing including `Callable` as a global. So if you previously included this in a `TYPE_CHECKING` block so it wasn't
     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)

docs/changelog.md

@@ -157,8 +157,7 @@ If you want to have this work with the static type checker, you can define your
 the `Expr` class as the first item in the union. For example, in this case you could then define:
 
 ```{code-cell} python
-from typing import Union
-MathLike = Union[Math, i64Like, StringLike]
+MathLike = Math | i64Like | StringLike
 
 @function
 def some_math_fn(x: MathLike) -> MathLike:

docs/reference/python-integration.md

@@ -9,7 +9,7 @@
 from fractions import Fraction
 from functools import partial, reduce
 from types import FunctionType, MethodType
-from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Union, cast, overload
+from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, cast, overload
 
 from typing_extensions import TypeVarTuple, Unpack
 
@@ -101,8 +101,6 @@ def join(*strings: StringLike) -> String: ...
 
 converter(str, String, String)
 
-BoolLike: TypeAlias = Union["Bool", bool]
-
 
 class Bool(BuiltinExpr, egg_sort="bool"):
     @method(preserve=True)
@@ -133,10 +131,10 @@ def __xor__(self, other: BoolLike) -> Bool: ...
     def implies(self, other: BoolLike) -> Bool: ...
 
 
-converter(bool, Bool, Bool)
+BoolLike: TypeAlias = Bool | bool
 
-# The types which can be convertered into an i64
-i64Like: TypeAlias = Union["i64", int]  # noqa: N816, PYI042
+
+converter(bool, Bool, Bool)
 
 
 class i64(BuiltinExpr):  # noqa: N801
@@ -248,16 +246,16 @@ def bool_le(self, other: i64Like) -> Bool: ...
     def bool_ge(self, other: i64Like) -> Bool: ...
 
 
+# The types which can be convertered into an i64
+i64Like: TypeAlias = i64 | int  # noqa: N816, PYI042
+
 converter(int, i64, i64)
 
 
 @function(builtin=True, egg_fn="count-matches")
 def count_matches(s: StringLike, pattern: StringLike) -> i64: ...
 
 
-f64Like: TypeAlias = Union["f64", float]  # noqa: N816, PYI042
-
-
 class f64(BuiltinExpr):  # noqa: N801
     @method(preserve=True)
     def eval(self) -> float:
@@ -337,6 +335,9 @@ def from_i64(cls, i: i64) -> f64: ...
     def to_string(self) -> String: ...
 
 
+f64Like: TypeAlias = f64 | float  # noqa: N816, PYI042
+
+
 converter(float, f64, f64)
 
 

python/egglog/builtins.py

@@ -569,16 +569,11 @@ def _fn_decl(
     if not isinstance(fn, FunctionType):
         raise NotImplementedError(f"Can only generate function decls for functions not {fn}  {type(fn)}")
 
-    hint_globals = fn.__globals__.copy()
-    # Copy Callable into global if not present bc sometimes it gets automatically removed by ruff to type only block
-    # https://docs.astral.sh/ruff/rules/typing-only-standard-library-import/
-    if "Callable" not in hint_globals:
-        hint_globals["Callable"] = Callable
     # Instead of passing both globals and locals, just pass the globals. Otherwise, for some reason forward references
     # won't be resolved correctly
     # We need this to be false so it returns "__forward_value__" https://github.com/python/cpython/blob/440ed18e08887b958ad50db1b823e692a747b671/Lib/typing.py#L919
     # https://github.com/egraphs-good/egglog-python/issues/210
-    hint_globals.update(hint_locals)
+    hint_globals = {**fn.__globals__, **hint_locals}
     hints = get_type_hints(fn, hint_globals)
 
     params = list(signature(fn).parameters.values())

python/egglog/egraph.py

@@ -5,12 +5,10 @@
 
 from __future__ import annotations
 
-from typing import TypeAlias, Union
+from typing import TypeAlias
 
 from egglog import *
 
-ProgramLike: TypeAlias = Union["Program", StringLike]
-
 
 class Program(Expr):
     """
@@ -97,6 +95,9 @@ def is_identifer(self) -> Bool:
         """
 
 
+ProgramLike: TypeAlias = Program | StringLike
+
+
 converter(String, Program, Program)
 
 

python/egglog/exp/program_gen.py

@@ -5,7 +5,7 @@
 import pathlib
 from copy import copy
 from fractions import Fraction
-from typing import ClassVar, TypeAlias, Union
+from typing import ClassVar, TypeAlias
 
 import pytest
 
@@ -451,8 +451,6 @@ def __init__(self, name: StringLike) -> None: ...
 
         def __add__(self, other: Int) -> Int: ...
 
-    NDArrayLike: TypeAlias = Union[Int, "NDArray"]
-
     class NDArray(Expr):
         def __init__(self, name: StringLike) -> None: ...
 
@@ -465,6 +463,8 @@ def to_int(self) -> Int: ...
         @classmethod
         def from_int(cls, other: Int) -> NDArray: ...
 
+    NDArrayLike: TypeAlias = NDArray | Int
+
     converter(Int, NDArray, NDArray.from_int)
     converter(NDArray, Int, lambda a: a.to_int(), 100)
 

python/tests/test_high_level.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import egglog as el
+from egglog import f64Like
 
 
 def test_no_import_star():
@@ -11,7 +12,18 @@ def test_no_import_star():
     class Num(el.Expr):
         def __init__(self, value: el.i64Like) -> None: ...
 
-    Num(1)  # gets an error "NameError: name 'i64' is not defined"
+    Num(1)
+
+
+def test_f64_import():
+    """
+    For some reason this wasn't working until we moved the union definition below the class
+    """
+
+    class Num(el.Expr):
+        def __init__(self, value: f64Like) -> None: ...
+
+    Num(1.0)
 
 
 def test_no_import_star_rulesset():