egglog 11.2.0__cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

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 (46) hide show
  1. egglog/__init__.py +13 -0
  2. egglog/bindings.cpython-314-x86_64-linux-gnu.so +0 -0
  3. egglog/bindings.pyi +734 -0
  4. egglog/builtins.py +1133 -0
  5. egglog/config.py +8 -0
  6. egglog/conversion.py +286 -0
  7. egglog/declarations.py +912 -0
  8. egglog/deconstruct.py +173 -0
  9. egglog/egraph.py +1875 -0
  10. egglog/egraph_state.py +680 -0
  11. egglog/examples/README.rst +5 -0
  12. egglog/examples/__init__.py +3 -0
  13. egglog/examples/bignum.py +32 -0
  14. egglog/examples/bool.py +38 -0
  15. egglog/examples/eqsat_basic.py +44 -0
  16. egglog/examples/fib.py +28 -0
  17. egglog/examples/higher_order_functions.py +42 -0
  18. egglog/examples/jointree.py +67 -0
  19. egglog/examples/lambda_.py +287 -0
  20. egglog/examples/matrix.py +175 -0
  21. egglog/examples/multiset.py +60 -0
  22. egglog/examples/ndarrays.py +144 -0
  23. egglog/examples/resolution.py +84 -0
  24. egglog/examples/schedule_demo.py +34 -0
  25. egglog/exp/__init__.py +3 -0
  26. egglog/exp/array_api.py +2019 -0
  27. egglog/exp/array_api_jit.py +51 -0
  28. egglog/exp/array_api_loopnest.py +74 -0
  29. egglog/exp/array_api_numba.py +69 -0
  30. egglog/exp/array_api_program_gen.py +510 -0
  31. egglog/exp/program_gen.py +425 -0
  32. egglog/exp/siu_examples.py +32 -0
  33. egglog/ipython_magic.py +41 -0
  34. egglog/pretty.py +509 -0
  35. egglog/py.typed +0 -0
  36. egglog/runtime.py +712 -0
  37. egglog/thunk.py +97 -0
  38. egglog/type_constraint_solver.py +113 -0
  39. egglog/version_compat.py +87 -0
  40. egglog/visualizer.css +1 -0
  41. egglog/visualizer.js +35777 -0
  42. egglog/visualizer_widget.py +39 -0
  43. egglog-11.2.0.dist-info/METADATA +74 -0
  44. egglog-11.2.0.dist-info/RECORD +46 -0
  45. egglog-11.2.0.dist-info/WHEEL +4 -0
  46. egglog-11.2.0.dist-info/licenses/LICENSE +21 -0
egglog/builtins.py ADDED
@@ -0,0 +1,1133 @@
1
+ # mypy: disable-error-code="empty-body"
2
+ """
3
+ Builtin sorts and function to egg.
4
+ """
5
+
6
+ from __future__ import annotations
7
+
8
+ from collections.abc import Callable
9
+ from dataclasses import dataclass
10
+ from fractions import Fraction
11
+ from functools import partial, reduce
12
+ from inspect import signature
13
+ from types import FunctionType, MethodType
14
+ from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, cast, overload
15
+
16
+ from typing_extensions import TypeVarTuple, Unpack, deprecated
17
+
18
+ from .conversion import convert, converter, get_type_args, resolve_literal
19
+ from .declarations import *
20
+ from .deconstruct import get_callable_args, get_literal_value
21
+ from .egraph import BaseExpr, BuiltinExpr, _add_default_rewrite_inner, expr_fact, function, get_current_ruleset, method
22
+ from .runtime import RuntimeExpr, RuntimeFunction, resolve_type_annotation_mutate
23
+ from .thunk import Thunk
24
+
25
+ if TYPE_CHECKING:
26
+ from collections.abc import Iterator
27
+
28
+
29
+ __all__ = [
30
+ "BigInt",
31
+ "BigIntLike",
32
+ "BigRat",
33
+ "BigRatLike",
34
+ "Bool",
35
+ "BoolLike",
36
+ "ExprValueError",
37
+ "Map",
38
+ "MapLike",
39
+ "MultiSet",
40
+ "PyObject",
41
+ "Rational",
42
+ "Set",
43
+ "SetLike",
44
+ "String",
45
+ "StringLike",
46
+ "Unit",
47
+ "UnstableFn",
48
+ "Vec",
49
+ "VecLike",
50
+ "f64",
51
+ "f64Like",
52
+ "i64",
53
+ "i64Like",
54
+ "join",
55
+ "py_eval",
56
+ "py_eval_fn",
57
+ "py_exec",
58
+ ]
59
+
60
+
61
+ @dataclass
62
+ class ExprValueError(AttributeError):
63
+ """
64
+ Raised when an expression cannot be converted to a Python value because the value is not a constructor.
65
+ """
66
+
67
+ expr: BaseExpr
68
+ allowed: str
69
+
70
+ def __str__(self) -> str:
71
+ return f"Cannot get Python value of {self.expr}, must be of form {self.allowed}. Try calling `extract` on it to get the underlying value."
72
+
73
+
74
+ class Unit(BuiltinExpr, egg_sort="Unit"):
75
+ """
76
+ The unit type. This is used to reprsent if a value exists in the e-graph or not.
77
+ """
78
+
79
+ def __init__(self) -> None: ...
80
+
81
+ @method(preserve=True)
82
+ def __bool__(self) -> bool:
83
+ return bool(expr_fact(self))
84
+
85
+
86
+ class String(BuiltinExpr):
87
+ def __init__(self, value: str) -> None: ...
88
+
89
+ @method(preserve=True)
90
+ @deprecated("use .value")
91
+ def eval(self) -> str:
92
+ return self.value
93
+
94
+ @method(preserve=True) # type: ignore[prop-decorator]
95
+ @property
96
+ def value(self) -> str:
97
+ if (value := get_literal_value(self)) is not None:
98
+ return value
99
+ raise ExprValueError(self, "String")
100
+
101
+ __match_args__ = ("value",)
102
+
103
+ @method(egg_fn="replace")
104
+ def replace(self, old: StringLike, new: StringLike) -> String: ...
105
+
106
+
107
+ StringLike: TypeAlias = String | str
108
+
109
+
110
+ @function(egg_fn="+", builtin=True)
111
+ def join(*strings: StringLike) -> String: ...
112
+
113
+
114
+ converter(str, String, String)
115
+
116
+
117
+ class Bool(BuiltinExpr, egg_sort="bool"):
118
+ def __init__(self, value: bool) -> None: ...
119
+
120
+ @method(preserve=True)
121
+ @deprecated("use .value")
122
+ def eval(self) -> bool:
123
+ return self.value
124
+
125
+ @method(preserve=True) # type: ignore[prop-decorator]
126
+ @property
127
+ def value(self) -> bool:
128
+ if (value := get_literal_value(self)) is not None:
129
+ return value
130
+ raise ExprValueError(self, "Bool")
131
+
132
+ __match_args__ = ("value",)
133
+
134
+ @method(preserve=True)
135
+ def __bool__(self) -> bool:
136
+ return self.value
137
+
138
+ @method(egg_fn="not")
139
+ def __invert__(self) -> Bool: ...
140
+
141
+ @method(egg_fn="and")
142
+ def __and__(self, other: BoolLike) -> Bool: ...
143
+
144
+ @method(egg_fn="or")
145
+ def __or__(self, other: BoolLike) -> Bool: ...
146
+
147
+ @method(egg_fn="xor")
148
+ def __xor__(self, other: BoolLike) -> Bool: ...
149
+
150
+ @method(egg_fn="=>")
151
+ def implies(self, other: BoolLike) -> Bool: ...
152
+
153
+
154
+ BoolLike: TypeAlias = Bool | bool
155
+
156
+
157
+ converter(bool, Bool, Bool)
158
+
159
+
160
+ class i64(BuiltinExpr): # noqa: N801
161
+ def __init__(self, value: int) -> None: ...
162
+
163
+ @method(preserve=True)
164
+ @deprecated("use .value")
165
+ def eval(self) -> int:
166
+ return self.value
167
+
168
+ @method(preserve=True) # type: ignore[prop-decorator]
169
+ @property
170
+ def value(self) -> int:
171
+ if (value := get_literal_value(self)) is not None:
172
+ return value
173
+ raise ExprValueError(self, "i64")
174
+
175
+ __match_args__ = ("value",)
176
+
177
+ @method(preserve=True)
178
+ def __index__(self) -> int:
179
+ return self.value
180
+
181
+ @method(preserve=True)
182
+ def __int__(self) -> int:
183
+ return self.value
184
+
185
+ @method(egg_fn="+")
186
+ def __add__(self, other: i64Like) -> i64: ...
187
+
188
+ @method(egg_fn="-")
189
+ def __sub__(self, other: i64Like) -> i64: ...
190
+
191
+ @method(egg_fn="*")
192
+ def __mul__(self, other: i64Like) -> i64: ...
193
+
194
+ @method(egg_fn="/")
195
+ def __truediv__(self, other: i64Like) -> i64: ...
196
+
197
+ @method(egg_fn="%")
198
+ def __mod__(self, other: i64Like) -> i64: ...
199
+
200
+ @method(egg_fn="&")
201
+ def __and__(self, other: i64Like) -> i64: ...
202
+
203
+ @method(egg_fn="|")
204
+ def __or__(self, other: i64Like) -> i64: ...
205
+
206
+ @method(egg_fn="^")
207
+ def __xor__(self, other: i64Like) -> i64: ...
208
+
209
+ @method(egg_fn="<<")
210
+ def __lshift__(self, other: i64Like) -> i64: ...
211
+
212
+ @method(egg_fn=">>")
213
+ def __rshift__(self, other: i64Like) -> i64: ...
214
+
215
+ def __radd__(self, other: i64Like) -> i64: ...
216
+
217
+ def __rsub__(self, other: i64Like) -> i64: ...
218
+
219
+ def __rmul__(self, other: i64Like) -> i64: ...
220
+
221
+ def __rtruediv__(self, other: i64Like) -> i64: ...
222
+
223
+ def __rmod__(self, other: i64Like) -> i64: ...
224
+
225
+ def __rand__(self, other: i64Like) -> i64: ...
226
+
227
+ def __ror__(self, other: i64Like) -> i64: ...
228
+
229
+ def __rxor__(self, other: i64Like) -> i64: ...
230
+
231
+ def __rlshift__(self, other: i64Like) -> i64: ...
232
+
233
+ def __rrshift__(self, other: i64Like) -> i64: ...
234
+
235
+ @method(egg_fn="log2")
236
+ def log2(self) -> i64: ...
237
+
238
+ @method(egg_fn="not-i64")
239
+ def __invert__(self) -> i64: ...
240
+
241
+ @method(egg_fn="<")
242
+ def __lt__(self, other: i64Like) -> Unit: # type: ignore[has-type]
243
+ ...
244
+
245
+ @method(egg_fn=">")
246
+ def __gt__(self, other: i64Like) -> Unit: ...
247
+
248
+ @method(egg_fn="<=")
249
+ def __le__(self, other: i64Like) -> Unit: # type: ignore[has-type]
250
+ ...
251
+
252
+ @method(egg_fn=">=")
253
+ def __ge__(self, other: i64Like) -> Unit: ...
254
+
255
+ @method(egg_fn="min")
256
+ def min(self, other: i64Like) -> i64: ...
257
+
258
+ @method(egg_fn="max")
259
+ def max(self, other: i64Like) -> i64: ...
260
+
261
+ @method(egg_fn="to-string")
262
+ def to_string(self) -> String: ...
263
+
264
+ @method(egg_fn="bool-<")
265
+ def bool_lt(self, other: i64Like) -> Bool: ...
266
+
267
+ @method(egg_fn="bool->")
268
+ def bool_gt(self, other: i64Like) -> Bool: ...
269
+
270
+ @method(egg_fn="bool-<=")
271
+ def bool_le(self, other: i64Like) -> Bool: ...
272
+
273
+ @method(egg_fn="bool->=")
274
+ def bool_ge(self, other: i64Like) -> Bool: ...
275
+
276
+
277
+ # The types which can be convertered into an i64
278
+ i64Like: TypeAlias = i64 | int # noqa: N816, PYI042
279
+
280
+ converter(int, i64, i64)
281
+
282
+
283
+ @function(builtin=True, egg_fn="count-matches")
284
+ def count_matches(s: StringLike, pattern: StringLike) -> i64: ...
285
+
286
+
287
+ class f64(BuiltinExpr): # noqa: N801
288
+ def __init__(self, value: float) -> None: ...
289
+
290
+ @method(preserve=True)
291
+ @deprecated("use .value")
292
+ def eval(self) -> float:
293
+ return self.value
294
+
295
+ @method(preserve=True) # type: ignore[prop-decorator]
296
+ @property
297
+ def value(self) -> float:
298
+ if (value := get_literal_value(self)) is not None:
299
+ return value
300
+ raise ExprValueError(self, "f64")
301
+
302
+ __match_args__ = ("value",)
303
+
304
+ @method(preserve=True)
305
+ def __float__(self) -> float:
306
+ return self.value
307
+
308
+ @method(preserve=True)
309
+ def __int__(self) -> int:
310
+ return int(self.value)
311
+
312
+ @method(egg_fn="neg")
313
+ def __neg__(self) -> f64: ...
314
+
315
+ @method(egg_fn="+")
316
+ def __add__(self, other: f64Like) -> f64: ...
317
+
318
+ @method(egg_fn="-")
319
+ def __sub__(self, other: f64Like) -> f64: ...
320
+
321
+ @method(egg_fn="*")
322
+ def __mul__(self, other: f64Like) -> f64: ...
323
+
324
+ @method(egg_fn="/")
325
+ def __truediv__(self, other: f64Like) -> f64: ...
326
+
327
+ @method(egg_fn="%")
328
+ def __mod__(self, other: f64Like) -> f64: ...
329
+
330
+ @method(egg_fn="^")
331
+ def __pow__(self, other: f64Like) -> f64: ...
332
+
333
+ def __radd__(self, other: f64Like) -> f64: ...
334
+
335
+ def __rsub__(self, other: f64Like) -> f64: ...
336
+
337
+ def __rmul__(self, other: f64Like) -> f64: ...
338
+
339
+ def __rtruediv__(self, other: f64Like) -> f64: ...
340
+
341
+ def __rmod__(self, other: f64Like) -> f64: ...
342
+
343
+ @method(egg_fn="<")
344
+ def __lt__(self, other: f64Like) -> Unit: # type: ignore[has-type]
345
+ ...
346
+
347
+ @method(egg_fn=">")
348
+ def __gt__(self, other: f64Like) -> Unit: ...
349
+
350
+ @method(egg_fn="<=")
351
+ def __le__(self, other: f64Like) -> Unit: # type: ignore[has-type]
352
+ ...
353
+
354
+ @method(egg_fn=">=")
355
+ def __ge__(self, other: f64Like) -> Unit: ...
356
+
357
+ @method(egg_fn="min")
358
+ def min(self, other: f64Like) -> f64: ...
359
+
360
+ @method(egg_fn="max")
361
+ def max(self, other: f64Like) -> f64: ...
362
+
363
+ @method(egg_fn="to-i64")
364
+ def to_i64(self) -> i64: ...
365
+
366
+ @method(egg_fn="to-f64")
367
+ @classmethod
368
+ def from_i64(cls, i: i64) -> f64: ...
369
+
370
+ @method(egg_fn="to-string")
371
+ def to_string(self) -> String: ...
372
+
373
+
374
+ f64Like: TypeAlias = f64 | float # noqa: N816, PYI042
375
+
376
+
377
+ converter(float, f64, f64)
378
+
379
+
380
+ T = TypeVar("T", bound=BaseExpr)
381
+ V = TypeVar("V", bound=BaseExpr)
382
+
383
+
384
+ class Map(BuiltinExpr, Generic[T, V]):
385
+ @method(preserve=True)
386
+ @deprecated("use .value")
387
+ def eval(self) -> dict[T, V]:
388
+ return self.value
389
+
390
+ @method(preserve=True) # type: ignore[prop-decorator]
391
+ @property
392
+ def value(self) -> dict[T, V]:
393
+ d = {}
394
+ while args := get_callable_args(self, Map[T, V].insert):
395
+ self, k, v = args # noqa: PLW0642
396
+ d[k] = v
397
+ if get_callable_args(self, Map.empty) is None:
398
+ raise ExprValueError(self, "Map.empty or Map.insert")
399
+ return d
400
+
401
+ __match_args__ = ("value",)
402
+
403
+ @method(preserve=True)
404
+ def __iter__(self) -> Iterator[T]:
405
+ return iter(self.value)
406
+
407
+ @method(preserve=True)
408
+ def __len__(self) -> int:
409
+ return len(self.value)
410
+
411
+ @method(preserve=True)
412
+ def __contains__(self, key: T) -> bool:
413
+ return key in self.value
414
+
415
+ @method(egg_fn="map-empty")
416
+ @classmethod
417
+ def empty(cls) -> Map[T, V]: ...
418
+
419
+ @method(egg_fn="map-insert")
420
+ def insert(self, key: T, value: V) -> Map[T, V]: ...
421
+
422
+ @method(egg_fn="map-get")
423
+ def __getitem__(self, key: T) -> V: ...
424
+
425
+ @method(egg_fn="map-not-contains")
426
+ def not_contains(self, key: T) -> Unit: ...
427
+
428
+ @method(egg_fn="map-contains")
429
+ def contains(self, key: T) -> Unit: ...
430
+
431
+ @method(egg_fn="map-remove")
432
+ def remove(self, key: T) -> Map[T, V]: ...
433
+
434
+ @method(egg_fn="rebuild")
435
+ def rebuild(self) -> Map[T, V]: ...
436
+
437
+
438
+ TO = TypeVar("TO")
439
+ VO = TypeVar("VO")
440
+
441
+ converter(
442
+ dict,
443
+ Map,
444
+ lambda t: reduce(
445
+ (lambda acc, kv: acc.insert(convert(kv[0], get_type_args()[0]), convert(kv[1], get_type_args()[1]))),
446
+ t.items(),
447
+ Map[get_type_args()].empty(), # type: ignore[misc]
448
+ ),
449
+ )
450
+
451
+ MapLike: TypeAlias = Map[T, V] | dict[TO, VO]
452
+
453
+
454
+ class Set(BuiltinExpr, Generic[T]):
455
+ @method(preserve=True)
456
+ @deprecated("use .value")
457
+ def eval(self) -> set[T]:
458
+ return self.value
459
+
460
+ @method(preserve=True) # type: ignore[prop-decorator]
461
+ @property
462
+ def value(self) -> set[T]:
463
+ if (args := get_callable_args(self, Set[T])) is not None:
464
+ return set(args)
465
+ raise ExprValueError(self, "Set(*xs)")
466
+
467
+ __match_args__ = ("value",)
468
+
469
+ @method(preserve=True)
470
+ def __iter__(self) -> Iterator[T]:
471
+ return iter(self.value)
472
+
473
+ @method(preserve=True)
474
+ def __len__(self) -> int:
475
+ return len(self.value)
476
+
477
+ @method(preserve=True)
478
+ def __contains__(self, key: T) -> bool:
479
+ return key in self.value
480
+
481
+ @method(egg_fn="set-of")
482
+ def __init__(self, *args: T) -> None: ...
483
+
484
+ @method(egg_fn="set-empty")
485
+ @classmethod
486
+ def empty(cls) -> Set[T]: ...
487
+
488
+ @method(egg_fn="set-insert")
489
+ def insert(self, value: T) -> Set[T]: ...
490
+
491
+ @method(egg_fn="set-not-contains")
492
+ def not_contains(self, value: T) -> Unit: ...
493
+
494
+ @method(egg_fn="set-contains")
495
+ def contains(self, value: T) -> Unit: ...
496
+
497
+ @method(egg_fn="set-remove")
498
+ def remove(self, value: T) -> Set[T]: ...
499
+
500
+ @method(egg_fn="set-union")
501
+ def __or__(self, other: Set[T]) -> Set[T]: ...
502
+
503
+ @method(egg_fn="set-diff")
504
+ def __sub__(self, other: Set[T]) -> Set[T]: ...
505
+
506
+ @method(egg_fn="set-intersect")
507
+ def __and__(self, other: Set[T]) -> Set[T]: ...
508
+
509
+ @method(egg_fn="rebuild")
510
+ def rebuild(self) -> Set[T]: ...
511
+
512
+
513
+ converter(
514
+ set,
515
+ Set,
516
+ lambda t: Set[get_type_args()[0]]( # type: ignore[misc,operator]
517
+ *(convert(x, get_type_args()[0]) for x in t)
518
+ ),
519
+ )
520
+
521
+ SetLike: TypeAlias = Set[T] | set[TO]
522
+
523
+
524
+ class MultiSet(BuiltinExpr, Generic[T]):
525
+ @method(preserve=True)
526
+ @deprecated("use .value")
527
+ def eval(self) -> list[T]:
528
+ return self.value
529
+
530
+ @method(preserve=True) # type: ignore[prop-decorator]
531
+ @property
532
+ def value(self) -> list[T]:
533
+ if (args := get_callable_args(self, MultiSet[T])) is not None:
534
+ return list(args)
535
+ raise ExprValueError(self, "MultiSet")
536
+
537
+ __match_args__ = ("value",)
538
+
539
+ @method(preserve=True)
540
+ def __iter__(self) -> Iterator[T]:
541
+ return iter(self.value)
542
+
543
+ @method(preserve=True)
544
+ def __len__(self) -> int:
545
+ return len(self.value)
546
+
547
+ @method(preserve=True)
548
+ def __contains__(self, key: T) -> bool:
549
+ return key in self.value
550
+
551
+ @method(egg_fn="multiset-of")
552
+ def __init__(self, *args: T) -> None: ...
553
+
554
+ @method(egg_fn="multiset-insert")
555
+ def insert(self, value: T) -> MultiSet[T]: ...
556
+
557
+ @method(egg_fn="multiset-not-contains")
558
+ def not_contains(self, value: T) -> Unit: ...
559
+
560
+ @method(egg_fn="multiset-contains")
561
+ def contains(self, value: T) -> Unit: ...
562
+
563
+ @method(egg_fn="multiset-remove")
564
+ def remove(self, value: T) -> MultiSet[T]: ...
565
+
566
+ @method(egg_fn="multiset-length")
567
+ def length(self) -> i64: ...
568
+
569
+ @method(egg_fn="multiset-pick")
570
+ def pick(self) -> T: ...
571
+
572
+ @method(egg_fn="multiset-sum")
573
+ def __add__(self, other: MultiSet[T]) -> MultiSet[T]: ...
574
+
575
+ @method(egg_fn="unstable-multiset-map", reverse_args=True)
576
+ def map(self, f: Callable[[T], T]) -> MultiSet[T]: ...
577
+
578
+
579
+ class Rational(BuiltinExpr):
580
+ @method(preserve=True)
581
+ @deprecated("use .value")
582
+ def eval(self) -> Fraction:
583
+ return self.value
584
+
585
+ @method(preserve=True) # type: ignore[prop-decorator]
586
+ @property
587
+ def value(self) -> Fraction:
588
+ match get_callable_args(self, Rational):
589
+ case (i64(num), i64(den)):
590
+ return Fraction(num, den)
591
+ raise ExprValueError(self, "Rational(i64(num), i64(den))")
592
+
593
+ __match_args__ = ("value",)
594
+
595
+ @method(preserve=True)
596
+ def __float__(self) -> float:
597
+ return float(self.value)
598
+
599
+ @method(preserve=True)
600
+ def __int__(self) -> int:
601
+ return int(self.value)
602
+
603
+ @method(egg_fn="rational")
604
+ def __init__(self, num: i64Like, den: i64Like) -> None: ...
605
+
606
+ @method(egg_fn="to-f64")
607
+ def to_f64(self) -> f64: ...
608
+
609
+ @method(egg_fn="+")
610
+ def __add__(self, other: Rational) -> Rational: ...
611
+
612
+ @method(egg_fn="-")
613
+ def __sub__(self, other: Rational) -> Rational: ...
614
+
615
+ @method(egg_fn="*")
616
+ def __mul__(self, other: Rational) -> Rational: ...
617
+
618
+ @method(egg_fn="/")
619
+ def __truediv__(self, other: Rational) -> Rational: ...
620
+
621
+ @method(egg_fn="min")
622
+ def min(self, other: Rational) -> Rational: ...
623
+
624
+ @method(egg_fn="max")
625
+ def max(self, other: Rational) -> Rational: ...
626
+
627
+ @method(egg_fn="neg")
628
+ def __neg__(self) -> Rational: ...
629
+
630
+ @method(egg_fn="abs")
631
+ def __abs__(self) -> Rational: ...
632
+
633
+ @method(egg_fn="floor")
634
+ def floor(self) -> Rational: ...
635
+
636
+ @method(egg_fn="ceil")
637
+ def ceil(self) -> Rational: ...
638
+
639
+ @method(egg_fn="round")
640
+ def round(self) -> Rational: ...
641
+
642
+ @method(egg_fn="pow")
643
+ def __pow__(self, other: Rational) -> Rational: ...
644
+
645
+ @method(egg_fn="log")
646
+ def log(self) -> Rational: ...
647
+
648
+ @method(egg_fn="sqrt")
649
+ def sqrt(self) -> Rational: ...
650
+
651
+ @method(egg_fn="cbrt")
652
+ def cbrt(self) -> Rational: ...
653
+
654
+ @method(egg_fn="numer") # type: ignore[prop-decorator]
655
+ @property
656
+ def numer(self) -> i64: ...
657
+
658
+ @method(egg_fn="denom") # type: ignore[prop-decorator]
659
+ @property
660
+ def denom(self) -> i64: ...
661
+
662
+
663
+ class BigInt(BuiltinExpr):
664
+ @method(preserve=True)
665
+ @deprecated("use .value")
666
+ def eval(self) -> int:
667
+ return self.value
668
+
669
+ @method(preserve=True) # type: ignore[prop-decorator]
670
+ @property
671
+ def value(self) -> int:
672
+ match get_callable_args(self, BigInt.from_string):
673
+ case (String(s),):
674
+ return int(s)
675
+ raise ExprValueError(self, "BigInt.from_string(String(s))")
676
+
677
+ __match_args__ = ("value",)
678
+
679
+ @method(preserve=True)
680
+ def __index__(self) -> int:
681
+ return self.value
682
+
683
+ @method(preserve=True)
684
+ def __int__(self) -> int:
685
+ return self.value
686
+
687
+ @method(egg_fn="from-string")
688
+ @classmethod
689
+ def from_string(cls, s: StringLike) -> BigInt: ...
690
+
691
+ @method(egg_fn="bigint")
692
+ def __init__(self, value: i64Like) -> None: ...
693
+
694
+ @method(egg_fn="+")
695
+ def __add__(self, other: BigIntLike) -> BigInt: ...
696
+
697
+ @method(egg_fn="-")
698
+ def __sub__(self, other: BigIntLike) -> BigInt: ...
699
+
700
+ @method(egg_fn="*")
701
+ def __mul__(self, other: BigIntLike) -> BigInt: ...
702
+
703
+ @method(egg_fn="/")
704
+ def __truediv__(self, other: BigIntLike) -> BigInt: ...
705
+
706
+ @method(egg_fn="%")
707
+ def __mod__(self, other: BigIntLike) -> BigInt: ...
708
+
709
+ @method(egg_fn="&")
710
+ def __and__(self, other: BigIntLike) -> BigInt: ...
711
+
712
+ @method(egg_fn="|")
713
+ def __or__(self, other: BigIntLike) -> BigInt: ...
714
+
715
+ @method(egg_fn="^")
716
+ def __xor__(self, other: BigIntLike) -> BigInt: ...
717
+
718
+ @method(egg_fn="<<")
719
+ def __lshift__(self, other: i64Like) -> BigInt: ...
720
+
721
+ @method(egg_fn=">>")
722
+ def __rshift__(self, other: i64Like) -> BigInt: ...
723
+
724
+ def __radd__(self, other: BigIntLike) -> BigInt: ...
725
+
726
+ def __rsub__(self, other: BigIntLike) -> BigInt: ...
727
+
728
+ def __rmul__(self, other: BigIntLike) -> BigInt: ...
729
+
730
+ def __rtruediv__(self, other: BigIntLike) -> BigInt: ...
731
+
732
+ def __rmod__(self, other: BigIntLike) -> BigInt: ...
733
+
734
+ def __rand__(self, other: BigIntLike) -> BigInt: ...
735
+
736
+ def __ror__(self, other: BigIntLike) -> BigInt: ...
737
+
738
+ def __rxor__(self, other: BigIntLike) -> BigInt: ...
739
+
740
+ @method(egg_fn="not-Z")
741
+ def __invert__(self) -> BigInt: ...
742
+
743
+ @method(egg_fn="bits")
744
+ def bits(self) -> BigInt: ...
745
+
746
+ @method(egg_fn="<")
747
+ def __lt__(self, other: BigIntLike) -> Unit: # type: ignore[has-type]
748
+ ...
749
+
750
+ @method(egg_fn=">")
751
+ def __gt__(self, other: BigIntLike) -> Unit: ...
752
+
753
+ @method(egg_fn="<=")
754
+ def __le__(self, other: BigIntLike) -> Unit: # type: ignore[has-type]
755
+ ...
756
+
757
+ @method(egg_fn=">=")
758
+ def __ge__(self, other: BigIntLike) -> Unit: ...
759
+
760
+ @method(egg_fn="min")
761
+ def min(self, other: BigIntLike) -> BigInt: ...
762
+
763
+ @method(egg_fn="max")
764
+ def max(self, other: BigIntLike) -> BigInt: ...
765
+
766
+ @method(egg_fn="to-string")
767
+ def to_string(self) -> String: ...
768
+
769
+ @method(egg_fn="bool-=")
770
+ def bool_eq(self, other: BigIntLike) -> Bool: ...
771
+
772
+ @method(egg_fn="bool-<")
773
+ def bool_lt(self, other: BigIntLike) -> Bool: ...
774
+
775
+ @method(egg_fn="bool->")
776
+ def bool_gt(self, other: BigIntLike) -> Bool: ...
777
+
778
+ @method(egg_fn="bool-<=")
779
+ def bool_le(self, other: BigIntLike) -> Bool: ...
780
+
781
+ @method(egg_fn="bool->=")
782
+ def bool_ge(self, other: BigIntLike) -> Bool: ...
783
+
784
+
785
+ converter(i64, BigInt, lambda i: BigInt(i))
786
+
787
+ BigIntLike: TypeAlias = BigInt | i64Like
788
+
789
+
790
+ class BigRat(BuiltinExpr):
791
+ @method(preserve=True)
792
+ @deprecated("use .value")
793
+ def eval(self) -> Fraction:
794
+ return self.value
795
+
796
+ @method(preserve=True) # type: ignore[prop-decorator]
797
+ @property
798
+ def value(self) -> Fraction:
799
+ match get_callable_args(self, BigRat):
800
+ case (BigInt(num), BigInt(den)):
801
+ return Fraction(num, den)
802
+ raise ExprValueError(self, "BigRat(BigInt(num), BigInt(den))")
803
+
804
+ __match_args__ = ("value",)
805
+
806
+ @method(preserve=True)
807
+ def __float__(self) -> float:
808
+ return float(self.value)
809
+
810
+ @method(preserve=True)
811
+ def __int__(self) -> int:
812
+ return int(self.value)
813
+
814
+ @method(egg_fn="bigrat")
815
+ def __init__(self, num: BigIntLike, den: BigIntLike) -> None: ...
816
+
817
+ @method(egg_fn="to-f64")
818
+ def to_f64(self) -> f64: ...
819
+
820
+ @method(egg_fn="+")
821
+ def __add__(self, other: BigRatLike) -> BigRat: ...
822
+
823
+ @method(egg_fn="-")
824
+ def __sub__(self, other: BigRatLike) -> BigRat: ...
825
+
826
+ @method(egg_fn="*")
827
+ def __mul__(self, other: BigRatLike) -> BigRat: ...
828
+
829
+ @method(egg_fn="/")
830
+ def __truediv__(self, other: BigRatLike) -> BigRat: ...
831
+
832
+ @method(egg_fn="min")
833
+ def min(self, other: BigRatLike) -> BigRat: ...
834
+
835
+ @method(egg_fn="max")
836
+ def max(self, other: BigRatLike) -> BigRat: ...
837
+
838
+ @method(egg_fn="neg")
839
+ def __neg__(self) -> BigRat: ...
840
+
841
+ @method(egg_fn="abs")
842
+ def __abs__(self) -> BigRat: ...
843
+
844
+ @method(egg_fn="floor")
845
+ def floor(self) -> BigRat: ...
846
+
847
+ @method(egg_fn="ceil")
848
+ def ceil(self) -> BigRat: ...
849
+
850
+ @method(egg_fn="round")
851
+ def round(self) -> BigRat: ...
852
+
853
+ @method(egg_fn="pow")
854
+ def __pow__(self, other: BigRatLike) -> BigRat: ...
855
+
856
+ @method(egg_fn="log")
857
+ def log(self) -> BigRat: ...
858
+
859
+ @method(egg_fn="sqrt")
860
+ def sqrt(self) -> BigRat: ...
861
+
862
+ @method(egg_fn="cbrt")
863
+ def cbrt(self) -> BigRat: ...
864
+
865
+ @method(egg_fn="numer") # type: ignore[prop-decorator]
866
+ @property
867
+ def numer(self) -> BigInt: ...
868
+
869
+ @method(egg_fn="denom") # type: ignore[prop-decorator]
870
+ @property
871
+ def denom(self) -> BigInt: ...
872
+
873
+ @method(egg_fn="<")
874
+ def __lt__(self, other: BigRatLike) -> Unit: ... # type: ignore[has-type]
875
+
876
+ @method(egg_fn=">")
877
+ def __gt__(self, other: BigRatLike) -> Unit: ...
878
+
879
+ @method(egg_fn=">=")
880
+ def __ge__(self, other: BigRatLike) -> Unit: ... # type: ignore[has-type]
881
+
882
+ @method(egg_fn="<=")
883
+ def __le__(self, other: BigRatLike) -> Unit: ...
884
+
885
+
886
+ converter(Fraction, BigRat, lambda f: BigRat(f.numerator, f.denominator))
887
+ BigRatLike: TypeAlias = BigRat | Fraction
888
+
889
+
890
+ class Vec(BuiltinExpr, Generic[T]):
891
+ @method(preserve=True)
892
+ @deprecated("use .value")
893
+ def eval(self) -> tuple[T, ...]:
894
+ return self.value
895
+
896
+ @method(preserve=True) # type: ignore[prop-decorator]
897
+ @property
898
+ def value(self) -> tuple[T, ...]:
899
+ if get_callable_args(self, Vec.empty) is not None:
900
+ return ()
901
+ if (args := get_callable_args(self, Vec[T])) is not None:
902
+ return args
903
+ raise ExprValueError(self, "Vec(*xs) or Vec.empty()")
904
+
905
+ __match_args__ = ("value",)
906
+
907
+ @method(preserve=True)
908
+ def __iter__(self) -> Iterator[T]:
909
+ return iter(self.value)
910
+
911
+ @method(preserve=True)
912
+ def __len__(self) -> int:
913
+ return len(self.value)
914
+
915
+ @method(preserve=True)
916
+ def __contains__(self, key: T) -> bool:
917
+ return key in self.value
918
+
919
+ @method(egg_fn="vec-of")
920
+ def __init__(self, *args: T) -> None: ...
921
+
922
+ @method(egg_fn="vec-empty")
923
+ @classmethod
924
+ def empty(cls) -> Vec[T]: ...
925
+
926
+ @method(egg_fn="vec-append")
927
+ def append(self, *others: Vec[T]) -> Vec[T]: ...
928
+
929
+ @method(egg_fn="vec-push")
930
+ def push(self, value: T) -> Vec[T]: ...
931
+
932
+ @method(egg_fn="vec-pop")
933
+ def pop(self) -> Vec[T]: ...
934
+
935
+ @method(egg_fn="vec-not-contains")
936
+ def not_contains(self, value: T) -> Unit: ...
937
+
938
+ @method(egg_fn="vec-contains")
939
+ def contains(self, value: T) -> Unit: ...
940
+
941
+ @method(egg_fn="vec-length")
942
+ def length(self) -> i64: ...
943
+
944
+ @method(egg_fn="vec-get")
945
+ def __getitem__(self, index: i64Like) -> T: ...
946
+
947
+ @method(egg_fn="rebuild")
948
+ def rebuild(self) -> Vec[T]: ...
949
+
950
+ @method(egg_fn="vec-remove")
951
+ def remove(self, index: i64Like) -> Vec[T]: ...
952
+
953
+ @method(egg_fn="vec-set")
954
+ def set(self, index: i64Like, value: T) -> Vec[T]: ...
955
+
956
+
957
+ for sequence_type in (list, tuple):
958
+ converter(
959
+ sequence_type,
960
+ Vec,
961
+ lambda t: Vec[get_type_args()[0]]( # type: ignore[misc,operator]
962
+ *(convert(x, get_type_args()[0]) for x in t)
963
+ ),
964
+ )
965
+
966
+ VecLike: TypeAlias = Vec[T] | tuple[TO, ...] | list[TO]
967
+
968
+
969
+ class PyObject(BuiltinExpr):
970
+ @method(preserve=True)
971
+ @deprecated("use .value")
972
+ def eval(self) -> object:
973
+ return self.value
974
+
975
+ @method(preserve=True) # type: ignore[prop-decorator]
976
+ @property
977
+ def value(self) -> object:
978
+ expr = cast("RuntimeExpr", self).__egg_typed_expr__.expr
979
+ if not isinstance(expr, PyObjectDecl):
980
+ raise ExprValueError(self, "PyObject(x)")
981
+ return expr.value
982
+
983
+ __match_args__ = ("value",)
984
+
985
+ def __init__(self, value: object) -> None: ...
986
+
987
+ @method(egg_fn="py-from-string")
988
+ @classmethod
989
+ def from_string(cls, s: StringLike) -> PyObject: ...
990
+
991
+ @method(egg_fn="py-to-string")
992
+ def to_string(self) -> String: ...
993
+
994
+ @method(egg_fn="py-to-bool")
995
+ def to_bool(self) -> Bool: ...
996
+
997
+ @method(egg_fn="py-dict-update")
998
+ def dict_update(self, *keys_and_values: object) -> PyObject: ...
999
+
1000
+ @method(egg_fn="py-from-int")
1001
+ @classmethod
1002
+ def from_int(cls, i: i64Like) -> PyObject: ...
1003
+
1004
+ @method(egg_fn="py-dict")
1005
+ @classmethod
1006
+ def dict(cls, *keys_and_values: object) -> PyObject: ...
1007
+
1008
+
1009
+ converter(object, PyObject, PyObject)
1010
+
1011
+
1012
+ @function(builtin=True, egg_fn="py-eval")
1013
+ def py_eval(code: StringLike, globals: object = PyObject.dict(), locals: object = PyObject.dict()) -> PyObject: ...
1014
+
1015
+
1016
+ class PyObjectFunction(Protocol):
1017
+ def __call__(self, *__args: PyObject) -> PyObject: ...
1018
+
1019
+
1020
+ def py_eval_fn(fn: Callable) -> PyObjectFunction:
1021
+ """
1022
+ Takes a python callable and maps it to a callable which takes and returns PyObjects.
1023
+
1024
+ It translates it to a call which uses `py_eval` to call the function, passing in the
1025
+ args as locals, and using the globals from function.
1026
+ """
1027
+
1028
+ def inner(*__args: PyObject, __fn: Callable = fn) -> PyObject:
1029
+ new_kvs: list[object] = []
1030
+ eval_str = "__fn("
1031
+ for i, arg in enumerate(__args):
1032
+ new_kvs.extend((f"__arg_{i}", arg))
1033
+ eval_str += f"__arg_{i}, "
1034
+ eval_str += ")"
1035
+ return py_eval(eval_str, PyObject({"__fn": __fn}).dict_update(*new_kvs), __fn.__globals__)
1036
+
1037
+ return inner
1038
+
1039
+
1040
+ @function(builtin=True, egg_fn="py-exec")
1041
+ def py_exec(code: StringLike, globals: object = PyObject.dict(), locals: object = PyObject.dict()) -> PyObject:
1042
+ """
1043
+ Copies the locals, execs the Python code, and returns the locals with any updates.
1044
+ """
1045
+
1046
+
1047
+ TS = TypeVarTuple("TS")
1048
+
1049
+ T1 = TypeVar("T1")
1050
+ T2 = TypeVar("T2")
1051
+ T3 = TypeVar("T3")
1052
+
1053
+
1054
+ class UnstableFn(BuiltinExpr, Generic[T, Unpack[TS]]):
1055
+ @overload
1056
+ def __init__(self, f: Callable[[Unpack[TS]], T]) -> None: ...
1057
+
1058
+ @overload
1059
+ def __init__(self, f: Callable[[T1, Unpack[TS]], T], _a: T1, /) -> None: ...
1060
+
1061
+ @overload
1062
+ def __init__(self, f: Callable[[T1, T2, Unpack[TS]], T], _a: T1, _b: T2, /) -> None: ...
1063
+
1064
+ # Removing due to bug in MyPy
1065
+ # https://github.com/python/mypy/issues/17212
1066
+ # @overload
1067
+ # def __init__(self, f: Callable[[T1, T2, T3, Unpack[TS]], T], _a: T1, _b: T2, _c: T3, /) -> None: ...
1068
+
1069
+ # etc, for partial application
1070
+
1071
+ @method(egg_fn="unstable-fn")
1072
+ def __init__(self, f, *partial) -> None: ...
1073
+
1074
+ @method(preserve=True)
1075
+ @deprecated("use .value")
1076
+ def eval(self) -> Callable[[Unpack[TS]], T]:
1077
+ return self.value
1078
+
1079
+ @method(preserve=True) # type: ignore[prop-decorator]
1080
+ @property
1081
+ def value(self) -> Callable[[Unpack[TS]], T]:
1082
+ """
1083
+ If this is a constructor, returns either the callable directly or a `functools.partial` function if args are provided.
1084
+ """
1085
+ if (fn := get_literal_value(self)) is not None:
1086
+ return fn
1087
+ raise ExprValueError(self, "UnstableFn(f, *args)")
1088
+
1089
+ __match_args__ = ("value",)
1090
+
1091
+ @method(egg_fn="unstable-app")
1092
+ def __call__(self, *args: Unpack[TS]) -> T: ...
1093
+
1094
+
1095
+ # Method Type is for builtins like __getitem__
1096
+ converter(MethodType, UnstableFn, lambda m: UnstableFn(m.__func__, m.__self__))
1097
+ converter(RuntimeFunction, UnstableFn, UnstableFn)
1098
+ converter(partial, UnstableFn, lambda p: UnstableFn(p.func, *p.args))
1099
+
1100
+
1101
+ def _convert_function(fn: FunctionType) -> UnstableFn:
1102
+ """
1103
+ Converts a function type to an unstable function. This function will be an anon function in egglog.
1104
+
1105
+ Would just be UnstableFn(function(a)) but we have to account for unbound vars within the body.
1106
+
1107
+ This means that we have to turn all of those unbound vars into args to the function, and then
1108
+ partially apply them, alongside creating a default rewrite for the function.
1109
+ """
1110
+ decls = Declarations()
1111
+ return_type, *arg_types = [resolve_type_annotation_mutate(decls, tp) for tp in get_type_args()]
1112
+ arg_names = [p.name for p in signature(fn).parameters.values()]
1113
+ arg_decls = [
1114
+ TypedExprDecl(tp.to_just(), UnboundVarDecl(name)) for name, tp in zip(arg_names, arg_types, strict=True)
1115
+ ]
1116
+ res = resolve_literal(
1117
+ return_type, fn(*(RuntimeExpr.__from_values__(decls, a) for a in arg_decls)), Thunk.value(decls)
1118
+ )
1119
+ res_expr = res.__egg_typed_expr__
1120
+ decls |= res
1121
+ # these are all the args that appear in the body that are not bound by the args of the function
1122
+ unbound_vars = list(collect_unbound_vars(res_expr) - set(arg_decls))
1123
+ # prefix the args with them
1124
+ fn_ref = UnnamedFunctionRef(tuple(unbound_vars + arg_decls), res_expr)
1125
+ rewrite_decl = DefaultRewriteDecl(fn_ref, res_expr.expr, subsume=True)
1126
+ ruleset_decls = _add_default_rewrite_inner(decls, rewrite_decl, get_current_ruleset())
1127
+ ruleset_decls |= res
1128
+
1129
+ fn = RuntimeFunction(Thunk.value(decls), Thunk.value(fn_ref))
1130
+ return UnstableFn(fn, *(RuntimeExpr.__from_values__(decls, v) for v in unbound_vars))
1131
+
1132
+
1133
+ converter(FunctionType, UnstableFn, _convert_function)