Fix typing issues in aesara.graph.type

aesara/graph/type.py

@@ -223,7 +223,7 @@ class Type(MetaObject):
         return utils.add_tag_trace(self.make_variable(name))
 
     @classmethod
-    def values_eq(cls, a: Any, b: Any) -> bool:
+    def values_eq(cls, a: "Type", b: "Type") -> bool:
         """Return ``True`` if `a` and `b` can be considered exactly equal.
 
         `a` and `b` are assumed to be valid values of this `Type`.

aesara/graph/utils.py

@@ -3,7 +3,13 @@ import sys
 import traceback
 from abc import ABCMeta
 from io import StringIO
-from typing import List, Optional, Sequence, Tuple
+from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple, TypeVar, Union
+
+
+if TYPE_CHECKING:
+    from aesara.graph.basic import Apply, Variable
+
+T = TypeVar("T", bound=Union["Apply", "Variable"])
 
 
 def simple_extract_stack(
@@ -67,7 +73,7 @@ def simple_extract_stack(
     return trace
 
 
-def add_tag_trace(thing, user_line: Optional[int] = None):
+def add_tag_trace(thing: T, user_line: Optional[int] = None) -> T:
     """Add tag.trace to a node or variable.
 
     The argument is returned after being affected (inplace).

setup.cfg

@@ -123,10 +123,6 @@ check_untyped_defs = False
 ignore_errors = True
 check_untyped_defs = False
 
-[mypy-aesara.graph.type]
-ignore_errors = True
-check_untyped_defs = False
-
 [mypy-aesara.graph.optdb]
 ignore_errors = True
 check_untyped_defs = False