egglog 7.1.0__tar.gz → 7.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of egglog might be problematic. Click here for more details.

Files changed (117) hide show
  1. {egglog-7.1.0 → egglog-7.2.0}/.github/workflows/version.yml +1 -1
  2. {egglog-7.1.0 → egglog-7.2.0}/Cargo.lock +1 -1
  3. {egglog-7.1.0 → egglog-7.2.0}/Cargo.toml +1 -1
  4. {egglog-7.1.0 → egglog-7.2.0}/PKG-INFO +7 -7
  5. {egglog-7.1.0 → egglog-7.2.0}/docs/changelog.md +5 -0
  6. egglog-7.2.0/docs/reference/contributing.md +91 -0
  7. {egglog-7.1.0 → egglog-7.2.0}/docs/reference/python-integration.md +84 -0
  8. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/declarations.py +66 -7
  9. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/egraph.py +141 -75
  10. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/egraph_state.py +51 -20
  11. egglog-7.2.0/python/egglog/examples/higher_order_functions.py +50 -0
  12. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/exp/array_api.py +4 -1
  13. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/pretty.py +15 -5
  14. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/runtime.py +9 -6
  15. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/thunk.py +16 -4
  16. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/type_constraint_solver.py +5 -4
  17. {egglog-7.1.0 → egglog-7.2.0}/python/tests/conftest.py +2 -0
  18. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_high_level.py +115 -2
  19. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_runtime.py +1 -1
  20. egglog-7.1.0/docs/reference/contributing.md +0 -15
  21. {egglog-7.1.0 → egglog-7.2.0}/.github/dependabot.yml +0 -0
  22. {egglog-7.1.0 → egglog-7.2.0}/.github/workflows/CI.yml +0 -0
  23. {egglog-7.1.0 → egglog-7.2.0}/.gitignore +0 -0
  24. {egglog-7.1.0 → egglog-7.2.0}/.pre-commit-config.yaml +0 -0
  25. {egglog-7.1.0 → egglog-7.2.0}/.readthedocs.yaml +0 -0
  26. {egglog-7.1.0 → egglog-7.2.0}/CITATION.cff +0 -0
  27. {egglog-7.1.0 → egglog-7.2.0}/LICENSE +0 -0
  28. {egglog-7.1.0 → egglog-7.2.0}/README.md +0 -0
  29. {egglog-7.1.0 → egglog-7.2.0}/conftest.py +0 -0
  30. {egglog-7.1.0 → egglog-7.2.0}/docs/.gitignore +0 -0
  31. {egglog-7.1.0 → egglog-7.2.0}/docs/_templates/comments.html +0 -0
  32. {egglog-7.1.0 → egglog-7.2.0}/docs/conf.py +0 -0
  33. {egglog-7.1.0 → egglog-7.2.0}/docs/environment.yml +0 -0
  34. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/.gitignore +0 -0
  35. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/2023_07_presentation.ipynb +0 -0
  36. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/2023_11_09_portland_state.ipynb +0 -0
  37. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/2023_11_17_pytensor.ipynb +0 -0
  38. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/2023_11_pydata_lightning_talk.ipynb +0 -0
  39. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/2023_12_02_congruence_closure-1.png +0 -0
  40. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/2023_12_02_congruence_closure-2.png +0 -0
  41. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/2023_12_02_congruence_closure.md +0 -0
  42. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/2024_03_17_community_talk.ipynb +0 -0
  43. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/2024_03_17_map.svg +0 -0
  44. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/big_graph.svg +0 -0
  45. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/define_and_define.md +0 -0
  46. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/ecosystem-graph.png +0 -0
  47. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/egg.png +0 -0
  48. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/indexing_pushdown.ipynb +0 -0
  49. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/moa.png +0 -0
  50. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/optional_values.md +0 -0
  51. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation/pldi_2023_presentation.ipynb +0 -0
  52. {egglog-7.1.0 → egglog-7.2.0}/docs/explanation.md +0 -0
  53. {egglog-7.1.0 → egglog-7.2.0}/docs/how-to-guides.md +0 -0
  54. {egglog-7.1.0 → egglog-7.2.0}/docs/index.md +0 -0
  55. {egglog-7.1.0 → egglog-7.2.0}/docs/reference/bindings.md +0 -0
  56. {egglog-7.1.0 → egglog-7.2.0}/docs/reference/egglog-translation.md +0 -0
  57. {egglog-7.1.0 → egglog-7.2.0}/docs/reference/high-level.md +0 -0
  58. {egglog-7.1.0 → egglog-7.2.0}/docs/reference/usage.md +0 -0
  59. {egglog-7.1.0 → egglog-7.2.0}/docs/reference.md +0 -0
  60. {egglog-7.1.0 → egglog-7.2.0}/docs/tutorials/getting-started.ipynb +0 -0
  61. {egglog-7.1.0 → egglog-7.2.0}/docs/tutorials/screenshot-1.png +0 -0
  62. {egglog-7.1.0 → egglog-7.2.0}/docs/tutorials/screenshot-2.png +0 -0
  63. {egglog-7.1.0 → egglog-7.2.0}/docs/tutorials/sklearn.ipynb +0 -0
  64. {egglog-7.1.0 → egglog-7.2.0}/docs/tutorials.md +0 -0
  65. {egglog-7.1.0 → egglog-7.2.0}/increment_version.py +0 -0
  66. {egglog-7.1.0 → egglog-7.2.0}/pyproject.toml +0 -0
  67. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/__init__.py +0 -0
  68. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/bindings.pyi +0 -0
  69. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/builtins.py +0 -0
  70. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/config.py +0 -0
  71. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/conversion.py +0 -0
  72. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/examples/README.rst +0 -0
  73. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/examples/__init__.py +0 -0
  74. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/examples/bool.py +0 -0
  75. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/examples/eqsat_basic.py +0 -0
  76. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/examples/fib.py +0 -0
  77. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/examples/lambda_.py +0 -0
  78. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/examples/matrix.py +0 -0
  79. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/examples/ndarrays.py +0 -0
  80. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/examples/resolution.py +0 -0
  81. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/examples/schedule_demo.py +0 -0
  82. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/exp/__init__.py +0 -0
  83. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/exp/array_api_jit.py +0 -0
  84. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/exp/array_api_numba.py +0 -0
  85. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/exp/array_api_program_gen.py +0 -0
  86. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/exp/program_gen.py +0 -0
  87. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/exp/siu_examples.py +0 -0
  88. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/graphviz_widget.py +0 -0
  89. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/ipython_magic.py +0 -0
  90. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/py.typed +0 -0
  91. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/widget.css +0 -0
  92. {egglog-7.1.0 → egglog-7.2.0}/python/egglog/widget.js +0 -0
  93. {egglog-7.1.0 → egglog-7.2.0}/python/tests/__init__.py +0 -0
  94. {egglog-7.1.0 → egglog-7.2.0}/python/tests/__snapshots__/test_array_api/TestLDA.test_optimize.py +0 -0
  95. {egglog-7.1.0 → egglog-7.2.0}/python/tests/__snapshots__/test_array_api/TestLDA.test_source_optimized.py +0 -0
  96. {egglog-7.1.0 → egglog-7.2.0}/python/tests/__snapshots__/test_array_api/TestLDA.test_trace.py +0 -0
  97. {egglog-7.1.0 → egglog-7.2.0}/python/tests/__snapshots__/test_program_gen/test_to_string.py +0 -0
  98. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_array_api.py +0 -0
  99. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_bindings.py +0 -0
  100. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_convert.py +0 -0
  101. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_modules.py +0 -0
  102. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_pretty.py +0 -0
  103. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_program_gen.py +0 -0
  104. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_py_object_sort.py +0 -0
  105. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_type_constraint_solver.py +0 -0
  106. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_typing.py +0 -0
  107. {egglog-7.1.0 → egglog-7.2.0}/python/tests/test_unstable_fn.py +0 -0
  108. {egglog-7.1.0 → egglog-7.2.0}/rust-toolchain.toml +0 -0
  109. {egglog-7.1.0 → egglog-7.2.0}/src/conversions.rs +0 -0
  110. {egglog-7.1.0 → egglog-7.2.0}/src/egraph.rs +0 -0
  111. {egglog-7.1.0 → egglog-7.2.0}/src/error.rs +0 -0
  112. {egglog-7.1.0 → egglog-7.2.0}/src/lib.rs +0 -0
  113. {egglog-7.1.0 → egglog-7.2.0}/src/py_object_sort.rs +0 -0
  114. {egglog-7.1.0 → egglog-7.2.0}/src/serialize.rs +0 -0
  115. {egglog-7.1.0 → egglog-7.2.0}/src/utils.rs +0 -0
  116. {egglog-7.1.0 → egglog-7.2.0}/stubtest_allow +0 -0
  117. {egglog-7.1.0 → egglog-7.2.0}/test-data/unit/check-high-level.test +0 -0
@@ -128,7 +128,7 @@ jobs:
128
128
  - run: |
129
129
  git tag "v$VERSION"
130
130
  git push --tags
131
- gh pr merge --admin --delete-branch
131
+ gh pr merge --admin --delete-branch --merge
132
132
  env:
133
133
  VERSION: ${{ needs.bump.outputs.version }}
134
134
  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -312,7 +312,7 @@ dependencies = [
312
312
 
313
313
  [[package]]
314
314
  name = "egglog-python"
315
- version = "7.1.0"
315
+ version = "7.2.0"
316
316
  dependencies = [
317
317
  "egglog",
318
318
  "egraph-serialize",
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "egglog-python"
3
- version = "7.1.0"
3
+ version = "7.2.0"
4
4
  edition = "2021"
5
5
 
6
6
  # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: egglog
3
- Version: 7.1.0
3
+ Version: 7.2.0
4
4
  Classifier: Environment :: MacOS X
5
5
  Classifier: Environment :: Win32 (MS Windows)
6
6
  Classifier: Intended Audience :: Developers
@@ -30,6 +30,11 @@ Requires-Dist: scikit-learn; extra == 'array'
30
30
  Requires-Dist: array_api_compat; extra == 'array'
31
31
  Requires-Dist: numba==0.59.1; extra == 'array'
32
32
  Requires-Dist: llvmlite==0.42.0; extra == 'array'
33
+ Requires-Dist: pre-commit; extra == 'dev'
34
+ Requires-Dist: ruff; extra == 'dev'
35
+ Requires-Dist: mypy; extra == 'dev'
36
+ Requires-Dist: anywidget[dev]; extra == 'dev'
37
+ Requires-Dist: egglog[docs,test]; extra == 'dev'
33
38
  Requires-Dist: pydata-sphinx-theme; extra == 'docs'
34
39
  Requires-Dist: myst-nb; extra == 'docs'
35
40
  Requires-Dist: sphinx-autodoc-typehints; extra == 'docs'
@@ -42,15 +47,10 @@ Requires-Dist: egglog[array]; extra == 'docs'
42
47
  Requires-Dist: line-profiler; extra == 'docs'
43
48
  Requires-Dist: sphinxcontrib-mermaid; extra == 'docs'
44
49
  Requires-Dist: ablog; extra == 'docs'
45
- Requires-Dist: pre-commit; extra == 'dev'
46
- Requires-Dist: ruff; extra == 'dev'
47
- Requires-Dist: mypy; extra == 'dev'
48
- Requires-Dist: anywidget[dev]; extra == 'dev'
49
- Requires-Dist: egglog[docs,test]; extra == 'dev'
50
50
  Provides-Extra: test
51
51
  Provides-Extra: array
52
- Provides-Extra: docs
53
52
  Provides-Extra: dev
53
+ Provides-Extra: docs
54
54
  License-File: LICENSE
55
55
  Summary: e-graphs in Python built around the the egglog rust library
56
56
  License: MIT
@@ -4,6 +4,11 @@ _This project uses semantic versioning_
4
4
 
5
5
  ## UNRELEASED
6
6
 
7
+ ## 7.2.0 (2024-05-23)
8
+
9
+ - Adds ability to use function bodies as default rewrites ([#167](https://github.com/egraphs-good/egglog-python/pull/167))
10
+ - Fixed bug with creating empty maps and adding to maps ([#168](https://github.com/egraphs-good/egglog-python/pull/168))
11
+
7
12
  ## 7.1.0 (2024-05-03)
8
13
 
9
14
  ## New Feaatures
@@ -0,0 +1,91 @@
1
+ # Contributing
2
+
3
+ ## Development
4
+
5
+ This package is in active development and welcomes contributions!
6
+
7
+ Feel free to bring up any questions/comments on [the Zulip chat](https://egraphs.zulipchat.com/) or open an issue.
8
+
9
+ All feedback is welcome and encouraged, it's great to hear about anything that works well or could be improved.
10
+
11
+ ### Getting Started
12
+
13
+ To get started locally developing with this project, fork it and clone it to your local machine.
14
+
15
+ Using [the Github CLI](https://github.com/cli/cli#installation) this would be:
16
+
17
+ ```bash
18
+ brew install gh
19
+ gh repo fork egraphs-good/egglog-python --clone
20
+ cd egglog-python
21
+ ```
22
+
23
+ Then [install Rust](https://www.rust-lang.org/tools/install) and get a Python environment set up with a compatible version. Using [miniconda](https://formulae.brew.sh/cask/miniconda) this would be:
24
+
25
+ ```bash
26
+ curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
27
+ brew install --cask miniconda
28
+ conda create -n egglog-python python
29
+ conda activate egglog-python
30
+ ```
31
+
32
+ Then install the package in editable mode with the development dependencies:
33
+
34
+ ```bash
35
+ pip install -e .[dev]
36
+ ```
37
+
38
+ Anytime you change the rust code, you can run `pip install -e .` to recompile the rust code.
39
+
40
+ ### Running Tests
41
+
42
+ To run the tests, you can use the `pytest` command:
43
+
44
+ ```bash
45
+ pytest
46
+ ```
47
+
48
+ All code must pass ruff linters and formaters. This will be checked automatically by the pre-commit if you run `pre-commit install`.
49
+
50
+ To run it manually, you can use:
51
+
52
+ ```bash
53
+ pre-commit run --all-files ruff
54
+ ```
55
+
56
+ If you make changes to the rust bindings, you can check that the stub files accurately reflect the rust code by running:
57
+
58
+ ```bash
59
+ pre-commit run --all-files --hook-stage manual stubtest
60
+ ```
61
+
62
+ All code must all pass MyPy type checking. To run that locally use:
63
+
64
+ ```bash
65
+ pre-commit run --all-files --hook-stage manual mypy
66
+ ```
67
+
68
+ Finally, to build the docs locally and test that they work, you can run:
69
+
70
+ ```bash
71
+ pre-commit run --all-files --hook-stage manual docs
72
+ ```
73
+
74
+ ## Making changes
75
+
76
+ All changes that impact users should be documented in the `docs/changelog.md` file. Please also add tests for any new features
77
+ or bug fixes.
78
+
79
+ When you are ready to submit your changes, please open a pull request. The CI will run the tests and check the code style.
80
+
81
+ ## Documentation
82
+
83
+ We use the [Diátaxis framework](https://diataxis.fr/) to organize our documentation. The "explanation" section has
84
+ been renamed to "Blog" since most of the content there is more like a blog post than a reference manual. It uses
85
+ the [ABlog](https://ablog.readthedocs.io/en/stable/index.html#how-it-works) extension.
86
+
87
+ ## Governance
88
+
89
+ The governance is currently informal, with Saul Shanabrook as the lead maintainer. If the project grows and there
90
+ are more contributors, we will formalize the governance structure in a way to allow it to be multi-stakeholder and
91
+ to spread out the power and responsibility.
@@ -135,6 +135,7 @@ class Math(Expr):
135
135
  def var(cls, name: StringLike) -> Math: ...
136
136
 
137
137
  def __add__(self, other: Math) -> Math: ...
138
+ def __mul__(self, other: Math) -> Math: ...
138
139
 
139
140
  converter(i64, Math, Math)
140
141
  converter(String, Math, Math.var)
@@ -426,3 +427,86 @@ check_eq(added_two, MathList.EMPTY.append(Math(2) + Math(1)), math_list_ruleset.
426
427
  Note that this is all built on the [unstable function support added as a sort to egglog](https://github.com/egraphs-good/egglog/pull/348).
427
428
  While this sort is exposed directly at the high level with the `UnstableFn` class, we don't reccomend depending on it directly, and instead
428
429
  using the builtin Python type annotations. This will allow us to change the implementation in the future without breaking user code.
430
+
431
+ ## Default Replacements
432
+
433
+ When defining a function or a constant, you can also provide a default replacement value. This is useful when
434
+ you might want both the original value and the replaced value in the e-graph, so that later rules could reference either.
435
+
436
+ ```{code-cell} python
437
+ @function
438
+ def math_float(f: f64Like) -> Math:
439
+ ...
440
+
441
+
442
+ # Can add a default replacement value for a constants
443
+ pi = constant("pi", Math, math_float(3.14))
444
+
445
+
446
+ # or for a function by providing a body
447
+ @function
448
+ def square(x: Math) -> Math:
449
+ return x * x
450
+
451
+ # thse rewrites will be added to the e-graph under the default ruleset
452
+ egraph = EGraph()
453
+ egraph.register(pi)
454
+ egraph.register(square(Math.var('x')))
455
+ egraph.run(1)
456
+ egraph.check(eq(pi).to(math_float(3.14)))
457
+ egraph.check(eq(square(Math.var('x'))).to(Math.var('x') * Math.var('x')))
458
+ egraph
459
+ ```
460
+
461
+ This is equivalent to adding the rewrite rules to the e-graph directly, like this, but just more succinct:
462
+
463
+ ```python
464
+ x = var("x", Math)
465
+ egraph.register(rewrite(pi).to(math_float(3.14)))
466
+ egraph.register(rewrite(square(x)).to(x * x))
467
+ ```
468
+
469
+ You can also specify a ruleset to add the rewrites to, by passing in the `ruleset` keyword argument:
470
+
471
+ ```{code-cell} python
472
+ math_ruleset = ruleset()
473
+
474
+ e_constant = constant("e", Math, math_float(2.71), ruleset=math_ruleset)
475
+
476
+ @function(ruleset=math_ruleset)
477
+ def cube(x: Math) -> Math:
478
+ return x * x * x
479
+
480
+
481
+ egraph.register(e_constant)
482
+ egraph.register(cube(Math.var('x')))
483
+ egraph.run(math_ruleset)
484
+ egraph.check(eq(e_constant).to(math_float(2.71)))
485
+ egraph.check(eq(cube(Math.var('x'))).to(Math.var('x') * Math.var('x') * Math.var('x')))
486
+ ```
487
+
488
+ ### Default Replacement for Classes
489
+
490
+ In classes, you can also provide a default replacement value for constants and methods, and an optional ruleset on the class constructor:
491
+
492
+ ```{code-cell} python
493
+ other_math_ruleset = ruleset()
494
+
495
+
496
+ class WrappedMath(Expr, ruleset=other_math_ruleset):
497
+ PI: ClassVar[Math] = math_float(3.14)
498
+
499
+ def __init__(self, x: Math) -> None: ...
500
+
501
+ def double(self) -> WrappedMath:
502
+ return self + self
503
+
504
+ def __add__(self, other: WrappedMath) -> WrappedMath: ...
505
+
506
+ x = WrappedMath(WrappedMath.PI).double()
507
+ egraph = EGraph()
508
+ egraph.register(x)
509
+ egraph.run(other_math_ruleset * 2)
510
+ egraph.check(eq(x).to(WrappedMath(math_float(3.14)) + WrappedMath(math_float(3.14))))
511
+ egraph
512
+ ```
@@ -71,6 +71,8 @@ __all__ = [
71
71
  "CommandDecl",
72
72
  "SpecialFunctions",
73
73
  "FunctionSignature",
74
+ "DefaultRewriteDecl",
75
+ "InitRef",
74
76
  ]
75
77
 
76
78
 
@@ -80,7 +82,13 @@ class DelayedDeclerations:
80
82
 
81
83
  @property
82
84
  def __egg_decls__(self) -> Declarations:
83
- return self.__egg_decls_thunk__()
85
+ try:
86
+ return self.__egg_decls_thunk__()
87
+ # Catch attribute error, so that it isn't bubbled up as a missing attribute and fallbacks on `__getattr__`
88
+ # instead raise explicitly
89
+ except AttributeError as err:
90
+ msg = "Error resolving declerations"
91
+ raise RuntimeError(msg) from err
84
92
 
85
93
 
86
94
  @runtime_checkable
@@ -113,6 +121,12 @@ class Declarations:
113
121
  _classes: dict[str, ClassDecl] = field(default_factory=dict)
114
122
  _rulesets: dict[str, RulesetDecl | CombinedRulesetDecl] = field(default_factory=lambda: {"": RulesetDecl([])})
115
123
 
124
+ @property
125
+ def default_ruleset(self) -> RulesetDecl:
126
+ ruleset = self._rulesets[""]
127
+ assert isinstance(ruleset, RulesetDecl)
128
+ return ruleset
129
+
116
130
  @classmethod
117
131
  def create(cls, *others: DeclerationsLike) -> Declarations:
118
132
  others = upcast_declerations(others)
@@ -127,7 +141,7 @@ class Declarations:
127
141
 
128
142
  def copy(self) -> Declarations:
129
143
  new = Declarations()
130
- new |= self
144
+ self.update_other(new)
131
145
  return new
132
146
 
133
147
  def update(self, *others: DeclerationsLike) -> None:
@@ -154,9 +168,13 @@ class Declarations:
154
168
  other._functions |= self._functions
155
169
  other._classes |= self._classes
156
170
  other._constants |= self._constants
171
+ # Must combine rulesets bc the empty ruleset might be different, bc DefaultRewriteDecl
172
+ # is added to functions.
173
+ combined_default_rules: set[RewriteOrRuleDecl] = {*self.default_ruleset.rules, *other.default_ruleset.rules}
157
174
  other._rulesets |= self._rulesets
175
+ other._rulesets[""] = RulesetDecl(list(combined_default_rules))
158
176
 
159
- def get_callable_decl(self, ref: CallableRef) -> CallableDecl:
177
+ def get_callable_decl(self, ref: CallableRef) -> CallableDecl: # noqa: PLR0911
160
178
  match ref:
161
179
  case FunctionRef(name):
162
180
  return self._functions[name]
@@ -170,8 +188,29 @@ class Declarations:
170
188
  return self._classes[class_name].class_methods[name]
171
189
  case PropertyRef(class_name, property_name):
172
190
  return self._classes[class_name].properties[property_name]
191
+ case InitRef(class_name):
192
+ init_fn = self._classes[class_name].init
193
+ assert init_fn
194
+ return init_fn
173
195
  assert_never(ref)
174
196
 
197
+ def set_function_decl(
198
+ self, ref: FunctionRef | MethodRef | ClassMethodRef | PropertyRef | InitRef, decl: FunctionDecl
199
+ ) -> None:
200
+ match ref:
201
+ case FunctionRef(name):
202
+ self._functions[name] = decl
203
+ case MethodRef(class_name, method_name):
204
+ self._classes[class_name].methods[method_name] = decl
205
+ case ClassMethodRef(class_name, name):
206
+ self._classes[class_name].class_methods[name] = decl
207
+ case PropertyRef(class_name, property_name):
208
+ self._classes[class_name].properties[property_name] = decl
209
+ case InitRef(class_name):
210
+ self._classes[class_name].init = decl
211
+ case _:
212
+ assert_never(ref)
213
+
175
214
  def has_method(self, class_name: str, method_name: str) -> bool | None:
176
215
  """
177
216
  Returns whether the given class has the given method, or None if we cant find the class.
@@ -183,12 +222,20 @@ class Declarations:
183
222
  def get_class_decl(self, name: str) -> ClassDecl:
184
223
  return self._classes[name]
185
224
 
225
+ def get_paramaterized_class(self, name: str) -> TypeRefWithVars:
226
+ """
227
+ Returns a class reference with type parameters, if the class is paramaterized.
228
+ """
229
+ type_vars = self._classes[name].type_vars
230
+ return TypeRefWithVars(name, tuple(map(ClassTypeVarRef, type_vars)))
231
+
186
232
 
187
233
  @dataclass
188
234
  class ClassDecl:
189
235
  egg_name: str | None = None
190
236
  type_vars: tuple[str, ...] = ()
191
237
  builtin: bool = False
238
+ init: FunctionDecl | None = None
192
239
  class_methods: dict[str, FunctionDecl] = field(default_factory=dict)
193
240
  # These have to be seperate from class_methods so that printing them can be done easily
194
241
  class_variables: dict[str, ConstantDecl] = field(default_factory=dict)
@@ -293,6 +340,11 @@ class ClassMethodRef:
293
340
  method_name: str
294
341
 
295
342
 
343
+ @dataclass(frozen=True)
344
+ class InitRef:
345
+ class_name: str
346
+
347
+
296
348
  @dataclass(frozen=True)
297
349
  class ClassVariableRef:
298
350
  class_name: str
@@ -305,7 +357,9 @@ class PropertyRef:
305
357
  property_name: str
306
358
 
307
359
 
308
- CallableRef: TypeAlias = FunctionRef | ConstantRef | MethodRef | ClassMethodRef | ClassVariableRef | PropertyRef
360
+ CallableRef: TypeAlias = (
361
+ FunctionRef | ConstantRef | MethodRef | ClassMethodRef | InitRef | ClassVariableRef | PropertyRef
362
+ )
309
363
 
310
364
 
311
365
  ##
@@ -378,7 +432,6 @@ class FunctionSignature:
378
432
  @dataclass(frozen=True)
379
433
  class FunctionDecl:
380
434
  signature: FunctionSignature | SpecialFunctions = field(default_factory=FunctionSignature)
381
-
382
435
  # Egg params
383
436
  builtin: bool = False
384
437
  egg_name: str | None = None
@@ -458,7 +511,7 @@ class CallDecl:
458
511
  bound_tp_params: tuple[JustTypeRef, ...] | None = None
459
512
 
460
513
  def __post_init__(self) -> None:
461
- if self.bound_tp_params and not isinstance(self.callable, ClassMethodRef):
514
+ if self.bound_tp_params and not isinstance(self.callable, ClassMethodRef | InitRef):
462
515
  msg = "Cannot bind type parameters to a non-class method callable."
463
516
  raise ValueError(msg)
464
517
 
@@ -629,7 +682,13 @@ class RuleDecl:
629
682
  name: str | None
630
683
 
631
684
 
632
- RewriteOrRuleDecl: TypeAlias = RewriteDecl | BiRewriteDecl | RuleDecl
685
+ @dataclass(frozen=True)
686
+ class DefaultRewriteDecl:
687
+ ref: CallableRef
688
+ expr: ExprDecl
689
+
690
+
691
+ RewriteOrRuleDecl: TypeAlias = RewriteDecl | BiRewriteDecl | RuleDecl | DefaultRewriteDecl
633
692
 
634
693
 
635
694
  @dataclass(frozen=True)