docarray · JoanFM · Jun 27, 2023 · Jun 26, 2023 · Jun 26, 2023 · Jun 26, 2023
diff --git a/docarray/documents/text.py b/docarray/documents/text.py
@@ -24,7 +24,7 @@ class TextDoc(BaseDoc):
     from docarray.documents import TextDoc
 
     # use it directly
-    txt_doc = TextDoc(url='http://www.jina.ai/')
+    txt_doc = TextDoc(url='https://www.gutenberg.org/files/1065/1065-0.txt')
     txt_doc.text = txt_doc.url.load()
     # model = MyEmbeddingModel()
     # txt_doc.embedding = model(txt_doc.text)
@@ -51,7 +51,7 @@ class MyText(TextDoc):
         second_embedding: Optional[AnyEmbedding]
 
 
-    txt_doc = MyText(url='http://www.jina.ai/')
+    txt_doc = MyText(url='https://www.gutenberg.org/files/1065/1065-0.txt')
     txt_doc.text = txt_doc.url.load()
     # model = MyEmbeddingModel()
     # txt_doc.embedding = model(txt_doc.text)
@@ -93,8 +93,8 @@ class MultiModalDoc(BaseDoc):
     ```python
     from docarray.documents import TextDoc
 
-    doc = TextDoc(text='This is the main text', url='exampleurl.com')
-    doc2 = TextDoc(text='This is the main text', url='exampleurl.com')
+    doc = TextDoc(text='This is the main text', url='exampleurl.com/file')
+    doc2 = TextDoc(text='This is the main text', url='exampleurl.com/file')
 
     doc == 'This is the main text'  # True
     doc == doc2  # True

diff --git a/docarray/typing/url/any_url.py b/docarray/typing/url/any_url.py
@@ -1,8 +1,9 @@
+import mimetypes
 import os
 import urllib
 import urllib.parse
 import urllib.request
-from typing import TYPE_CHECKING, Any, Optional, Type, TypeVar, Union
+from typing import TYPE_CHECKING, Any, List, Optional, Type, TypeVar, Union
 
 import numpy as np
 from pydantic import AnyUrl as BaseAnyUrl
@@ -20,13 +21,26 @@
 
 T = TypeVar('T', bound='AnyUrl')
 
+mimetypes.init([])
+
 
 @_register_proto(proto_type_name='any_url')
 class AnyUrl(BaseAnyUrl, AbstractType):
     host_required = (
         False  # turn off host requirement to allow passing of local paths as URL
     )
 
+    @classmethod
+    def mime_type(cls) -> str:
+        """Returns the mime type associated with the class."""
+        raise NotImplementedError
+
+    @classmethod
+    def extra_extensions(cls) -> List[str]:
+        """Returns a list of allowed file extensions for the class
+        that are not covered by the mimetypes library."""
+        raise NotImplementedError
+
     def _to_node_protobuf(self) -> 'NodeProto':
         """Convert Document into a NodeProto protobuf message. This function should
         be called when the Document is nested into another Document that need to
@@ -38,6 +52,48 @@ def _to_node_protobuf(self) -> 'NodeProto':
 
         return NodeProto(text=str(self), type=self._proto_type_name)
 
+    @staticmethod
+    def _get_url_extension(url: str) -> str:
+        """
+        Extracts and returns the file extension from a given URL.
+        If no file extension is present, the function returns an empty string.
+
+
+        :param url: The URL to extract the file extension from.
+        :return: The file extension without the period, if one exists,
+            otherwise an empty string.
+        """
+
+        parsed_url = urllib.parse.urlparse(url)
+        ext = os.path.splitext(parsed_url.path)[1]
+        ext = ext[1:] if ext.startswith('.') else ext
+        return ext
+
+    @classmethod
+    def is_extension_allowed(cls, value: Any) -> bool:
+        """
+        Check if the file extension of the URL is allowed for this class.
+        First, it guesses the mime type of the file. If it fails to detect the
+        mime type, it then checks the extra file extensions.
+        Note: This method assumes that any URL without an extension is valid.
+
+        :param value: The URL or file path.
+        :return: True if the extension is allowed, False otherwise
+        """
+        if cls is AnyUrl:
+            return True
+
+        url_parts = value.split('?')
+        extension = cls._get_url_extension(value)
+        if not extension:
+            return True
+
+        mimetype, _ = mimetypes.guess_type(url_parts[0])
+        if mimetype and mimetype.startswith(cls.mime_type()):
+            return True
+
+        return extension in cls.extra_extensions()
+
     @classmethod
     def validate(
         cls: Type[T],
@@ -61,10 +117,12 @@ def validate(
 
         url = super().validate(abs_path, field, config)  # basic url validation
 
-        if input_is_relative_path:
-            return cls(str(value), scheme=None)
-        else:
-            return cls(str(url), scheme=None)
+        if not cls.is_extension_allowed(value):
+            raise ValueError(
+                f"The file '{value}' is not in a valid format for class '{cls.__name__}'."
+            )
+
+        return cls(str(value if input_is_relative_path else url), scheme=None)
 
     @classmethod
     def validate_parts(cls, parts: 'Parts', validate_port: bool = True) -> 'Parts':

diff --git a/docarray/typing/url/audio_url.py b/docarray/typing/url/audio_url.py
@@ -1,10 +1,11 @@
 import warnings
-from typing import Optional, Tuple, TypeVar
+from typing import List, Optional, Tuple, TypeVar
 
 from docarray.typing import AudioNdArray
 from docarray.typing.bytes.audio_bytes import AudioBytes
 from docarray.typing.proto_register import _register_proto
 from docarray.typing.url.any_url import AnyUrl
+from docarray.typing.url.mimetypes import AUDIO_MIMETYPE
 from docarray.utils._internal.misc import is_notebook
 
 T = TypeVar('T', bound='AudioUrl')
@@ -17,6 +18,18 @@ class AudioUrl(AnyUrl):
     Can be remote (web) URL, or a local file path.
     """
 
+    @classmethod
+    def mime_type(cls) -> str:
+        return AUDIO_MIMETYPE
+
+    @classmethod
+    def extra_extensions(cls) -> List[str]:
+        """
+        Returns a list of additional file extensions that are valid for this class
+        but cannot be identified by the mimetypes library.
+        """
+        return []
+
     def load(self: T) -> Tuple[AudioNdArray, int]:
         """
         Load the data from the url into an [`AudioNdArray`][docarray.typing.AudioNdArray]

diff --git a/docarray/typing/url/image_url.py b/docarray/typing/url/image_url.py
@@ -1,10 +1,11 @@
 import warnings
-from typing import TYPE_CHECKING, Optional, Tuple, TypeVar
+from typing import TYPE_CHECKING, List, Optional, Tuple, TypeVar
 
 from docarray.typing import ImageBytes
 from docarray.typing.proto_register import _register_proto
 from docarray.typing.tensor.image import ImageNdArray
 from docarray.typing.url.any_url import AnyUrl
+from docarray.typing.url.mimetypes import IMAGE_MIMETYPE
 from docarray.utils._internal.misc import is_notebook
 
 if TYPE_CHECKING:
@@ -20,6 +21,18 @@ class ImageUrl(AnyUrl):
     Can be remote (web) URL, or a local file path.
     """
 
+    @classmethod
+    def mime_type(cls) -> str:
+        return IMAGE_MIMETYPE
+
+    @classmethod
+    def extra_extensions(cls) -> List[str]:
+        """
+        Returns a list of additional file extensions that are valid for this class
+        but cannot be identified by the mimetypes library.
+        """
+        return []
+
     def load_pil(self, timeout: Optional[float] = None) -> 'PILImage.Image':
         """
         Load the image from the bytes into a `PIL.Image.Image` instance

diff --git a/docarray/typing/url/mimetypes.py b/docarray/typing/url/mimetypes.py
@@ -0,0 +1,94 @@
+TEXT_MIMETYPE = 'text'
+AUDIO_MIMETYPE = 'audio'
+IMAGE_MIMETYPE = 'image'
+OBJ_MIMETYPE = 'application/x-tgif'
+VIDEO_MIMETYPE = 'video'
+
+MESH_EXTRA_EXTENSIONS = [
+    '3ds',
+    '3mf',
+    'ac',
+    'ac3d',
+    'amf',
+    'assimp',
+    'bvh',
+    'cob',
+    'collada',
+    'ctm',
+    'dxf',
+    'e57',
+    'fbx',
+    'gltf',
+    'glb',
+    'ifc',
+    'lwo',
+    'lws',
+    'lxo',
+    'md2',
+    'md3',
+    'md5',
+    'mdc',
+    'm3d',
+    'mdl',
+    'ms3d',
+    'nff',
+    'obj',
+    'off',
+    'pcd',
+    'pod',
+    'pmd',
+    'pmx',
+    'ply',
+    'q3o',
+    'q3s',
+    'raw',
+    'sib',
+    'smd',
+    'stl',
+    'ter',
+    'terragen',
+    'vtk',
+    'vrml',
+    'x3d',
+    'xaml',
+    'xgl',
+    'xml',
+    'xyz',
+    'zgl',
+    'vta',
+]
+
+TEXT_EXTRA_EXTENSIONS = ['md', 'log']
+
+POINT_CLOUD_EXTRA_EXTENSIONS = [
+    'ascii',
+    'bin',
+    'b3dm',
+    'bpf',
+    'dp',
+    'dxf',
+    'e57',
+    'fls',
+    'fls',
+    'glb',
+    'ply',
+    'gpf',
+    'las',
+    'obj',
+    'osgb',
+    'pcap',
+    'pcd',
+    'pdal',
+    'pfm',
+    'ply',
+    'ply2',
+    'pod',
+    'pods',
+    'pnts',
+    'ptg',
+    'ptx',
+    'pts',
+    'rcp',
+    'xyz',
+    'zfs',
+]
diff --git a/docarray/typing/url/text_url.py b/docarray/typing/url/text_url.py
@@ -1,7 +1,8 @@
-from typing import Optional, TypeVar
+from typing import List, Optional, TypeVar
 
 from docarray.typing.proto_register import _register_proto
 from docarray.typing.url.any_url import AnyUrl
+from docarray.typing.url.mimetypes import TEXT_EXTRA_EXTENSIONS, TEXT_MIMETYPE
 
 T = TypeVar('T', bound='TextUrl')
 
@@ -13,6 +14,18 @@ class TextUrl(AnyUrl):
     Can be remote (web) URL, or a local file path.
     """
 
+    @classmethod
+    def mime_type(cls) -> str:
+        return TEXT_MIMETYPE
+
+    @classmethod
+    def extra_extensions(cls) -> List[str]:
+        """
+        Returns a list of additional file extensions that are valid for this class
+        but cannot be identified by the mimetypes library.
+        """
+        return TEXT_EXTRA_EXTENSIONS
+
     def load(self, charset: str = 'utf-8', timeout: Optional[float] = None) -> str:
         """
         Load the text file into a string.

diff --git a/docarray/typing/url/url_3d/mesh_url.py b/docarray/typing/url/url_3d/mesh_url.py
@@ -1,10 +1,11 @@
-from typing import TYPE_CHECKING, Any, Dict, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypeVar
 
 import numpy as np
 from pydantic import parse_obj_as
 
 from docarray.typing.proto_register import _register_proto
 from docarray.typing.tensor.ndarray import NdArray
+from docarray.typing.url.mimetypes import MESH_EXTRA_EXTENSIONS
 from docarray.typing.url.url_3d.url_3d import Url3D
 
 if TYPE_CHECKING:
@@ -20,6 +21,14 @@ class Mesh3DUrl(Url3D):
     Can be remote (web) URL, or a local file path.
     """
 
+    @classmethod
+    def extra_extensions(cls) -> List[str]:
+        """
+        Returns a list of additional file extensions that are valid for this class
+        but cannot be identified by the mimetypes library.
+        """
+        return MESH_EXTRA_EXTENSIONS
+
     def load(
         self: T,
         skip_materials: bool = True,

diff --git a/docarray/typing/url/url_3d/point_cloud_url.py b/docarray/typing/url/url_3d/point_cloud_url.py
@@ -1,10 +1,11 @@
-from typing import TYPE_CHECKING, Any, Dict, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypeVar
 
 import numpy as np
 from pydantic import parse_obj_as
 
 from docarray.typing.proto_register import _register_proto
 from docarray.typing.tensor.ndarray import NdArray
+from docarray.typing.url.mimetypes import POINT_CLOUD_EXTRA_EXTENSIONS
 from docarray.typing.url.url_3d.url_3d import Url3D
 
 if TYPE_CHECKING:
@@ -21,6 +22,14 @@ class PointCloud3DUrl(Url3D):
     Can be remote (web) URL, or a local file path.
     """
 
+    @classmethod
+    def extra_extensions(cls) -> List[str]:
+        """
+        Returns a list of additional file extensions that are valid for this class
+        but cannot be identified by the mimetypes library.
+        """
+        return POINT_CLOUD_EXTRA_EXTENSIONS
+
     def load(
         self: T,
         samples: int,

diff --git a/docarray/typing/url/url_3d/url_3d.py b/docarray/typing/url/url_3d/url_3d.py
@@ -3,6 +3,7 @@
 
 from docarray.typing.proto_register import _register_proto
 from docarray.typing.url.any_url import AnyUrl
+from docarray.typing.url.mimetypes import OBJ_MIMETYPE
 from docarray.utils._internal.misc import import_library
 
 if TYPE_CHECKING:
@@ -18,6 +19,10 @@ class Url3D(AnyUrl, ABC):
     Can be remote (web) URL, or a local file path.
     """
 
+    @classmethod
+    def mime_type(cls) -> str:
+        return OBJ_MIMETYPE
+
     def _load_trimesh_instance(
         self: T,
         force: Optional[str] = None,