fix: load audio url to audio ndarray instead of np ndarray

Signed-off-by: anna-charlotte <[email protected]>
docarray · JoanFM · Jan 3, 2023 · Dec 14, 2022 · Dec 14, 2022 · Dec 15, 2022
commit 83ece31fce90d8f2ff6c1e58d9e1b93fd9f9916e
diff --git a/docarray/typing/tensor/audio/audio_ndarray.py b/docarray/typing/tensor/audio/audio_ndarray.py
@@ -2,7 +2,8 @@
 
 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
+
+MAX_INT_16 = 2**15
 
 T = TypeVar('T', bound='AudioNdArray')
 

diff --git a/docarray/typing/tensor/audio/audio_torch_tensor.py b/docarray/typing/tensor/audio/audio_torch_tensor.py
@@ -1,8 +1,8 @@
 from typing import TypeVar
 
 from docarray.typing.tensor.audio.abstract_audio_tensor import AbstractAudioTensor
+from docarray.typing.tensor.audio.audio_ndarray import MAX_INT_16
 from docarray.typing.tensor.torch_tensor import TorchTensor, metaTorchAndNode
-from docarray.typing.url.audio_url import MAX_INT_16
 
 T = TypeVar('T', bound='AudioTorchTensor')
 

diff --git a/docarray/typing/url/audio_url.py b/docarray/typing/url/audio_url.py
@@ -2,7 +2,9 @@
 from typing import TYPE_CHECKING, Any, Type, TypeVar, Union
 
 import numpy as np
+from pydantic import parse_obj_as
 
+from docarray.typing.tensor.audio.audio_ndarray import MAX_INT_16, AudioNdArray
 from docarray.typing.url.any_url import AnyUrl
 
 if TYPE_CHECKING:
@@ -14,7 +16,6 @@
 T = TypeVar('T', bound='AudioUrl')
 
 AUDIO_FILE_FORMATS = ['wav']
-MAX_INT_16 = 2**15
 
 
 class AudioUrl(AnyUrl):
@@ -50,9 +51,9 @@ def validate(
             )
         return cls(str(url), scheme=None)
 
-    def load(self: T) -> np.ndarray:
+    def load(self: T) -> AudioNdArray:
         """
-        Load the data from the url into a numpy.ndarray.
+        Load the data from the url into an AudioNdArray.
 
         EXAMPLE USAGE
 
@@ -66,14 +67,14 @@ def load(self: T) -> np.ndarray:
 
             class MyDoc(Document):
                 audio_url: AudioUrl
+                audio_tensor: AudioNdArray
 
 
             doc = MyDoc(audio_url="toydata/hello.wav")
+            doc.audio_tensor = doc.audio_url.load()
+            assert isinstance(doc.audio_tensor, np.ndarray)
 
-            audio_tensor = doc.audio_url.load()
-            assert isinstance(audio_tensor, np.ndarray)
-
-        :return: np.ndarray representing the audio file content
+        :return: AudioNdArray representing the audio file content
         """
         import io
 
@@ -109,6 +110,6 @@ class MyDoc(Document):
                 audio_stereo[:, 0] = audio_norm[range(0, len(audio_norm), 2)]
                 audio_stereo[:, 1] = audio_norm[range(1, len(audio_norm), 2)]
 
-                return audio_stereo
+                return parse_obj_as(AudioNdArray, audio_stereo)
             else:
-                return audio_norm
+                return parse_obj_as(AudioNdArray, audio_norm)
diff --git a/tests/units/typing/url/test_audio_url.py b/tests/units/typing/url/test_audio_url.py
@@ -1,9 +1,12 @@
+from typing import Optional
+
 import numpy as np
 import pytest
 from pydantic.tools import parse_obj_as, schema_json_of
 
+from docarray import Document
 from docarray.document.io.json import orjson_dumps
-from docarray.typing import AudioUrl
+from docarray.typing import AudioNdArray, AudioTorchTensor, AudioUrl
 from tests import TOYDATA_DIR
 
 AUDIO_FILES = [
@@ -23,6 +26,25 @@ def test_audio_url(file_url):
     uri = parse_obj_as(AudioUrl, file_url)
     tensor = uri.load()
     assert isinstance(tensor, np.ndarray)
+    assert isinstance(tensor, AudioNdArray)
+
+
+@pytest.mark.slow
+@pytest.mark.internet
+@pytest.mark.parametrize(
+    'file_url',
+    [*AUDIO_FILES, REMOTE_AUDIO_FILE],
+)
+def test_load_audio_url_to_audio_torch_tensor_field(file_url):
+    class MyAudioDoc(Document):
+        audio_url: AudioUrl
+        tensor: Optional[AudioTorchTensor]
+
+    doc = MyAudioDoc(audio_url=file_url)
+    doc.tensor = doc.audio_url.load()
+
+    assert isinstance(doc.tensor, np.ndarray)
+    assert isinstance(doc.tensor, AudioNdArray)
 
 
 @pytest.mark.slow