Merge pull request #12 from metadsl/increase-line-length

Increase max line length to 120

Author: saulshanabrook

pyproject.toml

@@ -42,6 +42,9 @@ docs = [
 ]
 
 
+[tool.black]
+line-length = 120
+
 [tool.isort]
 profile = "black"
 skip_gitignore = true

python/egg_smol/bindings.pyi

@@ -7,17 +7,13 @@ def parse(input: str) -> list[_Command]: ...
 class EGraph:
     def parse_and_run_program(self, input: str) -> list[str]: ...
     def declare_constructor(self, variant: Variant, sort: str) -> None: ...
-    def declare_sort(
-        self, name: str, presort_and_args: tuple[str, list[_Expr]] | None = None
-    ) -> None: ...
+    def declare_sort(self, name: str, presort_and_args: tuple[str, list[_Expr]] | None = None) -> None: ...
     def declare_function(self, decl: FunctionDecl) -> None: ...
     def define(self, name: str, expr: _Expr, cost: int | None = None) -> None: ...
     def add_rewrite(self, rewrite: Rewrite) -> str: ...
     def run_rules(self, limit: int) -> tuple[timedelta, timedelta, timedelta]: ...
     def check_fact(self, fact: _Fact) -> None: ...
-    def extract_expr(
-        self, expr: _Expr, variants: int = 0
-    ) -> tuple[int, _Expr, list[_Expr]]: ...
+    def extract_expr(self, expr: _Expr, variants: int = 0) -> tuple[int, _Expr, list[_Expr]]: ...
     def add_rule(self, rule: Rule) -> str: ...
     def eval_actions(self, *actions: _Action) -> None: ...
     def push(self) -> None: ...
@@ -133,9 +129,7 @@ class FunctionDecl:
 
 @final
 class Variant:
-    def __init__(
-        self, name: str, types: list[str], cost: int | None = None
-    ) -> None: ...
+    def __init__(self, name: str, types: list[str], cost: int | None = None) -> None: ...
     name: str
     types: list[str]
     cost: int | None
@@ -158,9 +152,7 @@ class Rewrite:
     rhs: _Expr
     conditions: list[_Fact]
 
-    def __init__(
-        self, lhs: _Expr, rhs: _Expr, conditions: list[_Fact] = []
-    ) -> None: ...
+    def __init__(self, lhs: _Expr, rhs: _Expr, conditions: list[_Fact] = []) -> None: ...
 
 @final
 class Datatype:

python/egg_smol/declarations.py

@@ -87,9 +87,7 @@ class Declarations:
     # Bidirectional mapping between egg function names and python callable references.
     # Note that there are possibly mutliple callable references for a single egg function name, like `+`
     # for both int and rational classes.
-    egg_fn_to_callable_refs: defaultdict[str, set[CallableRef]] = field(
-        default_factory=lambda: defaultdict(set)
-    )
+    egg_fn_to_callable_refs: defaultdict[str, set[CallableRef]] = field(default_factory=lambda: defaultdict(set))
     callable_ref_to_egg_fn: dict[CallableRef, str] = field(default_factory=dict)
 
     # Bidirectional mapping between egg sort names and python type references.
@@ -119,9 +117,7 @@ def get_function_decl(self, ref: CallableRef) -> FunctionDecl:
             return self.constants[ref.name].to_function_decl()
         assert_never(ref)
 
-    def register_sort(
-        self, type_ref: JustTypeRef, egg_name: Optional[str] = None
-    ) -> str:
+    def register_sort(self, type_ref: JustTypeRef, egg_name: Optional[str] = None) -> str:
         egg_name = egg_name or type_ref.generate_egg_name()
         if egg_name in self.egg_sort_to_type_ref:
             raise ValueError(f"Sort {egg_name} is already registered.")
@@ -154,10 +150,7 @@ def to_egg(self, decls: Declarations, egraph: bindings.EGraph) -> str:
             raise ValueError(f"Type {self.name} is not registered.")
         # If this is a type with arguments and it is not registered, then we need to register it
         egg_name = decls.register_sort(self)
-        arg_sorts = [
-            cast(bindings._Expr, bindings.Var(a.to_egg(decls, egraph)))
-            for a in self.args
-        ]
+        arg_sorts = [cast(bindings._Expr, bindings.Var(a.to_egg(decls, egraph))) for a in self.args]
         egraph.declare_sort(egg_name, (self.name, arg_sorts))
         return egg_name
 
@@ -265,9 +258,7 @@ class FunctionDecl:
     default: Optional[ExprDecl] = None
     merge: Optional[ExprDecl] = None
 
-    def to_egg(
-        self, decls: Declarations, egraph: bindings.EGraph, ref: CallableRef
-    ) -> bindings.FunctionDecl:
+    def to_egg(self, decls: Declarations, egraph: bindings.EGraph, ref: CallableRef) -> bindings.FunctionDecl:
         return bindings.FunctionDecl(
             decls.callable_ref_to_egg_fn[ref],
             # Remove all vars from the type refs, raising an errory if we find one,
@@ -288,9 +279,7 @@ class VarDecl:
 
     @classmethod
     def from_egg(cls, var: bindings.Var) -> tuple[JustTypeRef, LitDecl]:
-        raise NotImplementedError(
-            "Cannot turn var into egg type because typing unknown."
-        )
+        raise NotImplementedError("Cannot turn var into egg type because typing unknown.")
 
     def to_egg(self, _decls: Declarations) -> bindings.Var:
         return bindings.Var(self.name)
@@ -349,14 +338,10 @@ class CallDecl:
 
     def __post_init__(self):
         if self.bound_tp_params and not isinstance(self.callable, ClassMethodRef):
-            raise ValueError(
-                "Cannot bind type parameters to a non-class method callable."
-            )
+            raise ValueError("Cannot bind type parameters to a non-class method callable.")
 
     @classmethod
-    def from_egg(
-        cls, decls: Declarations, call: bindings.Call
-    ) -> tuple[JustTypeRef, CallDecl]:
+    def from_egg(cls, decls: Declarations, call: bindings.Call) -> tuple[JustTypeRef, CallDecl]:
         from .type_constraint_solver import TypeConstraintSolver
 
         results = [tp_and_expr_decl_from_egg(decls, a) for a in call.args]
@@ -367,9 +352,7 @@ def from_egg(
         for callable_ref in decls.egg_fn_to_callable_refs[call.name]:
             tcs = TypeConstraintSolver()
             fn_decl = decls.get_function_decl(callable_ref)
-            return_tp = tcs.infer_return_type(
-                fn_decl.arg_types, fn_decl.return_type, arg_types
-            )
+            return_tp = tcs.infer_return_type(fn_decl.arg_types, fn_decl.return_type, arg_types)
             return return_tp, cls(callable_ref, arg_decls)
         raise ValueError(f"Could not find callable ref for call {call}")
 
@@ -421,34 +404,12 @@ def test_expr_pretty():
     assert LitDecl("foo").pretty() == 'String("foo")'
     assert LitDecl(None).pretty() == "unit()"
     assert CallDecl(FunctionRef("foo"), (VarDecl("x"),)).pretty() == "foo(x)"
-    assert (
-        CallDecl(
-            FunctionRef("foo"), (VarDecl("x"), VarDecl("y"), VarDecl("z"))
-        ).pretty()
-        == "foo(x, y, z)"
-    )
-    assert (
-        CallDecl(MethodRef("foo", "__add__"), (VarDecl("x"), VarDecl("y"))).pretty()
-        == "x + y"
-    )
-    assert (
-        CallDecl(MethodRef("foo", "__getitem__"), (VarDecl("x"), VarDecl("y"))).pretty()
-        == "x[y]"
-    )
-    assert (
-        CallDecl(
-            ClassMethodRef("foo", "__init__"), (VarDecl("x"), VarDecl("y"))
-        ).pretty()
-        == "foo(x, y)"
-    )
-    assert (
-        CallDecl(ClassMethodRef("foo", "bar"), (VarDecl("x"), VarDecl("y"))).pretty()
-        == "foo.bar(x, y)"
-    )
-    assert (
-        CallDecl(MethodRef("foo", "__call__"), (VarDecl("x"), VarDecl("y"))).pretty()
-        == "x(y)"
-    )
+    assert CallDecl(FunctionRef("foo"), (VarDecl("x"), VarDecl("y"), VarDecl("z"))).pretty() == "foo(x, y, z)"
+    assert CallDecl(MethodRef("foo", "__add__"), (VarDecl("x"), VarDecl("y"))).pretty() == "x + y"
+    assert CallDecl(MethodRef("foo", "__getitem__"), (VarDecl("x"), VarDecl("y"))).pretty() == "x[y]"
+    assert CallDecl(ClassMethodRef("foo", "__init__"), (VarDecl("x"), VarDecl("y"))).pretty() == "foo(x, y)"
+    assert CallDecl(ClassMethodRef("foo", "bar"), (VarDecl("x"), VarDecl("y"))).pretty() == "foo.bar(x, y)"
+    assert CallDecl(MethodRef("foo", "__call__"), (VarDecl("x"), VarDecl("y"))).pretty() == "x(y)"
     assert (
         CallDecl(
             ClassMethodRef("Map", "__init__"),
@@ -462,9 +423,7 @@ def test_expr_pretty():
 ExprDecl = Union[VarDecl, LitDecl, CallDecl]
 
 
-def tp_and_expr_decl_from_egg(
-    decls: Declarations, expr: bindings._Expr
-) -> tuple[JustTypeRef, ExprDecl]:
+def tp_and_expr_decl_from_egg(decls: Declarations, expr: bindings._Expr) -> tuple[JustTypeRef, ExprDecl]:
     if isinstance(expr, bindings.Var):
         return VarDecl.from_egg(expr)
     if isinstance(expr, bindings.Lit):
@@ -557,9 +516,7 @@ class DeleteDecl:
     call: CallDecl
 
     def to_egg(self, decls: Declarations) -> bindings.Delete:
-        return bindings.Delete(
-            self.call.callable.to_egg(decls), [a.to_egg(decls) for a in self.call.args]
-        )
+        return bindings.Delete(self.call.callable.to_egg(decls), [a.to_egg(decls) for a in self.call.args])
 
 
 @dataclass(frozen=True)

python/egg_smol/egraph.py

@@ -59,16 +59,9 @@ def extract_multiple(self, expr: EXPR, n: int) -> list[EXPR]:
         """
         tp, decl = expr_parts(expr)
         egg_expr = decl.to_egg(self._decls)
-        _cost, new_egg_expr, variants = self._egraph.extract_expr(
-            egg_expr, variants=n + 1
-        )
-        new_decls = [
-            tp_and_expr_decl_from_egg(self._decls, egg_expr)[1]
-            for egg_expr in variants[::-1]
-        ]
-        return [
-            cast(EXPR, RuntimeExpr(self._decls, tp, new_decl)) for new_decl in new_decls
-        ]
+        _cost, new_egg_expr, variants = self._egraph.extract_expr(egg_expr, variants=n + 1)
+        new_decls = [tp_and_expr_decl_from_egg(self._decls, egg_expr)[1] for egg_expr in variants[::-1]]
+        return [cast(EXPR, RuntimeExpr(self._decls, tp, new_decl)) for new_decl in new_decls]
 
     def define(self, name: str, expr: EXPR, cost: Optional[int] = None) -> EXPR:
         """

python/egg_smol/monkeypatch.py

@@ -28,8 +28,6 @@ def _evaluate_monkeypatch(self, globalns, localns):
             "Forward references must evaluate to types.",
             is_argument=self.__forward_is_argument__,
         )
-        self.__forward_value__ = typing._eval_type(  # type: ignore
-            type_, globalns, localns
-        )
+        self.__forward_value__ = typing._eval_type(type_, globalns, localns)  # type: ignore
         self.__forward_evaluated__ = True
     return self.__forward_value__

python/egg_smol/registry.py

@@ -132,9 +132,7 @@ def _class(
         """
         cls_name = cls.__name__
         # Get all the methods from the class
-        cls_dict: dict[str, Any] = {
-            k: v for k, v in cls.__dict__.items() if k not in IGNORED_ATTRIBUTES
-        }
+        cls_dict: dict[str, Any] = {k: v for k, v in cls.__dict__.items() if k not in IGNORED_ATTRIBUTES}
         parameters: list[TypeVar] = cls_dict.pop("__parameters__", [])
 
         # Register class first
@@ -147,9 +145,7 @@ def _class(
         self._on_register_sort(egg_sort or cls_name)
 
         # The type ref of self is paramterized by the type vars
-        slf_type_ref = TypeRefWithVars(
-            cls_name, tuple(ClassTypeVarRef(i) for i in range(n_type_vars))
-        )
+        slf_type_ref = TypeRefWithVars(cls_name, tuple(ClassTypeVarRef(i) for i in range(n_type_vars)))
 
         # Then register each of its methods
         for method_name, method in cls_dict.items():
@@ -321,9 +317,7 @@ def _generate_function_decl(
         cls_type_and_name: Optional[tuple[type | RuntimeClass, str]] = None,
     ) -> FunctionDecl:
         if not isinstance(fn, FunctionType):
-            raise NotImplementedError(
-                f"Can only generate function decls for functions not {fn}  {type(fn)}"
-            )
+            raise NotImplementedError(f"Can only generate function decls for functions not {fn}  {type(fn)}")
 
         hint_globals = fn.__globals__.copy()
 
@@ -336,31 +330,20 @@ def _generate_function_decl(
                 raise ValueError("Init function must have a self type")
             return_type = first_arg
         else:
-            return_type = self._resolve_type_annotation(
-                hints["return"], cls_typevars, cls_type_and_name
-            )
+            return_type = self._resolve_type_annotation(hints["return"], cls_typevars, cls_type_and_name)
 
         params = list(signature(fn).parameters.values())
         # Remove first arg if this is a classmethod or a method, since it won't have an annotation
         if first_arg is not None:
             first, *params = params
             if first.annotation != Parameter.empty:
-                raise ValueError(
-                    f"First arg of a method must not have an annotation, not {first.annotation}"
-                )
+                raise ValueError(f"First arg of a method must not have an annotation, not {first.annotation}")
 
         for param in params:
             if param.kind != Parameter.POSITIONAL_OR_KEYWORD:
-                raise ValueError(
-                    f"Can only register functions with positional or keyword args, not {param.kind}"
-                )
+                raise ValueError(f"Can only register functions with positional or keyword args, not {param.kind}")
 
-        arg_types = tuple(
-            self._resolve_type_annotation(
-                hints[t.name], cls_typevars, cls_type_and_name
-            )
-            for t in params
-        )
+        arg_types = tuple(self._resolve_type_annotation(hints[t.name], cls_typevars, cls_type_and_name) for t in params)
         # If the first arg is a self, and this not an __init__ fn, add this as a typeref
         if isinstance(first_arg, (ClassTypeVarRef, TypeRefWithVars)) and not is_init:
             arg_types = (first_arg,) + arg_types
@@ -405,13 +388,9 @@ def _resolve_type_annotation(
                 raise TypeError("Union types are only supported for type promotion")
             fst, snd = args
             if fst in {int, str}:
-                return self._resolve_type_annotation(
-                    snd, cls_typevars, cls_type_and_name
-                )
+                return self._resolve_type_annotation(snd, cls_typevars, cls_type_and_name)
             if snd in {int, str}:
-                return self._resolve_type_annotation(
-                    fst, cls_typevars, cls_type_and_name
-                )
+                return self._resolve_type_annotation(fst, cls_typevars, cls_type_and_name)
             raise TypeError("Union types are only supported for type promotion")
 
         # If this is the type for the class, use the class name
@@ -471,9 +450,7 @@ class _WrappedMethod(Generic[P, EXPR]):
     fn: Callable[P, EXPR]
 
     def __call__(self, *args: P.args, **kwargs: P.kwargs) -> EXPR:
-        raise NotImplementedError(
-            "We should never call a wrapped method. Did you forget to wrap the class?"
-        )
+        raise NotImplementedError("We should never call a wrapped method. Did you forget to wrap the class?")
 
 
 class BaseExpr:
@@ -724,9 +701,7 @@ def __str__(self) -> str:
     def _to_decl(self) -> SetDecl:
         lhs = expr_parts(self.lhs)[1]
         if not isinstance(lhs, CallDecl):
-            raise ValueError(
-                f"Can only create a call with a call for the lhs, got {lhs}"
-            )
+            raise ValueError(f"Can only create a call with a call for the lhs, got {lhs}")
         return SetDecl(lhs, expr_parts(self.rhs)[1])