fix(type hints): improvements to type hints in ibis.expr

ibis/expr/analysis.py

@@ -30,7 +30,7 @@
 # compilation later
 
 
-def sub_for(node: ops.Node, substitutions: Mapping[ops.node, ops.Node]) -> ops.Node:
+def sub_for(node: ops.Node, substitutions: Mapping[ops.Node, ops.Node]) -> ops.Node:
     """Substitute operations in `node` with nodes in `substitutions`.
 
     Parameters

ibis/expr/deferred.py

@@ -174,6 +174,8 @@ def __neg__(self) -> Deferred:
 
 class DeferredAttr(Deferred):
     __slots__ = ("_value", "_attr")
+    _value: Any
+    _attr: str
 
     def __init__(self, value: Any, attr: str) -> None:
         self._value = value
@@ -189,6 +191,8 @@ def _resolve(self, param: Any) -> Any:
 
 class DeferredItem(Deferred):
     __slots__ = ("_value", "_key")
+    _value: Any
+    _key: Any
 
     def __init__(self, value: Any, key: Any) -> None:
         self._value = value
@@ -204,6 +208,9 @@ def _resolve(self, param: Any) -> Any:
 
 class DeferredCall(Deferred):
     __slots__ = ("_func", "_args", "_kwargs")
+    _func: Callable
+    _args: tuple[Any, ...]
+    _kwargs: dict[str, Any]
 
     def __init__(self, func: Any, *args: Any, **kwargs: Any) -> None:
         self._func = func
@@ -224,6 +231,9 @@ def _resolve(self, param: Any) -> Any:
 
 class DeferredBinaryOp(Deferred):
     __slots__ = ("_symbol", "_left", "_right")
+    _symbol: str
+    _left: Any
+    _right: Any
 
     def __init__(self, symbol: str, left: Any, right: Any) -> None:
         self._symbol = symbol
@@ -241,6 +251,8 @@ def _resolve(self, param: Any) -> Any:
 
 class DeferredUnaryOp(Deferred):
     __slots__ = ("_symbol", "_value")
+    _symbol: str
+    _value: Any
 
     def __init__(self, symbol: str, value: Any) -> None:
         self._symbol = symbol

ibis/expr/operations/relations.py

@@ -110,6 +110,8 @@ class SQLQueryResult(TableNode):
 
 class TableProxy(Immutable):
     __slots__ = ("_data", "_hash")
+    _data: Any
+    _hash: int
 
     def __init__(self, data) -> None:
         object.__setattr__(self, "_data", data)

ibis/expr/operations/tests/conftest.py

@@ -2,6 +2,6 @@
 
 import sys
 
-collect_ignore = []
+collect_ignore: list[str] = []
 if sys.version_info < (3, 10):
     collect_ignore.append("test_core_py310.py")

ibis/expr/types/core.py

@@ -6,19 +6,22 @@
 from typing import TYPE_CHECKING, Any, Mapping, NoReturn, Tuple
 
 from public import public
+from rich.jupyter import JupyterMixin
 
 import ibis.expr.operations as ops
+from ibis.common.annotations import ValidationError
 from ibis.common.exceptions import IbisError, TranslationError
 from ibis.common.grounds import Immutable
+from ibis.common.patterns import Coercible, CoercionError
 from ibis.config import _default_backend, options
 from ibis.util import experimental
-from ibis.common.annotations import ValidationError
-from rich.jupyter import JupyterMixin
-from ibis.common.patterns import Coercible, CoercionError
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     import pandas as pd
     import pyarrow as pa
+    import torch
 
     import ibis.expr.types as ir
     from ibis.backends.base import BaseBackend
@@ -41,6 +44,7 @@ class Expr(Immutable, Coercible):
     """Base expression class."""
 
     __slots__ = ("_arg",)
+    _arg: ops.Node
 
     def __init__(self, arg: ops.Node) -> None:
         object.__setattr__(self, "_arg", arg)

ibis/expr/types/generic.py

@@ -1245,7 +1245,7 @@ def topk(
         self,
         k: int,
         by: ir.Value | None = None,
-    ) -> ir.TopK:
+    ) -> ir.Table:
         """Return a "top k" expression.
 
         Parameters

ibis/expr/types/structs.py

@@ -338,7 +338,7 @@ def lift(self) -> ir.Table:
         table = an.find_first_base_table(self.op()).to_expr()
         return table[[self[name] for name in self.names]]
 
-    def destructure(self) -> list[ir.ValueExpr]:
+    def destructure(self) -> list[ir.Value]:
         """Destructure a ``StructValue`` into the corresponding struct fields.
 
         When assigned, a destruct value will be destructured and assigned to