Url

@_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
        be converted into a protobuf

        :return: the nested item protobuf message
        """
        from docarray.proto import 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],
        value: Union[T, np.ndarray, Any],
        field: 'ModelField',
        config: 'BaseConfig',
    ) -> T:
        import os

        abs_path: Union[T, np.ndarray, Any]
        if (
            isinstance(value, str)
            and not value.startswith('http')
            and not os.path.isabs(value)
        ):
            input_is_relative_path = True
            abs_path = os.path.abspath(value)
        else:
            input_is_relative_path = False
            abs_path = value

        url = super().validate(abs_path, field, config)  # basic url validation

        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':
        """
        A method used to validate parts of a URL.
        Our URLs should be able to function both in local and remote settings.
        Therefore, we allow missing `scheme`, making it possible to pass a file
        path without prefix.
        If `scheme` is missing, we assume it is a local file path.
        """
        scheme = parts['scheme']
        if scheme is None:
            # allow missing scheme, unlike pydantic
            pass

        elif cls.allowed_schemes and scheme.lower() not in cls.allowed_schemes:
            raise errors.UrlSchemePermittedError(set(cls.allowed_schemes))

        if validate_port:
            cls._validate_port(parts['port'])

        user = parts['user']
        if cls.user_required and user is None:
            raise errors.UrlUserInfoError()

        return parts

    @classmethod
    def build(
        cls,
        *,
        scheme: str,
        user: Optional[str] = None,
        password: Optional[str] = None,
        host: str,
        port: Optional[str] = None,
        path: Optional[str] = None,
        query: Optional[str] = None,
        fragment: Optional[str] = None,
        **_kwargs: str,
    ) -> str:
        """
        Build a URL from its parts.
        The only difference from the pydantic implementation is that we allow
        missing `scheme`, making it possible to pass a file path without prefix.
        """

        # allow missing scheme, unlike pydantic
        scheme_ = scheme if scheme is not None else ''
        url = super().build(
            scheme=scheme_,
            user=user,
            password=password,
            host=host,
            port=port,
            path=path,
            query=query,
            fragment=fragment,
            **_kwargs,
        )
        if scheme is None and url.startswith('://'):
            # remove the `://` prefix, since scheme is missing
            url = url[3:]
        return url

    @classmethod
    def from_protobuf(cls: Type[T], pb_msg: 'str') -> T:
        """
        Read url from a proto msg.
        :param pb_msg:
        :return: url
        """
        return parse_obj_as(cls, pb_msg)

    def load_bytes(self, timeout: Optional[float] = None) -> bytes:
        """Convert url to bytes. This will either load or download the file and save
        it into a bytes object.
        :param timeout: timeout for urlopen. Only relevant if URI is not local
        :return: bytes.
        """
        if urllib.parse.urlparse(self).scheme in {'http', 'https', 'data'}:
            req = urllib.request.Request(
                self, headers={'User-Agent': 'Mozilla/5.0'}
            )
            urlopen_kwargs = {'timeout': timeout} if timeout is not None else {}
            with urllib.request.urlopen(req, **urlopen_kwargs) as fp:  # type: ignore
                return fp.read()
        elif os.path.exists(self):
            with open(self, 'rb') as fp:
                return fp.read()
        else:
            raise FileNotFoundError(f'`{self}` is not a URL or a valid local path')

@_register_proto(proto_type_name='audio_url')
class AudioUrl(AnyUrl):
    """
    URL to an audio file.
    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]
        and the frame rate.

        ---

        ```python
        from typing import Optional

        from docarray import BaseDoc
        from docarray.typing import AudioNdArray, AudioUrl


        class MyDoc(BaseDoc):
            audio_url: AudioUrl
            audio_tensor: Optional[AudioNdArray] = None


        doc = MyDoc(audio_url='https://www.kozco.com/tech/piano2.wav')
        doc.audio_tensor, _ = doc.audio_url.load()
        assert isinstance(doc.audio_tensor, AudioNdArray)
        ```

        ---

        :return: tuple of an [`AudioNdArray`][docarray.typing.AudioNdArray] representing
            the audio file content, and an integer representing the frame rate.

        """
        bytes_ = self.load_bytes()
        return bytes_.load()

    def load_bytes(self, timeout: Optional[float] = None) -> AudioBytes:
        """
        Convert url to [`AudioBytes`][docarray.typing.AudioBytes]. This will either load or
        download the file and save it into an [`AudioBytes`][docarray.typing.AudioBytes] object.

        :param timeout: timeout for urlopen. Only relevant if url is not local
        :return: [`AudioBytes`][docarray.typing.AudioBytes] object
        """
        bytes_ = super().load_bytes(timeout=timeout)
        return AudioBytes(bytes_)

    def display(self):
        """
        Play the audio sound from url in notebook.
        """
        if is_notebook():
            from IPython.display import Audio, display

            remote_url = True if self.startswith('http') else False

            if remote_url:
                display(Audio(data=self))
            else:
                display(Audio(filename=self))
        else:
            warnings.warn('Display of audio is only possible in a notebook.')

@_register_proto(proto_type_name='image_url')
class ImageUrl(AnyUrl):
    """
    URL to an image file.
    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

        ---

        ```python
        from pydantic import parse_obj_as

        from docarray import BaseDoc
        from docarray.typing import ImageUrl

        img_url = "https://upload.wikimedia.org/wikipedia/commons/8/80/Dag_Sebastian_Ahlander_at_G%C3%B6teborg_Book_Fair_2012b.jpg"

        img_url = parse_obj_as(ImageUrl, img_url)
        img = img_url.load_pil()

        from PIL.Image import Image

        assert isinstance(img, Image)
        ```

        ---
        :return: a Pillow image
        """
        from docarray.typing.bytes.image_bytes import ImageBytes

        return ImageBytes(self.load_bytes(timeout=timeout)).load_pil()

    def load(
        self,
        width: Optional[int] = None,
        height: Optional[int] = None,
        axis_layout: Tuple[str, str, str] = ('H', 'W', 'C'),
        timeout: Optional[float] = None,
    ) -> ImageNdArray:
        """
        Load the data from the url into an [`ImageNdArray`][docarray.typing.ImageNdArray]

        ---

        ```python
        from docarray import BaseDoc
        from docarray.typing import ImageUrl, ImageNdArray


        class MyDoc(BaseDoc):
            img_url: ImageUrl


        doc = MyDoc(
            img_url="https://upload.wikimedia.org/wikipedia/commons/8/80/"
            "Dag_Sebastian_Ahlander_at_G%C3%B6teborg_Book_Fair_2012b.jpg"
        )

        img_tensor = doc.img_url.load()
        assert isinstance(img_tensor, ImageNdArray)

        img_tensor = doc.img_url.load(height=224, width=224)
        assert img_tensor.shape == (224, 224, 3)

        layout = ('C', 'W', 'H')
        img_tensor = doc.img_url.load(height=100, width=200, axis_layout=layout)
        assert img_tensor.shape == (3, 200, 100)
        ```

        ---

        :param width: width of the image tensor.
        :param height: height of the image tensor.
        :param axis_layout: ordering of the different image axes.
            'H' = height, 'W' = width, 'C' = color channel
        :param timeout: timeout (sec) for urlopen network request.
            Only relevant if URL is not local
        :return: [`ImageNdArray`][docarray.typing.ImageNdArray] representing the image as RGB values
        """
        from docarray.typing.bytes.image_bytes import ImageBytes

        buffer = ImageBytes(self.load_bytes(timeout=timeout))
        return buffer.load(width, height, axis_layout)

    def load_bytes(self, timeout: Optional[float] = None) -> ImageBytes:
        """
        Convert url to [`ImageBytes`][docarray.typing.ImageBytes]. This will either load or
        download the file and save it into an [`ImageBytes`][docarray.typing.ImageBytes] object.

        :param timeout: timeout for urlopen. Only relevant if url is not local
        :return: [`ImageBytes`][docarray.typing.ImageBytes] object
        """
        bytes_ = super().load_bytes(timeout=timeout)
        return ImageBytes(bytes_)

    def display(self) -> None:
        """
        Display image data from url in notebook.
        """
        if is_notebook():
            from IPython.display import Image, display

            remote_url = True if self.startswith('http') else False
            if remote_url:
                display(Image(url=self))
            else:
                display(Image(filename=self))
        else:
            warnings.warn('Display of image is only possible in a notebook.')