platypush/platypush/schemas/piped.py

import base64
from datetime import datetime

from marshmallow import EXCLUDE, fields, pre_dump
from marshmallow.schema import Schema

from platypush.schemas import StrippedString


class PipedVideoSchema(Schema):
    """
    Class for video items returned by the Piped API.
    """

    class Meta:
        """
        Exclude unknown fields.
        """

        unknown = EXCLUDE

    url = fields.Url(
        required=True,
        metadata={
            'description': 'Video URL',
            'example': 'https://youtube.com/watch?v=1234567890',
        },
    )

    title = StrippedString(
        missing='[No Title]',
        metadata={
            'description': 'Video title',
            'example': 'My Video Title',
        },
    )

    image = fields.Url(
        attribute='thumbnail',
        metadata={
            'description': 'Image URL',
            'example': 'https://i.ytimg.com/vi/1234567890/hqdefault.jpg',
        },
    )

    description = StrippedString(
        attribute='shortDescription',
        metadata={
            'description': 'Video description',
            'example': 'My video description',
        },
    )

    duration = fields.Int(
        missing=0,
        metadata={
            'description': 'Video duration in seconds',
            'example': 120,
        },
    )

    channel = fields.String(
        attribute='uploaderName',
        metadata={
            'description': 'Channel name',
            'example': 'My Channel',
        },
    )

    channel_url = fields.Url(
        attribute='uploaderUrl',
        metadata={
            'description': 'Channel URL',
            'example': 'https://youtube.com/channel/1234567890',
        },
    )

    channel_image = fields.Url(
        attribute='uploaderAvatar',
        metadata={
            'description': 'Channel image URL',
            'example': 'https://i.ytimg.com/vi/1234567890/hqdefault.jpg',
        },
    )

    created_at = fields.DateTime(
        attribute='uploaded',
        metadata={
            'description': 'Video upload date',
            'example': '2020-01-01T00:00:00',
        },
    )

    @pre_dump
    def fill_urls(self, data: dict, **_):
        for attr in ('url', 'uploaderUrl'):
            if data.get(attr) and not data[attr].startswith('https://'):
                data[attr] = f'https://youtube.com{data[attr]}'

        return data

    @pre_dump
    def normalize_timestamps(self, data: dict, **_):
        if data.get('uploaded') and isinstance(data['uploaded'], int):
            data['uploaded'] = datetime.fromtimestamp(data["uploaded"] / 1000)

        return data


class PipedPlaylistSchema(Schema):
    """
    Class for playlist items returned by the Piped API.
    """

    class Meta:
        """
        Exclude unknown fields.
        """

        unknown = EXCLUDE

    id = fields.UUID(
        required=True,
        metadata={
            'description': 'Playlist ID',
            'example': '12345678-1234-1234-1234-1234567890ab',
        },
    )

    name = StrippedString(
        missing='[No Name]',
        metadata={
            'description': 'Playlist name',
            'example': 'My Playlist Name',
        },
    )

    image = fields.Url(
        attribute='thumbnail',
        metadata={
            'description': 'Image URL',
            'example': 'https://i.ytimg.com/vi/1234567890/hqdefault.jpg',
        },
    )

    description = StrippedString(
        attribute='shortDescription',
        metadata={
            'description': 'Video description',
            'example': 'My video description',
        },
    )

    videos = fields.Int(
        missing=0,
        metadata={
            'description': 'Number of videos in the playlist',
            'example': 10,
        },
    )


class PipedChannelSchema(Schema):
    """
    Class for channel items returned by the Piped API.
    """

    class Meta:
        """
        Exclude unknown fields.
        """

        unknown = EXCLUDE

    id = fields.String(
        required=True,
        metadata={
            'description': 'Channel ID',
            'example': '1234567890',
        },
    )

    url = fields.String(
        required=True,
        metadata={
            'description': 'Channel URL',
            'example': 'https://youtube.com/channel/1234567890',
        },
    )

    name = StrippedString(
        missing='[No Name]',
        metadata={
            'description': 'Channel name',
            'example': 'My Channel Name',
        },
    )

    description = StrippedString(
        metadata={
            'description': 'Channel description',
            'example': 'My channel description',
        },
    )

    image = fields.Url(
        attribute='avatar',
        metadata={
            'description': 'Channel image URL',
            'example': 'https://i.ytimg.com/vi/1234567890/hqdefault.jpg',
        },
    )

    banner = fields.Url(
        attribute='bannerUrl',
        metadata={
            'description': 'Channel banner URL',
            'example': 'https://i.ytimg.com/vi/1234567890/hqdefault.jpg',
        },
    )

    subscribers = fields.Int(
        attribute='subscriberCount',
        missing=0,
        metadata={
            'description': 'Number of subscribers',
            'example': 1000,
        },
    )

    next_page_token = fields.String(
        attribute='nextpage',
        metadata={
            'description': 'The token that should be passed to get the next page of results',
            'example': '1234567890',
        },
    )

    items = fields.Nested(PipedVideoSchema, attribute='relatedStreams', many=True)

    @pre_dump
    def normalize_id_and_url(self, data: dict, **_):
        if data.get('id'):
            if not data.get('url'):
                data['url'] = f'https://youtube.com/channel/{data["id"]}'
        elif data.get('url'):
            data['id'] = data['url'].split('/')[-1]
            data['url'] = f'https://youtube.com{data["url"]}'
        else:
            raise AssertionError('Channel ID or URL not found')

        return data

    @pre_dump
    def normalize_avatar(self, data: dict, **_):
        if data.get('avatarUrl'):
            data['avatar'] = data.pop('avatarUrl')

        return data

    @pre_dump
    def serialize_next_page_token(self, data: dict, **_):
        if data.get('nextpage'):
            data['nextpage'] = base64.b64encode(data['nextpage'].encode()).decode()

        return data