diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index dca3c646..3f633557 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,12 +1,12 @@ exclude: ^python/tests/__snapshots__/ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.14.3 + rev: v0.16.0 hooks: - id: ruff-check args: [--fix] - id: ruff-format - repo: https://github.com/astral-sh/uv-pre-commit - rev: 0.9.7 + rev: 0.11.32 hooks: - id: uv-lock diff --git a/python/egglog/__init__.py b/python/egglog/__init__.py index 27b2bc61..5f937c99 100644 --- a/python/egglog/__init__.py +++ b/python/egglog/__init__.py @@ -4,7 +4,7 @@ from . import config, ipython_magic # noqa: F401 from .bindings import EggSmolError, StageInfo, TimeOnly, WithPlan # noqa: F401 -from .builtins import * # noqa: UP029 +from .builtins import * from .conversion import * from .deconstruct import * from .egraph import * diff --git a/python/egglog/declarations.py b/python/egglog/declarations.py index 9a5d08e2..5e53cf42 100644 --- a/python/egglog/declarations.py +++ b/python/egglog/declarations.py @@ -49,7 +49,6 @@ "ConstantRef", "ConstructorDecl", "Declarations", - "Declarations", "DeclarationsLike", "DefaultRewriteDecl", "DelayedDeclarations", @@ -128,7 +127,7 @@ class HasDeclarations(Protocol): def __egg_decls__(self) -> Declarations: ... -DeclarationsLike: TypeAlias = Union[HasDeclarations, None, "Declarations"] +DeclarationsLike: TypeAlias = Union[HasDeclarations, "Declarations", None] def upcast_declarations(declarations_like: Iterable[DeclarationsLike]) -> list[Declarations]: @@ -659,7 +658,7 @@ def signature(self) -> FunctionSignature: ) @property - def egg_name(self) -> None | str: + def egg_name(self) -> str | None: return None diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py index 6fe782d5..e9334371 100644 --- a/python/egglog/egraph.py +++ b/python/egglog/egraph.py @@ -60,13 +60,11 @@ "BaseExpr", "BuiltinExpr", "Command", - "Command", "CostModel", "EGraph", "Expr", "ExprCallable", "Fact", - "Fact", "GraphvizKwargs", "GreedyDagCost", "RewriteOrRule", @@ -1538,7 +1536,7 @@ def __call__(self, *args, **kwargs) -> Never: def ruleset( rule_or_generator: RewriteOrRule | RewriteOrRuleGenerator | None = None, *rules: RewriteOrRule, - name: None | str = None, + name: str | None = None, ) -> Ruleset: """ Creates a ruleset with the following rules. @@ -2086,7 +2084,7 @@ def run(ruleset: Ruleset | None = None, *until: FactLike, scheduler: BackOff | N ) -def back_off(match_limit: None | int = None, ban_length: None | int = None) -> BackOff: +def back_off(match_limit: int | None = None, ban_length: int | None = None) -> BackOff: """ Create a backoff scheduler configuration. diff --git a/python/egglog/exp/any_expr.py b/python/egglog/exp/any_expr.py index 4de0d3ff..07eb8097 100644 --- a/python/egglog/exp/any_expr.py +++ b/python/egglog/exp/any_expr.py @@ -753,8 +753,8 @@ def any_eval(self: A) -> Any: return res -_CURRENT_EGRAPH: None | EGraph = None -_LAST_ASSERT: None | A = None +_CURRENT_EGRAPH: EGraph | None = None +_LAST_ASSERT: A | None = None @contextlib.contextmanager diff --git a/python/egglog/exp/array_api.py b/python/egglog/exp/array_api.py index 42ea19a0..ba4d6290 100644 --- a/python/egglog/exp/array_api.py +++ b/python/egglog/exp/array_api.py @@ -1459,7 +1459,7 @@ def slice(cls, slice: Slice) -> MultiAxisIndexKeyItem: ... converter(Int, MultiAxisIndexKeyItem, lambda i: MultiAxisIndexKeyItem.int(i)) converter(Slice, MultiAxisIndexKeyItem, lambda s: MultiAxisIndexKeyItem.slice(s)) -MultiAxisIndexKeyItemLike: TypeAlias = MultiAxisIndexKeyItem | EllipsisType | None | IntLike | SliceLike +MultiAxisIndexKeyItemLike: TypeAlias = MultiAxisIndexKeyItem | EllipsisType | IntLike | SliceLike | None class MultiAxisIndexKey(Expr, ruleset=array_api_ruleset): @@ -2123,7 +2123,7 @@ def int(cls, value: Int) -> OptionalIntOrTuple: ... def tuple(cls, value: TupleIntLike) -> OptionalIntOrTuple: ... -OptionalIntOrTupleLike: TypeAlias = OptionalIntOrTuple | None | IntLike | TupleIntLike +OptionalIntOrTupleLike: TypeAlias = OptionalIntOrTuple | IntLike | TupleIntLike | None converter(type(None), OptionalIntOrTuple, lambda _: OptionalIntOrTuple.none) converter(Int, OptionalIntOrTuple, lambda v: OptionalIntOrTuple.int(v)) @@ -2379,7 +2379,8 @@ def vecdot(x1: NDArrayLike, x2: NDArrayLike) -> NDArray: x1.shape.drop_last(), x1.dtype, lambda idx: ( - TupleInt.range(x1.shape.last()) + TupleInt + .range(x1.shape.last()) .map_value(lambda i: x1.index(idx.append(i)) * x2.index((i,))) .foldl_value(Value.__add__, Value.from_float(0)) ), @@ -2740,7 +2741,8 @@ def unravel_index(flat_index: IntLike, shape: TupleIntLike) -> TupleInt: shape = cast("TupleInt", shape) return ( - shape.reverse() + shape + .reverse() .foldl_tuple_int( # Store remainder as last item in accumulator lambda acc, dim: acc.drop_last().append((r := acc.last()) % dim).append(r // dim), @@ -2754,7 +2756,7 @@ def unravel_index(flat_index: IntLike, shape: TupleIntLike) -> TupleInt: array_api_combined_ruleset = array_api_ruleset array_api_schedule = (array_api_combined_ruleset + run()).saturate() -_CURRENT_EGRAPH: None | EGraph = None +_CURRENT_EGRAPH: EGraph | None = None @contextlib.contextmanager diff --git a/python/egglog/runtime.py b/python/egglog/runtime.py index 2ec5a453..bebe9c55 100644 --- a/python/egglog/runtime.py +++ b/python/egglog/runtime.py @@ -649,7 +649,7 @@ def to_py_signature(sig: FunctionSignature, decls: Declarations, optional_args: return Signature(parameters) -ON_CREATE_EXPR: None | Callable[[Callable[[], TypedExprDecl]], None] = None +ON_CREATE_EXPR: Callable[[Callable[[], TypedExprDecl]], None] | None = None @dataclass diff --git a/python/tests/test_array_api.py b/python/tests/test_array_api.py index d9eea325..e401ae39 100644 --- a/python/tests/test_array_api.py +++ b/python/tests/test_array_api.py @@ -263,11 +263,14 @@ def linalg_norm(X: NDArray, axis: TupleIntLike) -> NDArray: return NDArray.fn( outshape, X.dtype, - lambda k: LoopNestAPI.from_tuple(reduce_axis) - .unwrap() - .indices() - .foldl_value(lambda carry, i: carry + ((x := X.index(i + k)).conj() * x).real(), init=0.0) - .sqrt(), + lambda k: ( + LoopNestAPI + .from_tuple(reduce_axis) + .unwrap() + .indices() + .foldl_value(lambda carry, i: carry + ((x := X.index(i + k)).conj() * x).real(), init=0.0) + .sqrt() + ), ) @@ -277,9 +280,11 @@ def linalg_norm_v2(X: NDArrayLike, axis: TupleIntLike) -> NDArray: return NDArray.fn( X.shape.deselect(axis), X.dtype, - lambda k: ndindex(X.shape.select(axis)) - .foldl_value(lambda carry, i: carry + ((x := X.index(i + k)).conj() * x).real(), init=0.0) - .sqrt(), + lambda k: ( + ndindex(X.shape.select(axis)) + .foldl_value(lambda carry, i: carry + ((x := X.index(i + k)).conj() * x).real(), init=0.0) + .sqrt() + ), )