The goal of this PR is to get @kszucs's work, which updates this repo to "support user defined python classes (not just named tuple) and dynamic appliers", in a spot to merge into the main repo. @kszucs if you prefer me not to include these changes in a PR, let me know, I am happy to add my own changes on the main here. I opted to build them off of yours because you had already done some work here and I didn't want to cause a merge conflict if you later decided to merge them upstream.

The TODO items before this is ready for merge currently are:

• [x] See if it is possible to avoid the change making Subst.vec public
• [x] Keep package named snake-egg instead of egg (Can switch to egg if folks prefer. There is already an egg pypi package, so the package name would have to be different).
• [x] Update the existing static typings to include new features
• [x] Make sure the tests pass in CI
• [x] Replace ibis tests with other dynamic applier example

Happy to chat more if anyone else has thoughts here.

EDIT: I have postponed the following items in this PR and could be tackled in follow up PRs to limit the scope in this one.

• [ ] Extend the ability to accept various types by including a way to specify different construct/deconstruct functions, keeping it backwards compatible with @kscucs's definition by using those as default.
• [ ] Possibly start adding some basic documentation that could then be refactored into a full sphinx doc later.
• [ ] Copy over lambda calculus example. This needs conditional rewrites and also the ability to define analysis merge behavior. It would serve as a good test for a more complicated dynamic application.

In this PR I have started adding a .pyi stub file for the library, to expose the types statically for tools like MyPy to consume. I have started adding stubs to not only help future downstream users of this library but to also increase my familiarity to the features this library exposes.

• Mypy docs on stubs
• Py03 docs on adding stub files

I also added MyPy as a development dependency, so that we can use its stubtest tool to verify the stubs match the actual implementation.

I also added a test to github actions to verify the stubs.

TODO:

• [x] Add all classes and functions
• [x] Make tests run on PRs to verify in CI
• [x] FIx submodule issue

Hello, thanks for adding this API!

I'm trying to run the tests on my machine but they seem to fail when I include tests that have a single variable on the right-hand-side. For example, the rewrite rule https://github.com/egraphs-good/snake-egg/blob/c593d6aeb45ac8047facd9ffcb63f59a84bfd996/snake_egg/tests/test_simple.py#L27 causes the following error when I run make test:

thread '<unnamed>' panicked at 'called `Result::unwrap()` on an `Err` value: "Rewrite add-0 refers to unbound var ?a"', src/core.rs:88:63
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
Traceback (most recent call last):
  File "/Users/fmora/Documents/software/snake-egg/snake_egg/tests/test_simple.py", line 27, in <module>
    Rewrite(Add(a, 0), a, name="add-0"),
pyo3_runtime.PanicException: called `Result::unwrap()` on an `Err` value: "Rewrite add-0 refers to unbound var ?a"
make: *** [test] Error 1

When I comment out all the rewrite rules in the test set that have a single variable on the rhs, this error goes away. Have you seen this before? Do you know what is going on?

Thank you!!

Any plans to allow for constant folding or the definition of other types of analyses from Python?

Happy to contribute, but for a full PR think I would need some guidance.

Stubtest was previously raising errors on the tests defined in the module by comparing the stubs for the built-in NamedTuple class with their runtime implementation. We did not implement this class and so we want to ignore any errors raised about it.

This error was raised in #7 but for some reason did not fail in CI until that was merged, so this is a follow-up PR.

It would be useful for egg to allow an extraction method that returns recexpers representing all enodes in an eclass. I am happy to make a pull request if needed.

The stubtest seems to have broken in #7. This disables it in CI for now, since the errors it raises are spurious. It replaces #8, which was overly broad and didn't do what was intended.

Hi, I'm an open source developer and got interested in your project to use it to preprocess IR of computational graphs for onnx, pytorch, tensorflow, jax and on. To make it possible to include your work as a dependency, python wheels have to be provided on pypi. Do you mind me helping to set up github CI publishing installable wheels for common operating systems like linux, windows, macos? So far I've checked there is no package snake-egg on pypi, so you might need to register it there