Mostly working code!

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,9 +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 vecs (TODO)
+- Open files with webbrowser instead of internal graphviz util for better support
 
 ## 8.0.1 (2024-10-24)
 

pyproject.toml

@@ -27,14 +27,16 @@ classifiers = [
     "Typing :: Typed",
 ]
 dependencies = ["typing-extensions", "black", "graphviz", "anywidget"]
+dynamic = ["version"]
 
 [project.optional-dependencies]
 
 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",
@@ -215,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

python/egglog/builtins.py

@@ -11,7 +11,7 @@
 
 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
@@ -418,6 +418,9 @@ def remove(self, index: i64Like) -> Vec[T]: ...
     def set(self, index: i64Like, value: T) -> Vec[T]: ...
 
 
+converter(tuple, Vec, lambda t: Vec(*(convert(x, get_type_args()[0]) for x in t)))
+
+
 class PyObject(Expr, builtin=True):
     def __init__(self, value: object) -> None: ...
 

python/egglog/conversion.py

@@ -148,11 +148,11 @@ def identity(x: object) -> object:
 TYPE_ARGS = ContextVar[tuple[RuntimeClass, ...]]("TYPE_ARGS")
 
 
-def get_type_args() -> tuple[RuntimeClass, ...]:
+def get_type_args() -> tuple[type, ...]:
     """
     Get the type args for the type being converted.
     """
-    return TYPE_ARGS.get()
+    return cast(tuple[type, ...], TYPE_ARGS.get())
 
 
 @contextmanager

python/egglog/egraph.py

@@ -8,6 +8,7 @@
 from collections.abc import Callable, Generator, Iterable
 from contextvars import ContextVar, Token
 from dataclasses import InitVar, dataclass, field
+from functools import partial
 from inspect import Parameter, currentframe, signature
 from types import FrameType, FunctionType
 from typing import (
@@ -1633,7 +1634,9 @@ class UnstableCombinedRuleset(Schedule):
 
     def __post_init__(self, rulesets: list[Ruleset | UnstableCombinedRuleset]) -> None:
         self.schedule = RunDecl(self.__egg_name__, ())
-        self.__egg_decls_thunk__ = Thunk.fn(self._create_egg_decls, *rulesets)
+        # Don't use thunk so that this is re-evaluated each time its requsted, so that additions inside will
+        # be added after its been evaluated once.
+        self.__egg_decls_thunk__ = partial(self._create_egg_decls, *rulesets)
 
     @property
     def __egg_name__(self) -> str: