controller-wisetopdataserver/python/biopro/file/segment.py

"""segment a file. used to upload file to network for someone can donwload.

Main
----

segment a file.

::

    python -m biopro.file.segment FILE [SIZE]

Class
-----

"""

from collections import deque
from math import ceil
from pathlib import Path
from typing import Dict, Optional
import struct
import base64

from biopro.util.json import JsonSerialize, JSON_OBJECT
from biopro.util.lock import Synchronized, synchronized


class FileSegment(Synchronized, JsonSerialize):
    """file segment information.

    **json format**

    ::

        segment_info = {
            'path': str
            'file_size': int
            'segment_total': int
            'segment_size': int
        }


    """
    DEFAULT_CACHE_SIZE = 16
    DEFAULT_SEGMENT_SIZE = 200 * 1024  # 200KB

    ALL = {}

    __slots__ = '_export_path', '_segment_size', '_file_size', '_total_segment', '_cache'

    def __init__(self,
                 export_path: Path,
                 segment_size = DEFAULT_SEGMENT_SIZE):
        """

        :param export_path: source path
        :param segment_size: size in bytes
        """
        super().__init__()

        if segment_size <= 0:
            raise ValueError('non-positive segment_size : ' + str(segment_size))

        if segment_size < 1024:
            raise ValueError('non-suggest segment_size, too small : ' + str(segment_size))

        self._export_path = export_path
        self._segment_size = segment_size
        self._file_size = 0

        self._total_segment = 0
        self._cache = deque([], maxlen=self.DEFAULT_CACHE_SIZE)

        self.reload()

    @synchronized
    def reload(self):
        """reload file cache"""
        if not self._export_path.exists():
            raise FileNotFoundError(str(self._export_path))

        self._file_size = file_size = self._export_path.stat().st_size
        self._total_segment = int(ceil(file_size / self._segment_size))
        self._cache.clear()

    @property
    def export_path(self) -> Path:
        """source path"""
        return self._export_path

    @property
    def segment_size(self) -> int:
        """segment size in bytes."""
        return self._segment_size

    @property
    def file_size(self) -> int:
        """original file size"""
        return self._file_size

    @property
    def total_size(self) -> int:
        """total segment number"""
        return self._total_segment

    def get_segment(self, index: int) -> 'FileSegmentContent':
        """get certain segment content.

        :param index:
        :return:
        """
        for i, content in self._cache:  # type: int, str
            if i == index:
                return FileSegmentContent(self, index, content)

        content = self._get_segment(index, 'BPRF')

        self._cache.append((index, content))

        return FileSegmentContent(self, index, content)

    @synchronized
    def _get_segment(self, index: int, _type: str) -> str:
        if not (0 <= index < self._total_segment):
            raise IndexError()

        if _type == 'BPRF':
            read = struct.calcsize('>4sIQ16s5B')
            with self._export_path.open('rb') as f:
                f.seek( read + index * self._segment_size )
                return base64.b64encode(f.read(self._segment_size)).decode() 
        else:
            with self._export_path.open('r') as f:
                f.seek(index * self._segment_size)
                return f.read(self._segment_size)

    @classmethod
    def get(cls,
            export_path: Path,
            segment_size = DEFAULT_SEGMENT_SIZE,
            create=False) -> Optional['FileSegment']:
        """get/create file segment.

        :param export_path:
        :param segment_size:
        :param create: create segment if not found
        :return: file segment.
        """
        print('*****1', export_path)
        ret = cls.ALL.get(export_path, None)

        if not create or ret is not None:
            return ret

        ret = FileSegment(export_path, segment_size)

        cls.ALL[export_path] = ret
        print('*****2', export_path)

        return ret

    def invalid(self):
        """invalid this segment and release the resource.

        :return:
        """
        self._cache.clear()

        try:
            del self.ALL[self._export_path]
        except KeyError:
            pass

    @classmethod
    def invalid_all(cls):
        cls.ALL.clear()

    def as_json(self) -> JSON_OBJECT:
        return {
            'path':          self._export_path,
            'file_size':     self._file_size,
            'segment_total': self._total_segment,
            'segment_size':  self._segment_size,
        }


class FileSegmentContent(JsonSerialize):
    """File segment with file segment content.

    **json format**

    ::

        segment_info = {
            'path': str
            'file_size': int
            'segment_total': int
            'segment_size': int
            'segment_index': int
            'segment': str
        }

    """
    __slots__ = '_segment', '_index', '_content'

    def __init__(self, segment: FileSegment, index: int, content: str):
        self._segment = segment
        self._index = index
        self._content = content

    def __len__(self):
        return len(self._content)

    def __str__(self):
        return self._content

    def __repr__(self):
        return 'FileSegment[%s:%d/%d]' % (self.export_path, self._index, self.total_size)

    @property
    def export_path(self) -> Path:
        """source path"""
        return self._segment.export_path

    @property
    def segment_size(self) -> int:
        """segment size in bytes."""
        return self._segment.segment_size

    @property
    def file_size(self) -> int:
        """original file size"""
        return self._segment.file_size

    @property
    def total_size(self) -> int:
        """total segment number"""
        return self._segment.total_size

    def as_json(self) -> JSON_OBJECT:
        ret = self._segment.as_json()
        ret['segment_index'] = self._index
        ret['segment'] = self._content
        return ret


if __name__ == '__main__':
    import sys


    def _main():
        try:
            sg_sz = int(sys.argv[2])
        except (IndexError, ValueError):
            sg_sz = 1024

        fs = FileSegment(Path(sys.argv[1]), sg_sz)

        for i in range(fs.total_size):
            fc = fs.get_segment(i)
            print('---', repr(fc), '---')
            print(str(fc))


    _main()