Merge pull request #257 from egraphs-good/fix-loopnest

Adds working loopnest example

Author: saulshanabrook

.pre-commit-config.yaml

@@ -6,3 +6,8 @@ repos:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
       - id: ruff-format
+  - repo: https://github.com/astral-sh/uv-pre-commit
+    # uv version.
+    rev: 0.5.31
+    hooks:
+      - id: uv-lock

docs/changelog.md

@@ -9,7 +9,11 @@ _This project uses semantic versioning_
 - Add better error message when using @function in class (thanks @shinawy)
 - Add error method if `@method` decorator is in wrong place
 - Subsumes lambda functions after replacing
-- Add working loopnest test
+- Add working loopnest test and rewrite array api suport to be more general
+- Improve tracebacks on failing conversions.
+- Use `add_note` for exception to add more context, instead of raising a new exception, to make it easier to debug.
+- Add conversions from generic types to be supported at runtime and typing level (so can go from `(1, 2, 3)` to `TupleInt`)
+- Open files with webbrowser instead of internal graphviz util for better support
 
 ## 8.0.1 (2024-10-24)
 

docs/how-to-guides.md

@@ -9,7 +9,7 @@ file_format: mystnb
 You can provide your program in a special DSL language. You can parse this with {meth}`egglog.bindings.parse_program` and then run the result with You can parse this with {meth}`egglog.bindings.EGraph.run_program`::
 
 ```{code-cell}
-from egglog.bindings import EGraph
+from egglog.bindings import EGraph, parse_program
 
 egraph = EGraph()
 commands = parse_program("(check (= (+ 1 2) 3))")

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "maturin"
 name = "egglog"
 description = "e-graphs in Python built around the the egglog rust library"
 readme = "README.md"
-dynamic = ["version"] 
+dynamic = ["version"]
 license = { text = "MIT" }
 requires-python = ">=3.10"
 classifiers = [
@@ -34,8 +34,9 @@ dependencies = ["typing-extensions", "black", "graphviz", "anywidget"]
 array = [
     "scikit-learn",
     "array_api_compat",
-    "numba==0.59.1",
-    "llvmlite==0.42.0",
+    "numba>=0.59.1",
+    "llvmlite>=0.42.0",
+    "numpy>2",
 ]
 dev = [
     "ruff",
@@ -216,7 +217,7 @@ preview = true
 
 [tool.ruff.lint.per-file-ignores]
 # Don't require annotations for tests
-"python/tests/**" = ["ANN001", "ANN201"]
+"python/tests/**" = ["ANN001", "ANN201", "INP001"]
 
 [tool.mypy]
 ignore_missing_imports = true
@@ -239,7 +240,6 @@ python_files = ["test_*.py", "test.py"]
 markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"]
 norecursedirs = ["__snapshots__"]
 filterwarnings = [
-    "error",
     "ignore::numba.core.errors.NumbaPerformanceWarning",
     "ignore::pytest_benchmark.logger.PytestBenchmarkWarning",
     # https://github.com/manzt/anywidget/blob/d38bb3f5f9cfc7e49e2ff1aa1ba994d66327cb02/pyproject.toml#L120

python/egglog/__init__.py

@@ -4,7 +4,7 @@
 
 from . import config, ipython_magic  # noqa: F401
 from .builtins import *  # noqa: UP029
-from .conversion import convert, converter  # noqa: F401
+from .conversion import ConvertError, convert, converter, get_type_args  # noqa: F401
 from .egraph import *
 
 del ipython_magic

python/egglog/builtins.py

@@ -5,13 +5,13 @@
 
 from __future__ import annotations
 
-from functools import partial
+from functools import partial, reduce
 from types import FunctionType
 from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Union, cast, overload
 
 from typing_extensions import TypeVarTuple, Unpack
 
-from .conversion import converter, get_type_args
+from .conversion import convert, converter, get_type_args
 from .egraph import Expr, Unit, function, get_current_ruleset, method
 from .functionalize import functionalize
 from .runtime import RuntimeClass, RuntimeExpr, RuntimeFunction
@@ -22,24 +22,27 @@
 
 
 __all__ = [
-    "i64",
-    "i64Like",
-    "f64",
-    "f64Like",
     "Bool",
     "BoolLike",
-    "String",
-    "StringLike",
     "Map",
+    "MapLike",
+    "PyObject",
     "Rational",
     "Set",
+    "SetLike",
+    "String",
+    "StringLike",
+    "UnstableFn",
     "Vec",
+    "VecLike",
+    "f64",
+    "f64Like",
+    "i64",
+    "i64Like",
     "join",
-    "PyObject",
     "py_eval",
-    "py_exec",
     "py_eval_fn",
-    "UnstableFn",
+    "py_exec",
 ]
 
 
@@ -210,6 +213,9 @@ def __truediv__(self, other: f64Like) -> f64: ...
     @method(egg_fn="%")
     def __mod__(self, other: f64Like) -> f64: ...
 
+    @method(egg_fn="^")
+    def __pow__(self, other: f64Like) -> f64: ...
+
     def __radd__(self, other: f64Like) -> f64: ...
 
     def __rsub__(self, other: f64Like) -> f64: ...
@@ -282,6 +288,22 @@ def remove(self, key: T) -> Map[T, V]: ...
     def rebuild(self) -> Map[T, V]: ...
 
 
+TO = TypeVar("TO")
+VO = TypeVar("VO")
+
+converter(
+    dict,
+    Map,
+    lambda t: reduce(
+        (lambda acc, kv: acc.insert(convert(kv[0], get_type_args()[0]), convert(kv[1], get_type_args()[1]))),
+        t.items(),
+        Map[get_type_args()].empty(),  # type: ignore[misc]
+    ),
+)
+
+MapLike: TypeAlias = Map[T, V] | dict[TO, VO]
+
+
 class Set(Expr, Generic[T], builtin=True):
     @method(egg_fn="set-of")
     def __init__(self, *args: T) -> None: ...
@@ -315,6 +337,17 @@ def __and__(self, other: Set[T]) -> Set[T]: ...
     def rebuild(self) -> Set[T]: ...
 
 
+converter(
+    set,
+    Set,
+    lambda t: Set[get_type_args()[0]](  # type: ignore[misc,operator]
+        *(convert(x, get_type_args()[0]) for x in t)
+    ),
+)
+
+SetLike: TypeAlias = Set[T] | set[TO]
+
+
 class Rational(Expr, builtin=True):
     @method(egg_fn="rational")
     def __init__(self, num: i64Like, den: i64Like) -> None: ...
@@ -415,6 +448,18 @@ def remove(self, index: i64Like) -> Vec[T]: ...
     def set(self, index: i64Like, value: T) -> Vec[T]: ...
 
 
+for sequence_type in (list, tuple):
+    converter(
+        sequence_type,
+        Vec,
+        lambda t: Vec[get_type_args()[0]](  # type: ignore[misc,operator]
+            *(convert(x, get_type_args()[0]) for x in t)
+        ),
+    )
+
+VecLike: TypeAlias = Vec[T] | tuple[TO, ...] | list[TO]
+
+
 class PyObject(Expr, builtin=True):
     def __init__(self, value: object) -> None: ...