lbry-sdk/lbry/file/source_manager.py

import os
import asyncio
import logging
import typing
from typing import Optional
from lbry.file.source import ManagedDownloadSource
if typing.TYPE_CHECKING:
    from lbry.conf import Config
    from lbry.extras.daemon.analytics import AnalyticsManager
    from lbry.extras.daemon.storage import SQLiteStorage

log = logging.getLogger(__name__)

COMPARISON_OPERATORS = {
    'eq': lambda a, b: a == b,
    'ne': lambda a, b: a != b,
    'g': lambda a, b: a > b,
    'l': lambda a, b: a < b,
    'ge': lambda a, b: a >= b,
    'le': lambda a, b: a <= b,
}


class SourceManager:
    filter_fields = {
        'rowid',
        'status',
        'file_name',
        'added_on',
        'download_path',
        'claim_name',
        'claim_height',
        'claim_id',
        'outpoint',
        'txid',
        'nout',
        'channel_claim_id',
        'channel_name',
        'completed'
    }

    set_filter_fields = {
        "claim_ids": "claim_id",
        "channel_claim_ids": "channel_claim_id",
        "outpoints": "outpoint"
    }

    source_class = ManagedDownloadSource

    def __init__(self, loop: asyncio.AbstractEventLoop, config: 'Config', storage: 'SQLiteStorage',
                 analytics_manager: Optional['AnalyticsManager'] = None):
        self.loop = loop
        self.config = config
        self.storage = storage
        self.analytics_manager = analytics_manager
        self._sources: typing.Dict[str, ManagedDownloadSource] = {}
        self.started = asyncio.Event(loop=self.loop)

    def add(self, source: ManagedDownloadSource):
        self._sources[source.identifier] = source

    def remove(self, source: ManagedDownloadSource):
        if source.identifier not in self._sources:
            return
        self._sources.pop(source.identifier)
        source.stop_tasks()

    async def initialize_from_database(self):
        raise NotImplementedError()

    async def start(self):
        await self.initialize_from_database()
        self.started.set()

    def stop(self):
        while self._sources:
            _, source = self._sources.popitem()
            source.stop_tasks()
        self.started.clear()

    async def create(self, file_path: str, key: Optional[bytes] = None,
                     iv_generator: Optional[typing.Generator[bytes, None, None]] = None) -> ManagedDownloadSource:
        raise NotImplementedError()

    async def delete(self, source: ManagedDownloadSource, delete_file: Optional[bool] = False):
        self.remove(source)
        if delete_file and source.output_file_exists:
            os.remove(source.full_path)

    def get_filtered(self, sort_by: Optional[str] = None, reverse: Optional[bool] = False,
                     comparison: Optional[str] = None, **search_by) -> typing.List[ManagedDownloadSource]:
        """
        Get a list of filtered and sorted ManagedStream objects

        :param sort_by: field to sort by
        :param reverse: reverse sorting
        :param comparison: comparison operator used for filtering
        :param search_by: fields and values to filter by
        """
        if sort_by and sort_by not in self.filter_fields:
            raise ValueError(f"'{sort_by}' is not a valid field to sort by")
        if comparison and comparison not in COMPARISON_OPERATORS:
            raise ValueError(f"'{comparison}' is not a valid comparison")
        if 'full_status' in search_by:
            del search_by['full_status']

        for search in search_by:
            if search not in self.filter_fields:
                raise ValueError(f"'{search}' is not a valid search operation")

        compare_sets = {}
        if isinstance(search_by.get('claim_id'), list):
            compare_sets['claim_ids'] = search_by.pop('claim_id')
        if isinstance(search_by.get('outpoint'), list):
            compare_sets['outpoints'] = search_by.pop('outpoint')
        if isinstance(search_by.get('channel_claim_id'), list):
            compare_sets['channel_claim_ids'] = search_by.pop('channel_claim_id')

        if search_by or compare_sets:
            comparison = comparison or 'eq'
            streams = []
            for stream in self._sources.values():
                if compare_sets and not all(
                        getattr(stream, self.set_filter_fields[set_search]) in val
                        for set_search, val in compare_sets.items()):
                    continue
                if search_by and not all(
                        COMPARISON_OPERATORS[comparison](getattr(stream, search), val)
                        for search, val in search_by.items()):
                    continue
                streams.append(stream)
        else:
            streams = list(self._sources.values())
        if sort_by:
            streams.sort(key=lambda s: getattr(s, sort_by))
            if reverse:
                streams.reverse()
        return streams