[python][typehint] Simplify find_paths_if and add typehints to _utils (#6497)

A while ago I brought up the idea of adding typehints to triton. I
started following through on that ambition in
triton-lang/triton#6467, and want to continue
the endeavor here.

I plan to keep taking the 'leafiest' modules and adding typehinting to
them. In order to validate types locally I'm using mypy with an lsp
connection to neovim. One there's a good amount of typehinting I can go
through and add `ignore` statements everywhere there's an error so that
we can get typechecking in our pre-commit. For now I'll address type
errors ad-hoc as I see them.

Author: danzimm

python/triton/_utils.py

@@ -1,35 +1,37 @@
+from __future__ import annotations
+
 from functools import reduce
+from typing import Any, Callable, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    from .language import core
+    IterableType = Union[list[Any], tuple[Any, ...], core.tuple, core.tuple_type]
+    ObjPath = tuple[int, ...]
 
 
-def get_iterable_path(iterable, path):
-    return reduce(lambda a, idx: a[idx], path, iterable)
+def get_iterable_path(iterable: IterableType, path: ObjPath) -> Any:
+    return reduce(lambda a, idx: a[idx], path, iterable)  # type: ignore[index]
 
 
-def set_iterable_path(iterable, path, val):
+def set_iterable_path(iterable: IterableType, path: tuple[int, ...], val: Any):
+    assert len(path) != 0
     prev = iterable if len(path) == 1 else get_iterable_path(iterable, path[:-1])
-    prev[path[-1]] = val
+    prev[path[-1]] = val  # type: ignore[index]
 
 
-def find_paths_if(iterable, pred):
+def find_paths_if(iterable: Union[IterableType, Any], pred: Callable[[ObjPath, Any], bool]) -> list[ObjPath]:
     from .language import core
-    is_iterable = lambda x: isinstance(x, (list, tuple, core.tuple, core.tuple_type))
-    ret = dict()
+    is_iterable: Callable[[Any], bool] = lambda x: isinstance(x, (list, tuple, core.tuple, core.tuple_type))
+    # We need to use dict so that ordering is maintained, while set doesn't guarantee order
+    ret: dict[ObjPath, None] = {}
 
-    def _impl(current, path):
-        path = (path[0], ) if len(path) == 1 else tuple(path)
+    def _impl(path: tuple[int, ...], current: Any):
         if is_iterable(current):
             for idx, item in enumerate(current):
-                _impl(item, path + (idx, ))
+                _impl((*path, idx), item)
         elif pred(path, current):
-            if len(path) == 1:
-                ret[(path[0], )] = None
-            else:
-                ret[tuple(path)] = None
-
-    if is_iterable(iterable):
-        _impl(iterable, [])
-    elif pred(list(), iterable):
-        ret = {tuple(): None}
-    else:
-        ret = dict()
+            ret[path] = None
+
+    _impl((), iterable)
+
     return list(ret.keys())