Source code for allel.chunked.storage_bcolz

# -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function, division
import tempfile
import atexit
import shutil
from types import MethodType


import bcolz


from allel.chunked import util as _util


def _table_append(ctbl, data):

    if hasattr(data, 'keys') and callable(data.keys):
        # normalise dict-like data
        data = [data[n] for n in ctbl.names]

    ctbl.append_original(data)


[docs]class BcolzStorage(object): """Storage layer using bcolz carray and ctable.""" def __init__(self, **kwargs): self.defaults = kwargs def _set_defaults(self, kwargs): for k, v in self.defaults.items(): kwargs.setdefault(k, v) return kwargs def array(self, data, expectedlen=None, **kwargs): data = _util.ensure_array_like(data) kwargs = self._set_defaults(kwargs) return bcolz.carray(data, expectedlen=expectedlen, **kwargs) def table(self, data, names=None, expectedlen=None, **kwargs): names, columns = _util.check_table_like(data, names=names) kwargs = self._set_defaults(kwargs) ctbl = bcolz.ctable(columns, names=names, expectedlen=expectedlen, **kwargs) # patch append method ctbl.append_original = ctbl.append ctbl.append = MethodType(_table_append, ctbl) return ctbl
[docs]class BcolzMemStorage(BcolzStorage): # noinspection PyShadowingBuiltins def _set_defaults(self, kwargs): for k, v in self.defaults.items(): kwargs.setdefault(k, v) kwargs['rootdir'] = None return kwargs
[docs]class BcolzTmpStorage(BcolzStorage): def _set_defaults(self, kwargs): for k, v in self.defaults.items(): kwargs.setdefault(k, v) suffix = kwargs.pop('suffix', '.bcolz') prefix = kwargs.pop('prefix', 'scikit_allel_') tempdir = kwargs.pop('dir', None) rootdir = tempfile.mkdtemp(suffix=suffix, prefix=prefix, dir=tempdir) atexit.register(shutil.rmtree, rootdir) kwargs['rootdir'] = rootdir kwargs['mode'] = 'w' return kwargs
bcolz_storage = BcolzStorage() """bcolz storage with default parameters""" bcolzmem_storage = BcolzMemStorage() """bcolz in-memory storage with default compression""" bcolztmp_storage = BcolzTmpStorage() """bcolz temporary file storage with default compression""" _zlib1 = bcolz.cparams(cname='zlib', clevel=1) bcolz_zlib1_storage = BcolzStorage(cparams=_zlib1) """bcolz storage with zlib level 1 compression""" bcolzmem_zlib1_storage = BcolzMemStorage(cparams=_zlib1) """bcolz in-memory storage with zlib level 1 compression""" bcolztmp_zlib1_storage = BcolzTmpStorage(cparams=_zlib1) """bcolz temporary file storage with zlib level 1 compression""" _util.storage_registry['bcolz'] = bcolz_storage _util.storage_registry['bcolzmem'] = bcolzmem_storage _util.storage_registry['bcolztmp'] = bcolztmp_storage _util.storage_registry['bcolz_zlib1'] = bcolz_zlib1_storage _util.storage_registry['bcolzmem_zlib1'] = bcolzmem_zlib1_storage _util.storage_registry['bcolztmp_zlib1'] = bcolztmp_zlib1_storage