Linting fixes for entire library.

Author: hameerabbasi

pixi.toml

@@ -24,6 +24,7 @@ sparse = "*"
 numba = ">=0.60"
 scipy = "*"
 numpy = "==2.*"
+mypy = ">=1.15.0,<2"
 
 [feature.test.tasks]
 test = { cmd = "pytest", depends-on = ["compile"] }

pyproject.toml

@@ -21,3 +21,23 @@ numba = "^0.61.0"
 [build-system]
 requires = ["poetry-core>=1.0.8"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.ruff.lint]
+select = ["F", "E", "W", "I", "B", "UP", "YTT", "BLE", "C4", "T10", "ISC", "ICN", "PIE", "PYI", "RSE", "RET", "SIM", "PGH", "FLY", "NPY", "PERF"]
+
+[tool.ruff.lint.isort.sections]
+numpy = ["numpy", "numpy.*", "scipy", "scipy.*"]
+
+[tool.ruff.format]
+quote-style = "double"
+docstring-code-format = true
+
+[tool.ruff.lint.isort]
+section-order = [
+    "future",
+    "standard-library",
+    "first-party",
+    "numpy",
+    "third-party",
+    "local-folder",
+]

src/finch/__init__.py

@@ -1,2 +1,31 @@
 from . import finch_logic
-from .interface import *
+from .interface import (
+    compute,
+    elementwise,
+    expand_dims,
+    fuse,
+    fused,
+    identify,
+    lazy,
+    multiply,
+    permute_dims,
+    prod,
+    reduce,
+    squeeze,
+)
+
+__all__ = [
+    "lazy",
+    "compute",
+    "finch_logic",
+    "fuse",
+    "fused",
+    "permute_dims",
+    "expand_dims",
+    "squeeze",
+    "identify",
+    "reduce",
+    "elementwise",
+    "prod",
+    "multiply",
+]

src/finch/algebra/__init__.py

@@ -1,4 +1,13 @@
-from .algebra import *
+from .algebra import (
+    element_type,
+    fill_value,
+    fixpoint_type,
+    init_value,
+    is_associative,
+    query_property,
+    register_property,
+    return_type,
+)
 
 __all__ = [
     "fill_value",
@@ -9,4 +18,4 @@
     "is_associative",
     "query_property",
     "register_property",
-]
+]

src/finch/algebra/algebra.py

@@ -1,6 +1,3 @@
-from typing import Any, Type
-from collections.abc import Hashable
-
 """
 Finch performs extensive rewriting and defining of functions.  The Finch
 compiler is designed to inspect objects and functions defined by other
@@ -22,33 +19,40 @@
 
 ```python
 from finch import register_property
-register_property(complex, '__add__', 'is_associative', lambda obj: True)
+
+register_property(complex, "__add__", "is_associative", lambda obj: True)
 ```
 
 Finch includes a convenience functions to query each property as well,
 for example:
 ```python
 from finch import query_property
 from operator import add
-query_property(complex, '__add__', 'is_associative')
+
+query_property(complex, "__add__", "is_associative")
 # True
 is_associative(add, complex, complex)
 # True
 ```
 
-Properties can be inherited in the same way as methods. First we check whether properties have been defined for the object itself (in the case of functions), then we check For example, if you
-register a property for a class, all subclasses of that class will inherit
-that property. This allows you to define properties for a class and have
-them automatically apply to all subclasses, without having to register the
-property for each subclass individually.
+Properties can be inherited in the same way as methods. First we check whether
+properties have been defined for the object itself (in the case of functions), then we
+check For example, if you register a property for a class, all subclasses of that class
+will inherit that property. This allows you to define properties for a class and have
+them automatically apply to all subclasses, without having to register the property for
+each subclass individually.
 """
+
 import operator
-from typing import Union
+from collections.abc import Hashable
+from typing import Any, Callable
+
 import numpy as np
 
-_properties = {}
+_properties: dict[tuple[Hashable, str, str], Any] = {}
 
-def query_property(obj, attr, prop, *args):
+
+def query_property(obj: Hashable, attr: str, prop: Hashable, *args: Any) -> Any:
     """Queries a property of an attribute of an object or class.  Properties can
     be overridden by calling register_property on the object or it's class.
 
@@ -64,22 +68,28 @@ def query_property(obj, attr, prop, *args):
     Raises:
         NotImplementedError: If the property is not implemented for the given type.
     """
-    if isinstance(obj, type):
-        T = obj
-    else:
-        if isinstance(obj, Hashable):
-            if (obj, attr, prop) in _properties:
-                return _properties[(obj, attr, prop)](obj, *args)
+    T = obj
+    if not isinstance(obj, Hashable):
         T = type(obj)
-    while True:
-        if (T, attr, prop) in _properties:
-            return _properties[(T, attr, prop)](obj, *args)
-        if T is object:
-            break
-        T = T.__base__
+    to_query = {T}
+    queried: set[type] = set()
+    while len(to_query) != len(queried):
+        to_query_new = to_query.copy()
+        for o in to_query:
+            method = _properties.get((o, attr, prop), None)
+            if method is not None:
+                return method(obj, *args)
+            queried.add(o)
+            if not isinstance(o, type):
+                to_query_new.update(type(o))
+                continue
+            to_query_new.update(o.__mro__)
+        to_query = to_query_new
+
     raise NotImplementedError(f"Property {prop} not implemented for {type(obj)}")
 
-def register_property(cls, attr, prop, f):
+
+def register_property(cls: type, attr: str, prop: str, f: Callable) -> None:
     """Registers a property for a class or object.
 
     Args:
@@ -90,6 +100,7 @@ def register_property(cls, attr, prop, f):
     """
     _properties[(cls, attr, prop)] = f
 
+
 def fill_value(arg: Any) -> Any:
     """The fill value for the given argument.  The fill value is the
     default value for a tensor when it is created with a given shape and dtype,
@@ -104,11 +115,15 @@ def fill_value(arg: Any) -> Any:
     Raises:
         NotImplementedError: If the fill value is not implemented for the given type.
     """
-    return query_property(arg, '__self__', 'fill_value')
+    return query_property(arg, "__self__", "fill_value")
+
+
+register_property(
+    np.ndarray, "__self__", "fill_value", lambda x: np.zeros((), dtype=x.dtype)[()]
+)
 
-register_property(np.ndarray, '__self__', 'fill_value', lambda x: np.zeros((), dtype=x.dtype)[()])
 
-def element_type(arg: Any) -> Type:
+def element_type(arg: Any) -> type:
     """The element type of the given argument.  The element type is the scalar type of
     the elements in a tensor, which may be different from the data type of the
     tensor.
@@ -122,9 +137,16 @@ def element_type(arg: Any) -> Type:
     Raises:
         NotImplementedError: If the element type is not implemented for the given type.
     """
-    return query_property(arg, '__self__', 'element_type')
+    return query_property(arg, "__self__", "element_type")
+
+
+register_property(
+    np.ndarray,
+    "__self__",
+    "element_type",
+    lambda x: type(np.zeros((), dtype=x.dtype)[()]),
+)
 
-register_property(np.ndarray, '__self__', 'element_type', lambda x: type(np.zeros((), dtype=x.dtype)[()]))
 
 def return_type(op: Any, *args: Any) -> Any:
     """The return type of the given function on the given argument types.
@@ -136,7 +158,8 @@ def return_type(op: Any, *args: Any) -> Any:
     Returns:
         The return type of op(*args: arg_types)
     """
-    return query_property(op, '__call__', 'return_type', *args)
+    return query_property(op, "__call__", "return_type", *args)
+
 
 StableNumber = (np.number, bool, int, float, complex)
 
@@ -158,34 +181,52 @@ def return_type(op: Any, *args: Any) -> Any:
 }
 
 for op, (meth, rmeth) in _reflexive_operators.items():
-    register_property(op, '__call__', 'return_type', lambda op, a, b: query_property(a, meth, 'return_type', b) if hasattr(a, meth) else query_property(b, rmeth, 'return_type', a)),
+    (
+        register_property(
+            op,
+            "__call__",
+            "return_type",
+            lambda op, a, b, meth=meth, rmeth=rmeth: query_property(
+                a, meth, "return_type", b
+            )
+            if hasattr(a, meth)
+            else query_property(b, rmeth, "return_type", a),
+        ),
+    )
+
     def _return_type(meth):
         def _return_type_closure(a, b):
             if issubclass(b, StableNumber):
                 return type(getattr(a(True), meth)(b(True)))
-            else:
-                raise TypeError(f"Unsupported operand type for {type(a)}.{meth}:  {type(b)}")
+            raise TypeError(
+                f"Unsupported operand type for {type(a)}.{meth}:  {type(b)}"
+            )
+
         return _return_type_closure
+
     for T in StableNumber:
-        register_property(T, meth, 'return_type', _return_type(meth))
-        register_property(T, rmeth, 'return_type', _return_type(rmeth))
+        register_property(T, meth, "return_type", _return_type(meth))
+        register_property(T, rmeth, "return_type", _return_type(rmeth))
+
 
 def is_associative(op: Any) -> bool:
     """Returns whether the given function is associative, that is, whether the
     op(op(a, b), c) == op(a, op(b, c)) for all a, b, c.
 
     Args:
         op: The function to check.
-    
+
     Returns:
         True if the function can be proven to be associative, False otherwise.
     """
-    return query_property(op, '__call__', 'is_associative')
+    return query_property(op, "__call__", "is_associative")
+
 
 for op in [operator.add, operator.mul, operator.and_, operator.xor, operator.or_]:
-    register_property(op, '__call__', 'is_associative', lambda op: True) 
+    register_property(op, "__call__", "is_associative", lambda op: True)
 
-def fixpoint_type(op: Any, z: Any, T: Type) -> Type:
+
+def fixpoint_type(op: Any, z: Any, T: type) -> type:
     """Determines the fixpoint type after repeated calling the given operation.
 
     Args:
@@ -200,9 +241,12 @@ def fixpoint_type(op: Any, z: Any, T: Type) -> Type:
     R = type(z)
     while R not in S:
         S.add(R)
-        R = return_type(op, type(z), T)  # Assuming `op` is a callable that takes `z` and `T` as arguments
+        R = return_type(
+            op, type(z), T
+        )  # Assuming `op` is a callable that takes `z` and `T` as arguments
     return R
 
+
 def init_value(op, arg) -> Any:
     """Returns the initial value for a reduction operation on the given type.
 
@@ -214,17 +258,24 @@ def init_value(op, arg) -> Any:
         The initial value for the given operation and type.
 
     Raises:
-        NotImplementedError: If the initial value is not implemented for the given type and operation.
+        NotImplementedError: If the initial value is not implemented for the given type
+        and operation.
     """
-    return query_property(op, '__call__', 'init_value', arg)
+    return query_property(op, "__call__", "init_value", arg)
+
 
 for op in [operator.add, operator.mul, operator.and_, operator.xor, operator.or_]:
     (meth, rmeth) = _reflexive_operators[op]
-    register_property(op, '__call__', 'init_value', lambda op, arg: query_property(arg, meth, 'init_value', arg))
+    register_property(
+        op,
+        "__call__",
+        "init_value",
+        lambda op, arg, meth=meth: query_property(arg, meth, "init_value", arg),
+    )
 
 for T in StableNumber:
-    register_property(T, '__add__', 'init_value', lambda a, b: a(False))
-    register_property(T, '__mul__', 'init_value', lambda a, b: a(True))
-    register_property(T, '__and__', 'init_value', lambda a, b: a(True))
-    register_property(T, '__xor__', 'init_value', lambda a, b: a(False))
-    register_property(T, '__or__', 'init_value', lambda a, b: a(False))
+    register_property(T, "__add__", "init_value", lambda a, b: a(False))
+    register_property(T, "__mul__", "init_value", lambda a, b: a(True))
+    register_property(T, "__and__", "init_value", lambda a, b: a(True))
+    register_property(T, "__xor__", "init_value", lambda a, b: a(False))
+    register_property(T, "__or__", "init_value", lambda a, b: a(False))

src/finch/autoschedule/__init__.py

@@ -14,12 +14,12 @@
     Subquery,
     Table,
 )
+from ..symbolic import PostOrderDFS, PostWalk, PreWalk
 from .optimize import (
+    lift_subqueries,
     optimize,
     propagate_map_queries,
-    lift_subqueries,
 )
-from ..symbolic import PostOrderDFS, PostWalk, PreWalk
 
 __all__ = [
     "Aggregate",