text: from __future__ import absolute_import, division, print_function
  import numpy as np
  import six
  
  import bcolz
  from pybedtools import BedTool
  from pybedtools import Interval
  from pysam import FastaFile
  import pyBigWig
  
  from . import backend
  from .util import nan_to_zero
  from .util import one_hot_encode_sequence
  
  NUM_SEQ_CHARS = 4
  
  
  class BaseExtractor(object):
      dtype = np.float32
  
      def __init__(self, datafile, **kwargs):
          self._datafile = datafile
  
      def __call__(self, intervals, out=None, **kwargs):
          data = self._check_or_create_output_array(intervals, out)
          self._extract(intervals, data, **kwargs)
          return data
  
      def _check_or_create_output_array(self, intervals, out):
          width = intervals[0].stop - intervals[0].start
          output_shape = self._get_output_shape(len(intervals), width)
  
          if out is None:
              out = np.zeros(output_shape, dtype=self.dtype)
          else:
              if out.shape != output_shape:
                  raise ValueError('out array has incorrect shape: {} '
                                   '(need {})'.format(out.shape, output_shape))
              if out.dtype != self.dtype:
                  raise ValueError('out array has incorrect dtype: {} '
                                   '(need {})'.format(out.dtype, self.dtype))
          return out
  
      def _extract(self, intervals, out, **kwargs):
          'Subclassses should implement this and return the data'
          raise NotImplementedError
  
      @staticmethod
      def _get_output_shape(num_intervals, width):
          'Subclasses should implement this and return the shape of output'
          raise NotImplementedError
  
  
  class ArrayExtractor(BaseExtractor):
  
      def __init__(self, datafile, in_memory=False, **kwargs):
          super(ArrayExtractor, self).__init__(datafile, **kwargs)
          self._data = backend.load_directory(datafile, in_memory=in_memory)
          self.multiprocessing_safe = in_memory
  
          arr = next(iter(self._data.values()))
          def _mm_extract(self, intervals, out, **kwargs):
              mm_data = self._data
              for index, interval in enumerate(intervals):
                  out[index] = mm_data[interval.chrom][interval.start:interval.stop]
  
          # output shape method
          shape = arr.shape
          if len(shape) == 1:
              def _get_output_shape(num_intervals, width):
                  return (num_intervals, width)
          elif len(shape) == 2:
              def _get_output_shape(num_intervals, width):
                  return (num_intervals, width, shape[1])
          else:
              raise ValueError('Can only extract from 1D/2D arrays')
  
          self._mm_extract = _mm_extract.__get__(self)
          self._extract = self._mm_extract
          self._get_output_shape = staticmethod(_get_output_shape).__get__(self)
  
  
  class FastaExtractor(BaseExtractor):
  
      def __init__(self, datafile, use_strand=False, **kwargs):
          """Fasta file extractor
          
          NOTE: The extractor is not thread-save.
          If you with to use it with multiprocessing,
          create a new extractor object in each process.
          
          Args:
            datafile (str): path to the bigwig file
            use_strand (bool): if True, the extracted sequence
              is reverse complemented in case interval.strand == "-"
          """
          super(FastaExtractor, self).__init__(datafile, **kwargs)
          self.use_strand = use_strand
          self.fasta = FastaFile(self._datafile)
  
      def _extract(self, intervals, out, **kwargs):    
          for index, interval in enumerate(intervals):
              seq = self.fasta.fetch(str(interval.chrom), interval.start,
                                         interval.stop)
              one_hot_encode_sequence(seq, out[index, :, :])
  
              # reverse-complement seq the negative strand
              if self.use_strand and interval.strand == "-":
                  out[index, :, :] = out[index, ::-1, ::-1]
  
          return out
  
      @staticmethod
      def _get_output_shape(num_intervals, width):
          return (num_intervals, width, NUM_SEQ_CHARS)
  
  
  class BigwigExtractor(BaseExtractor):
  
      def __init__(self, datafile, **kwargs):
          """Big-wig file extractor
          
          NOTE: The extractor is not thread-save.
          If you with to use it with multiprocessing,
          create a new extractor object in each process.
          
          Args:
            datafile: path to the bigwig file
          """
          super(BigwigExtractor, self).__init__(datafile, **kwargs)
          self._verbose = kwargs.get('verbose', False)
          self.bw = pyBigWig.open(datafile)
          
      def _extract(self, intervals, out, **kwargs):
          out[:] = self._bigwig_extractor(self.bw, intervals,
                                          **kwargs)
  
          return out
  
      @staticmethod
      def _get_output_shape(num_intervals, width):
          return (num_intervals, width)
  
      @staticmethod
      def _bigwig_extractor(bw, intervals, out=None, **kwargs):
          nan_as_zero = kwargs.get('nan_as_zero', True)
          if out is None:
              width = intervals[0].stop - intervals[0].start
              out = np.zeros((len(intervals), width), dtype=np.float32)
              
          for index, interval in enumerate(intervals):
              out[index] = bw.values(
                  interval.chrom, interval.start, interval.stop)
              if nan_as_zero:
                  nan_to_zero(out[index])
          return out
      
      def __del__(self):
          return self.close()
      
      def close(self):
          return self.bw.close()
  