fix: mypy checks

Signed-off-by: anna-charlotte <[email protected]>
docarray · JoanFM · Jan 3, 2023 · Dec 14, 2022 · Dec 14, 2022 · Dec 15, 2022
commit eecca41a00f885c9e9e76370d9e421968ca024cd
diff --git a/docarray/typing/tensor/audio/abstract_audio_tensor.py b/docarray/typing/tensor/audio/abstract_audio_tensor.py
@@ -1,15 +1,11 @@
 import wave
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, BinaryIO, TypeVar, Union
+from typing import BinaryIO, TypeVar, Union
 
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
-from docarray.typing.url.audio_url import MAX_INT_16
 
 T = TypeVar('T', bound='AbstractAudioTensor')
 
-if TYPE_CHECKING:
-    from docarray.proto import NodeProto
-
 
 class AbstractAudioTensor(AbstractTensor, ABC):
 
@@ -18,7 +14,7 @@ class AbstractAudioTensor(AbstractTensor, ABC):
     @abstractmethod
     def to_audio_bytes(self):
         """
-        Convert tensor to bytes.
+        Convert audio tensor to bytes.
         """
         raise NotImplementedError
 
@@ -36,23 +32,14 @@ def save_to_wav_file(
         :param sample_rate: sampling frequency
         :param sample_width: sample width in bytes
         """
-        n_channels = 2 if self.ndim() > 1 else 1
+        n_channels = 2 if self.n_dim() > 1 else 1
 
         with wave.open(file_path, 'w') as f:
             f.setnchannels(n_channels)
             f.setsampwidth(sample_width)
             f.setframerate(sample_rate)
-            f.writeframes((self * MAX_INT_16).to_audio_bytes())
+            f.writeframes(self.to_audio_bytes())
 
-    def _to_node_protobuf(self: T) -> 'NodeProto':
-        """
-        Convert itself into a NodeProto protobuf message. This function should
-        be called when the Document is nested into another Document that need to be
-        converted into a protobuf
-        :param field: field in which to store the content in the node proto
-        :return: the nested item protobuf message
-        """
-        from docarray.proto import NodeProto
-
-        nd_proto = self.to_protobuf()
-        return NodeProto(**{self.TENSOR_FIELD_NAME: nd_proto})
+    @abstractmethod
+    def n_dim(self) -> int:
+        ...
diff --git a/docarray/typing/tensor/audio/audio_ndarray.py b/docarray/typing/tensor/audio/audio_ndarray.py
@@ -1,10 +1,14 @@
-from typing import TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
 from docarray.typing import NdArray
 from docarray.typing.tensor.audio.abstract_audio_tensor import AbstractAudioTensor
+from docarray.typing.url.audio_url import MAX_INT_16
 
 T = TypeVar('T', bound='AudioNdArray')
 
+if TYPE_CHECKING:
+    from docarray.proto import NodeProto
+
 
 class AudioNdArray(AbstractAudioTensor, NdArray):
     """
@@ -52,7 +56,22 @@ class MyAudioDoc(Document):
 
     TENSOR_FIELD_NAME = 'audio_ndarray'
 
-    def to_audio_bytes(self):
-        tensor = self.astype('<h')
+    def _to_node_protobuf(self: T, field: str = TENSOR_FIELD_NAME) -> 'NodeProto':
+        """
+        Convert itself into a NodeProto protobuf message. This function should
+        be called when the Document is nested into another Document that need to be
+        converted into a protobuf
+        :param field: field in which to store the content in the node proto
+        :return: the nested item protobuf message
+        """
+        from docarray.proto import NodeProto
+
+        nd_proto = self.to_protobuf()
+        return NodeProto(**{field: nd_proto})
 
+    def to_audio_bytes(self):
+        tensor = (self * MAX_INT_16).astype('<h')
         return tensor.tobytes()
+
+    def n_dim(self) -> int:
+        return self.ndim
diff --git a/docarray/typing/tensor/audio/audio_torch_tensor.py b/docarray/typing/tensor/audio/audio_torch_tensor.py
@@ -1,10 +1,14 @@
-from typing import TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
 from docarray.typing.tensor.audio.abstract_audio_tensor import AbstractAudioTensor
 from docarray.typing.tensor.torch_tensor import TorchTensor, metaTorchAndNode
+from docarray.typing.url.audio_url import MAX_INT_16
 
 T = TypeVar('T', bound='AudioTorchTensor')
 
+if TYPE_CHECKING:
+    from docarray.proto import NodeProto
+
 
 class AudioTorchTensor(AbstractAudioTensor, TorchTensor, metaclass=metaTorchAndNode):
     """
@@ -51,8 +55,24 @@ class MyAudioDoc(Document):
 
     TENSOR_FIELD_NAME = 'audio_torch_tensor'
 
+    def _to_node_protobuf(self: T, field: str = TENSOR_FIELD_NAME) -> 'NodeProto':
+        """
+        Convert itself into a NodeProto protobuf message. This function should
+        be called when the Document is nested into another Document that need to be
+        converted into a protobuf
+        :param field: field in which to store the content in the node proto
+        :return: the nested item protobuf message
+        """
+        from docarray.proto import NodeProto
+
+        nd_proto = self.to_protobuf()
+        return NodeProto(**{field: nd_proto})
+
     def to_audio_bytes(self):
         import torch
 
-        tensor = self.to(dtype=torch.int16)
+        tensor = (self * MAX_INT_16).to(dtype=torch.int16)
         return tensor.cpu().detach().numpy().tobytes()
+
+    def n_dim(self) -> int:
+        return self.ndim