wrolpi/modules/videos/__init__.py

import json
import pathlib
from abc import ABC
from typing import List, Dict

from sqlalchemy.orm import Session

from modules.videos.models import Video, Channel
from wrolpi.captions import read_captions, extract_captions
from wrolpi.common import logger, register_modeler, register_after_refresh, limit_concurrent, split_lines_by_length
from wrolpi.db import get_db_curs
from wrolpi.files.indexers import Indexer, register_indexer
from wrolpi.files.lib import split_path_stem_and_suffix
from wrolpi.files.models import File
from wrolpi.vars import PYTEST
from .downloader import video_downloader  # Import downloaders so they are registered.

logger = logger.getChild(__name__)

__all__ = ['video_modeler', 'VideoIndexer']


def find_video_file_in_group(group: List[File]):
    return next((i for i in group if i.mimetype.startswith('video/') and i.path.suffix != '.part'), None)


@register_modeler
def video_modeler(groups: Dict[str, List[File]], session: Session):
    new_videos = []

    # Search all groups for video files.
    video_files = {stem: video_file for stem, group in groups.items() if
                   (video_file := find_video_file_in_group(group))}
    if not video_files:
        # No videos in these groups.
        return
    # Get all matching Video records (if any) in one query.
    video_paths = [i.path for i in video_files.values()]
    video_records = {i.video_path: i for i in session.query(Video).filter(Video.video_path.in_(video_paths))}

    for stem, video_file in video_files.items():
        group = groups[stem]

        session.flush(group)
        poster_file = info_json_file = caption_file = None
        for file in group:
            if file.mimetype.startswith('image/'):
                poster_file = file
            elif file.path.name.endswith('.info.json'):
                info_json_file = file
            # Prefer WebVTT over SRT.  (SRT cannot be displayed for HTML video).
            elif file.path.name.endswith('.en.vtt'):
                caption_file = file
            elif file.path.name.endswith('.en.srt'):
                caption_file = caption_file or file

        if poster_file:
            poster_file.associated = True
            poster_file.do_stats()
        if caption_file:
            caption_file.associated = True
            caption_file.do_stats()
        if info_json_file:
            info_json_file.associated = True
            info_json_file.do_stats()

        video: Video = video_records.get(video_file.path)
        if not video:
            video = Video(video_file=video_file)
            session.add(video)

        size = video_file.path.stat().st_size

        if poster_file != video.poster_file or \
                caption_file != video.caption_file or \
                info_json_file != video.info_json_file or \
                video.size != size:
            # Files might have been changed.  Re-index.
            video.video_file.indexed = False

        video.poster_file = poster_file
        video.caption_file = caption_file
        video.info_json_file = info_json_file
        video.size = size
        video_file.model = Video.__tablename__

        new_videos.append(video)

    if new_videos:
        session.flush(new_videos)
        for video in new_videos:
            video.video_file.do_index()
            video.validate(session)

            # Remove this group, it will not be processed by other modelers.
            stem, _ = split_path_stem_and_suffix(video.video_path)
            del groups[stem]


@register_after_refresh
@limit_concurrent(1)
def video_cleanup():
    # Claim all Videos in a Channel's directory for that Channel.  But, only if they have not yet been claimed.
    logger.info('Claiming Videos for their Channels')
    with get_db_curs(commit=True) as curs:
        curs.execute('''
            UPDATE video v
            SET channel_id = c.id
            FROM channel c
            WHERE
             v.video_path LIKE c.directory || '/%'::VARCHAR
             AND v.channel_id IS NULL
        ''')


EXTRACT_SUBTITLES = False


@register_indexer('video')
class VideoIndexer(Indexer, ABC):
    """Handles video files like mp4/ogg."""

    @staticmethod
    def get_description(file):
        video_path: pathlib.Path = file.path.path if hasattr(file.path, 'path') else file.path
        info_json_path = video_path.with_suffix('.info.json')
        if info_json_path.is_file():
            with info_json_path.open('rt') as fh:
                try:
                    return json.load(fh)['description']
                except Exception:
                    if not PYTEST:
                        logger.warning(f'Video info json file exists, but cannot get description. {file}')
                    return None

    @classmethod
    def create_index(cls, file):
        """
        Index the video file and it's associated files.

        a = title.
        b = <empty>
        c = description from info json.
        d = captions from info json, or extracted from video file.
        """
        a = cls.get_title(file)
        c = cls.get_description(file)

        # Detect the associated caption files.
        d = None
        if (en_vtt := file.path.with_suffix('.en.vtt')).is_file():
            d = read_captions(en_vtt)
        elif (vtt := file.path.with_suffix('.vtt')).is_file():
            d = read_captions(vtt)
        elif (en_srt := file.path.with_suffix('.en.srt')).is_file():
            d = read_captions(en_srt)
        elif (srt := file.path.with_suffix('.srt')).is_file():
            d = read_captions(srt)
        elif EXTRACT_SUBTITLES:
            d = extract_captions(file.path) or ''

        d = split_lines_by_length(d)
        return a, None, c, d