Add start of compiling expressions to strings

.gitignore

@@ -79,3 +79,7 @@ TODO
 
 .ipynb_checkpoints
 Source.*
+2
+3
+4
+inlined

.pre-commit-config.yaml

@@ -1,3 +1,4 @@
+exclude: ^python/tests/__snapshots__/
 default_language_version:
   python: python3.10
 ci:

Cargo.toml

@@ -24,4 +24,4 @@ name = "egglog.bindings"
 
 [patch.'https://github.com/egraphs-good/egraph-serialize']
 
-egraph-serialize = { git = "https://github.com/saulshanabrook/egraph-serialize", rev = "8dc11836fdade20bf2b20c5f1e2d75bcd005767d" }
+egraph-serialize = { git = "https://github.com/saulshanabrook/egraph-serialize", rev = "a3f6fef9b958a335367d80d51e028c6db886fb6e" }

docs/changelog.md

@@ -4,6 +4,17 @@ _This project uses semantic versioning. Before 1.0.0, this means that every brea
 
 ## Unreleased
 
+### Breaking Changes
+
+### New Features
+
+- Add ability to pass `seminaive` flag to Egraph to replicate `--naive` CLI flag [#48](https://github.com/metadsl/egglog-python/pull/48)
+- Add ability to inline leaves $n$ times instead of just once for visualization [#48](https://github.com/metadsl/egglog-python/pull/48)
+
+### Bug fixes
+
+### Uncategorized
+
 - Added initial supported for Python objects [#31](https://github.com/metadsl/egglog-python/pull/31)
 
   - Renamed `BaseExpr` to `Expr` for succinctness

pyproject.toml

@@ -31,7 +31,7 @@ dependencies = ["typing-extensions", "black", "graphviz"]
 [project.optional-dependencies]
 dev = ["pre-commit", "black", "mypy", "flake8", "isort"]
 
-test = ["pytest", "mypy", "scikit-learn", "array_api_compat"]
+test = ["pytest", "mypy", "scikit-learn", "array_api_compat", "syrupy"]
 
 docs = [
     "pydata-sphinx-theme",
@@ -61,6 +61,7 @@ strict_equality = true
 warn_unused_configs = true
 allow_redefinition = true
 enable_incomplete_feature = ["Unpack", "TypeVarTuple"]
+exclude = ["__snapshots__", "_build", "^conftest.py$"]
 
 [tool.maturin]
 python-source = "python"
@@ -70,3 +71,4 @@ addopts = ["--import-mode=importlib"]
 testpaths = ["python"]
 python_files = ["test_*.py", "test.py"]
 markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"]
+norecursedirs = ["__snapshots__"]

python/egglog/bindings.pyi

@@ -18,7 +18,7 @@ class EGraph:
         *,
         max_functions: Optional[int] = None,
         max_calls_per_function: Optional[int] = None,
-        inline_leaves: bool = False,
+        n_inline_leaves: int = 0,
     ) -> str: ...
     def save_object(self, __o: object, /) -> _Expr: ...
     def load_object(self, __e: _Expr, /) -> object: ...

python/egglog/egraph.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import inspect
 from abc import ABC, abstractmethod
 from contextvars import ContextVar, Token
 from copy import deepcopy
@@ -104,12 +105,12 @@ class _BaseModule(ABC):
     """
 
     # Any modules you want to depend on
-    deps: InitVar[list[Module]] = []
+    modules: InitVar[list[Module]] = []
     # All dependencies flattened
     _flatted_deps: list[Module] = field(init=False, default_factory=list)
     _mod_decls: ModuleDeclarations = field(init=False)
 
-    def __post_init__(self, modules: list[Module] = []) -> None:
+    def __post_init__(self, modules: list[Module]) -> None:
         included_decls = [_BUILTIN_DECLS] if _BUILTIN_DECLS else []
         # Traverse all the included modules to flatten all their dependencies and add to the included declerations
         for mod in modules:
@@ -617,7 +618,7 @@ def let(self, name: str, expr: EXPR) -> EXPR:
 
 @dataclass
 class _Builtins(_BaseModule):
-    def __post_init__(self, modules: list[Module] = []) -> None:
+    def __post_init__(self, modules: list[Module]) -> None:
         """
         Register these declarations as builtins, so others can use them.
         """
@@ -663,13 +664,16 @@ class EGraph(_BaseModule):
     Represents an EGraph instance at runtime
     """
 
-    _egraph: bindings.EGraph = field(repr=False, default_factory=bindings.EGraph)
+    seminaive: InitVar[bool] = True
+
+    _egraph: bindings.EGraph = field(repr=False, init=False)
     # The current declarations which have been pushed to the stack
     _decl_stack: list[Declarations] = field(default_factory=list, repr=False)
     _token: Optional[Token[EGraph]] = None
 
-    def __post_init__(self, modules: list[Module] = []) -> None:
+    def __post_init__(self, modules: list[Module], seminaive) -> None:
         super().__post_init__(modules)
+        self._egraph = bindings.EGraph(seminaive=seminaive)
         for m in self._flatted_deps:
             for o in m._py_objects:
                 self._egraph.save_object(o)
@@ -1430,7 +1434,15 @@ def _command_generator(gen: CommandGenerator) -> Iterable[Command]:
     """
     Calls the function with variables of the type and name of the arguments.
     """
-    hints = get_type_hints(gen)
+    # 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
+    current_frame = inspect.currentframe()
+    assert current_frame
+    register_frame = current_frame.f_back
+    assert register_frame
+    original_frame = register_frame.f_back
+    assert original_frame
+    hints = get_type_hints(gen, gen.__globals__, original_frame.f_locals)
     args = (_var(p.name, hints[p.name]) for p in signature(gen).parameters.values())
     return gen(*args)
 

python/egglog/examples/sklearn_lda.py

@@ -38,8 +38,10 @@ def fit(X, y):
 
 with EGraph([array_api_module]) as egraph:
     egraph.register(res)
-    egraph.run((run() * 10).saturate())
-    egraph.run((run() * 10).saturate())
+    egraph.run((run() * 10))
+    # egraph.run((run() * 10).saturate())
+    # egraph.graphviz(n_inline_leaves=3).render("3", view=True)
+
     res = egraph.extract(expr=res)
     print(res)
     # egraph.display()