fix: introduce compbackendinterface

Signed-off-by: anna-charlotte <[email protected]>
docarray · samsja · Jan 17, 2023 · Jan 3, 2023 · Jan 3, 2023 · Jan 4, 2023
commit d50ae67fcb2b9462304a1b99d8fd622b431680e6
diff --git a/docarray/typing/tensor/video/video_ndarray.py b/docarray/typing/tensor/video/video_ndarray.py
@@ -31,4 +31,4 @@ def validate(
         config: 'BaseConfig',
     ) -> T:
         tensor = super().validate(value=value, field=field, config=config)
-        return VideoTensorMixin.validate_shape(cls, value=tensor)
+        return cls.validate_shape(value=tensor)
diff --git a/docarray/typing/tensor/video/video_tensor_mixin.py b/docarray/typing/tensor/video/video_tensor_mixin.py
@@ -1,27 +1,33 @@
+import abc
 from typing import TYPE_CHECKING, BinaryIO, Optional, Type, TypeVar, Union
 
 import numpy as np
 
 from docarray.typing.tensor.audio.audio_tensor import AudioTensor
 
 if TYPE_CHECKING:
-    from docarray.typing import VideoNdArray, VideoTorchTensor
-
+    from docarray.typing.tensor.abstract_tensor import AbstractTensor
 
 T = TypeVar('T', bound='VideoTensorMixin')
+TT = TypeVar('TT', bound='AbstractTensor')
 
 
-class VideoTensorMixin:
+class CompBackendInterface(abc.ABC):
     @staticmethod
-    def validate_shape(
-        cls: Union[Type['VideoTorchTensor'], Type['VideoNdArray']], value: 'T'
-    ) -> 'T':
-        comp_backend = cls.get_comp_backend()
-        shape = comp_backend.shape(value)  # type: ignore
-        if comp_backend.n_dim(value) not in [3, 4] or shape[-1] != 3:  # type: ignore
+    @abc.abstractmethod
+    def get_comp_backend():
+        """The computational backend compatible with this tensor type."""
+        ...
+
+
+class VideoTensorMixin(CompBackendInterface, abc.ABC):
+    @classmethod
+    def validate_shape(cls: Type['T'], value: 'T') -> 'T':
+        comp_be = cls.get_comp_backend()
+        if comp_be.n_dim(value) not in [3, 4] or comp_be.shape(value)[-1] != 3:
             raise ValueError(
                 f'Expects tensor with 3 or 4 dimensions and the last dimension equal '
-                f'to 3, but received {shape}.'
+                f'to 3, but received {comp_be.shape(value)}.'
             )
         else:
             return value

diff --git a/docarray/typing/tensor/video/video_torch_tensor.py b/docarray/typing/tensor/video/video_torch_tensor.py
@@ -31,4 +31,4 @@ def validate(
         config: 'BaseConfig',
     ) -> T:
         tensor = super().validate(value=value, field=field, config=config)
-        return VideoTensorMixin.validate_shape(cls, value=tensor)
+        return cls.validate_shape(value=tensor)