turn HasDataType and HasShape into Protocol\s

Author: markusschmaus

aesara/graph/type.py

@@ -1,7 +1,7 @@
 from abc import abstractmethod
 from typing import Any, Generic, Optional, Text, Tuple, TypeVar, Union
 
-from typing_extensions import TypeAlias
+from typing_extensions import Protocol, TypeAlias, runtime_checkable
 
 from aesara.graph import utils
 from aesara.graph.basic import Constant, Variable
@@ -262,14 +262,22 @@ def values_eq_approx(cls, a: D, b: D) -> bool:
         return cls.values_eq(a, b)
 
 
-class HasDataType:
-    """A mixin for a type that has a :attr:`dtype` attribute."""
+DataType = str
 
-    dtype: str
 
+@runtime_checkable
+class HasDataType(Protocol):
+    """A protocol matching any class with :attr:`dtype` attribute."""
 
-class HasShape:
-    """A mixin for a type that has :attr:`shape` and :attr:`ndim` attributes."""
+    dtype: DataType
+
+
+ShapeType = Tuple[Optional[int], ...]
+
+
+@runtime_checkable
+class HasShape(Protocol):
+    """A protocol matching any class that has :attr:`shape` and :attr:`ndim` attributes."""
 
     ndim: int
-    shape: Tuple[Optional[int], ...]
+    shape: ShapeType

aesara/link/c/cmodule.py

@@ -2441,7 +2441,7 @@ def linking_patch(lib_dirs: List[str], libs: List[str]) -> List[str]:
         if sys.platform != "win32":
             return [f"-l{l}" for l in libs]
 
-        def sort_key(lib):  # type: ignore
+        def sort_key(lib):
             name, *numbers, extension = lib.split(".")
             return (extension == "dll", tuple(map(int, numbers)))
 

aesara/scalar/basic.py

@@ -28,7 +28,7 @@
 from aesara.graph.basic import Apply, Constant, Variable, clone, list_of_nodes
 from aesara.graph.fg import FunctionGraph
 from aesara.graph.rewriting.basic import MergeOptimizer
-from aesara.graph.type import HasDataType, HasShape
+from aesara.graph.type import DataType
 from aesara.graph.utils import MetaObject, MethodNotDefined
 from aesara.link.c.op import COp
 from aesara.link.c.type import CType
@@ -268,7 +268,7 @@ def convert(x, dtype=None):
     return x_
 
 
-class ScalarType(CType, HasDataType, HasShape):
+class ScalarType(CType):
 
     """
     Internal class, should not be used by clients.
@@ -284,6 +284,7 @@ class ScalarType(CType, HasDataType, HasShape):
     __props__ = ("dtype",)
     ndim = 0
     shape = ()
+    dtype: DataType
 
     def __init__(self, dtype):
         if isinstance(dtype, str) and dtype == "floatX":

aesara/sparse/type.py

@@ -7,7 +7,6 @@
 import aesara
 from aesara import scalar as aes
 from aesara.graph.basic import Variable
-from aesara.graph.type import HasDataType
 from aesara.tensor.type import DenseTensorType, TensorType
 
 
@@ -33,7 +32,7 @@ def _is_sparse(x):
     return isinstance(x, scipy.sparse.spmatrix)
 
 
-class SparseTensorType(TensorType, HasDataType):
+class SparseTensorType(TensorType):
     """A `Type` for sparse tensors.
 
     Notes

aesara/tensor/type.py

@@ -8,7 +8,7 @@
 from aesara import scalar as aes
 from aesara.configdefaults import config
 from aesara.graph.basic import Variable
-from aesara.graph.type import HasDataType, HasShape
+from aesara.graph.type import DataType, ShapeType
 from aesara.graph.utils import MetaType
 from aesara.link.c.type import CType
 from aesara.misc.safe_asarray import _asarray
@@ -48,11 +48,15 @@
 }
 
 
-class TensorType(CType[np.ndarray], HasDataType, HasShape):
+class TensorType(CType[np.ndarray]):
     r"""Symbolic `Type` representing `numpy.ndarray`\s."""
 
     __props__: Tuple[str, ...] = ("dtype", "shape")
 
+    ndim: int
+    shape: ShapeType
+    dtype: DataType
+
     dtype_specs_map = dtype_specs_map
     context_name = "cpu"
     filter_checks_isfinite = False