<commit_before>'''
Created on Aug 28, 2014

@author: David Zwicker <dzwicker@seas.harvard.edu>

Holds classes describing mouse burrows
'''

from __future__ import division

import copy
import itertools

import numpy as np
import cv2
import shapely
from shapely import geometry

from .utils import LazyHDFCollection, rtrim_nan
from video.analysis import curves, regions
from video.analysis.utils import cached_property

from ..debug import *  # @UnusedWildImport

# monkey patch shapely.geometry to get compatibility with older shapely versions
if not hasattr(geometry, 'LinearRing'):
    geometry.LinearRing = geometry.polygon.LinearRing

def bool_nan(number):
    """ returns the boolean value of the number.
    Returns False if number is not-a-number"""
    if np.isnan(number):
        return False
    else:
        return bool(number)

class Burrow(object):
    """ represents a single burrow.
    Note that the outline are always given in clockwise direction
    """

    storage_class = LazyHDFCollection
    array_columns = ('Outline X', 'Outline Y',
                     'Burrow length + Centerline X',
                     'Flag if burrow was refined + Centerline Y')

    def __init__(self, outline, centerline=None, length=None,
                 refined=False, two_exits=False):
        """ initialize the structure using line on its outline """
        if len(outline) < 3:
            raise ValueError("Burrow outline must be defined by at least "
                             "three points.")
        self.centerline = centerline
        self.outline = outline
        self.refined = refined
        self.two_exits = two_exits

        if length is not None:
            self.length = length

    def copy(self):
        return Burrow(copy.copy(self.outline), copy.copy(self.centerline),
                      self.length, self.refined)

    def __repr__(self):
        center = self.polygon.centroid
        flags = ['center=(%d, %d)' % (center.x, center.y),
                 'points=%d' % len(self._outline)]
        if self.length:
            flags.append('length=%d' % self.length)
        else:
            flags.append('area=%d' % self.polygon.area)
        if self.refined:
            flags.append('refined')
        if self.two_exits:
            flags.append('two_exits')

        return ('Burrow(%s)' % (', '.join(flags)))

    @property
    def centerline(self):
        return self._centerline

    @centerline.setter
    def centerline(self, points):
        if points is None:
            self._centerline = None
            self.length = 0
        else:
            self._centerline = np.array(points, np.double)
            self.length = curves.curve_length(self._centerline)
        self._cache = {}

    @property
    def outline(self):
        return self._outline

    @outline.setter
    def outline(self, points):
        """ set the outline of the burrow.
        `point_list` can be a list/array of points or a shapely LinearRing
        """
        if points is None:
            self._outline = None
        else:
            if isinstance(points, geometry.LinearRing):
                ring = points
            else:
                ring = geometry.LinearRing(points)

            # make sure that the outline is given in clockwise direction
            self._outline = np.array(points, np.double)
            if ring.is_ccw:
                self._outline = self._outline[::-1]

        self._cache = {}

    def merge(self, other):
        """ merge this burrow with another one """
        polygon = self.polygon.union(other.polygon)
        self.outline = regions.get_enclosing_outline(polygon)

        # set the centerline to the longest of the two
        if other.length > self.length:
            self.centerline = other.centerline

    @cached_property
    def outline_ring(self):
        """ return the linear ring of the burrow outline """
        return geometry.LinearRing(self.outline)

    @cached_property
    def polygon(self):
        """ return the polygon of the burrow outline """
        return geometry.Polygon(self.outline)

    @cached_property
    def linestring(self):
        if self.centerline is None:
            return geometry.LineString()
        else:
            return geometry.LineString(self.centerline)

    @cached_property
    def position(self):
        return self.polygon.representative_point()

    @cached_property
    def area(self):
        """ return the area of the burrow shape """
        return self.polygon.area

    @property
    def entry_point(self):
        if self.centerline is None:
            return None
        else:
            return self.centerline[0]

    @property
    def end_point(self):
        if self.centerline is None:
            return None
        else:
            return self.centerline[-1]

    @property
    def is_valid(self):
        return len(self._outline) > 3

    @cached_property
    def eccentricity(self):
        """ return the eccentricity of the burrow shape
        The eccentricity will be between 0 and 1, corresponding to a circle
        and a straight line, respectively.
        """
        m = cv2.moments(np.asarray(self.outline, np.uint8))
        a, b, c = m['mu20'], -m['mu11'], m['mu02']
        e1 = (a + c) + np.sqrt(4*b**2 + (a - c)**2)
        e2 = (a + c) - np.sqrt(4*b**2 + (a - c)**2)
        if e1 == 0:
            return 0
        else:
            return np.sqrt(1 - e2/e1)

    def contains(self, point):
        """ returns True if the point is inside the burrow """
        return self.polygon.contains(geometry.Point(point))

    def intersects(self, burrow_or_shape):
        """ returns True if polygon intersects the burrow """
        if isinstance(burrow_or_shape, Burrow):
            burrow_or_shape = burrow_or_shape.polygon
        try:
            return self.polygon.intersects(burrow_or_shape)
        except shapely.geos.TopologicalError:
            return False

    def simplify_outline(self, tolerance=0.1):
        """ simplifies the outline """
        outline = geometry.LineString(self.outline)
        tolerance *= outline.length
        outline = outline.simplify(tolerance, preserve_topology=True)
        self.outline = outline.coords

    def get_bounding_rect(self, margin=0):
        """ returns the bounding rectangle of the burrow """
        burrow_points = self.outline
        if self.centerline is not None:
            burrow_points = np.vstack((burrow_points, self.centerline))
        bounds = geometry.MultiPoint(burrow_points).bounds
        bound_rect = regions.corners_to_rect(bounds[:2], bounds[2:])
        if margin:
            bound_rect = regions.expand_rectangle(bound_rect, margin)
        return np.asarray(bound_rect, np.int)

    def extend_outline(self, extension_polygon, simplify_threshold):
        """ extends the outline of the burrow to also enclose the object given
        by polygon """
        # get the union of the burrow and the extension
        burrow = self.polygon.union(extension_polygon)

        # determine the outline of the union
        outline = regions.get_enclosing_outline(burrow)

        outline = outline.simplify(simplify_threshold*outline.length)

        self.outline = np.asarray(outline, np.int32)

    def get_mask(self, margin=0, dtype=np.uint8, ret_shift=False):
        """ builds a mask of the burrow """
        # prepare the array to store the mask into
        rect = self.get_bounding_rect(margin=margin)
        mask = np.zeros((rect[3], rect[2]), dtype)

        # draw the burrow into the mask
        outline = np.asarray(self.outline, np.int)
        offset = (-rect[0], -rect[1])
        cv2.fillPoly(mask, [outline], color=1, offset=offset)

        if ret_shift:
            return mask, (-offset[0], -offset[1])
        else:
            return mask

    def to_array(self):
        """ converts the internal representation to a single array """
        # collect the data for the first two columns
        if self.outline is None:
            data1 = np.zeros((0, 2), np.double)
        else:
            data1 = np.asarray(self.outline, np.double)

        # collect the data for the last two columns
        data2 = np.array([[self.length, self.refined],
                          [self.two_exits, np.nan]], np.double)
        if self.centerline is not None:
            data2 = np.concatenate((data2, self.centerline))

        # save the data in a convenient array
        l1, l2 = len(data1), len(data2)
        data_res = np.empty((max(l1, l2), 4), np.double)
        data_res.fill(np.nan)
        data_res[:l1, :2] = data1
        data_res[:l2, 2:] = data2
        return data_res

    @classmethod
    def create_from_array(cls, data):
        """ creates a burrow track from a single array """
        # load the data from the respective places
        data = np.asarray(data)
        data_attr = data[:2, 2:]
        data_outline = rtrim_nan(data[:, :2])
        data_centerline = rtrim_nan(data[2:, 2:])
        if len(data_centerline) == 0:
            data_centerline = None

        # create the object
        return cls(outline=data_outline,
                   centerline=data_centerline,
                   length=data_attr[0, 0],
                   refined=bool_nan(data_attr[0, 1]),
                   two_exits=bool_nan(data_attr[1, 0]))

class BurrowTrack(object):
    """ class that stores the evolution of a single burrow over time.
    """
    column_names = ('Time',) + Burrow.array_columns

    def __init__(self, time=None, burrow=None):
        self.times = [] if time is None else [time]
        self.burrows = [] if burrow is None else [burrow]

    def __repr__(self):
        if len(self.times) == 0:
            return 'BurrowTrack([])'
        elif len(self.times) == 1:
            return 'BurrowTrack(time=%d)' % (self.times[0])
        else:
            return 'BurrowTrack(times=%d..%d)' % (self.times[0], self.times[-1])

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

    def append(self, time, burrow):
        """ append a new burrow with a time code """
        self.times.append(time)
        self.burrows.append(burrow)

    def __delitem__(self, key):
        """ deletes a certain burrow from the track list """
        del self.times[key]
        del self.burrows[key]

    @property
    def last(self):
        """ return the last burrow in the track """
        return self.burrows[-1]

    @property
    def track_start(self): return self.times[0]
    @property
    def track_end(self): return self.times[-1]

    def get_burrow(self, time, ret_next_change=False):
        """ returns the burrow at a specific time.
        If ret_next_change is True, we also return the frame where the burrow
        changes next. """
        if not self.times[0] <= time <= self.times[-1]:
            raise IndexError

        idx = np.argmin(np.abs(np.asarray(self.times) - time))
        burrow = self.burrows[idx]
        if ret_next_change:
            if idx >= len(self.times) - 1:
                return burrow, time + 1
            else:
                return burrow, (self.times[idx] + self.times[idx + 1])/2
        else:
            return burrow

    def to_array(self):
        """ converts the internal representation to a single array
        useful for storing the data """
        res = []
        for time, burrow in itertools.izip(self.times, self.burrows):
            burrow_data = burrow.to_array()
            time_array = np.zeros((len(burrow_data), 1), np.int32) + time
            res.append(np.hstack((time_array, burrow_data)))
        if res:
            return np.vstack(res)
        else:
            return []

    @classmethod
    def create_from_array(cls, data):
        """ constructs an object from an array previously created by to_array() """
        burrow_track = cls()
        burrow_data = None
        time_cur = -1
        for d in data:
            if d[0] != time_cur:
                if burrow_data is not None:
                    burrow_track.append(time_cur, Burrow.create_from_array(burrow_data))
                time_cur = d[0]
                burrow_data = [d[1:]]
            else:
                burrow_data.append(d[1:])

        if burrow_data is not None:
            burrow_track.append(time_cur, Burrow.create_from_array(burrow_data))

        return burrow_track

    def save_to_hdf5(self, hdf_file, key):
        """ save the data of the current burrow to an HDF5 file """
        if key in hdf_file:
            del hdf_file[key]
        hdf_file.create_dataset(key, data=self.to_array(), track_times=True)

    @classmethod
    def create_from_hdf5(cls, hdf_file, key):
        """ creates a burrow track from data in a HDF5 file """
        return cls.create_from_array(hdf_file[key])

class BurrowTrackList(list):
    """ class that stores instances of BurrowTrack in a list """
    item_class = BurrowTrack
    storage_class = LazyHDFCollection
    hdf_attributes = {'column_names': BurrowTrack.column_names,
                      'remark':
                        'Each burrow is represented by its outline saved as a '
                        'list of points of the format (Time, X, Y, a, b), '
                        'where all points with the same Time belong to the same '
                        'burrow. The last two values (a, b) contain additional '
                        'data, e.g. the burrow length and the coordinates of '
                        'the centerline.'}

    def create_track(self, frame_id, burrow):
        """ creates a new burrow track and appends it to the list """
        self.append(BurrowTrack(frame_id, burrow))

    def find_burrows(self, frame_id, ret_next_change=False):
        """ returns a list of all burrows active in a given frame.
        If ret_next_change is True, the number of the frame where something
        will change in the returned burrows is also returned. This can be
        useful while iterating over all frames. """
        result = []
        for burrow_track in self:
            try:
                res = burrow_track.get_burrow(frame_id, ret_next_change)
            except IndexError:
                continue
            else:
                result.append(res)

        if ret_next_change:
            burrows = [res[0] for res in result]

            if burrows:
                next_change = min(res[1] for res in result)
            else:
                next_change = frame_id + 1
            return burrows, next_change
        else:
            return result

<commit_msg>Fixed bug occurring with strange shapes<commit_after>'''
Created on Aug 28, 2014

@author: David Zwicker <dzwicker@seas.harvard.edu>

Holds classes describing mouse burrows
'''

from __future__ import division

import copy
import itertools

import numpy as np
import cv2
import shapely
from shapely import geometry

from .utils import LazyHDFCollection, rtrim_nan
from video.analysis import curves, regions
from video.analysis.utils import cached_property

from ..debug import *  # @UnusedWildImport

# monkey patch shapely.geometry to get compatibility with older shapely versions
if not hasattr(geometry, 'LinearRing'):
    geometry.LinearRing = geometry.polygon.LinearRing

def bool_nan(number):
    """ returns the boolean value of the number.
    Returns False if number is not-a-number"""
    if np.isnan(number):
        return False
    else:
        return bool(number)

class Burrow(object):
    """ represents a single burrow.
    Note that the outline are always given in clockwise direction
    """

    storage_class = LazyHDFCollection
    array_columns = ('Outline X', 'Outline Y',
                     'Burrow length + Centerline X',
                     'Flag if burrow was refined + Centerline Y')

    def __init__(self, outline, centerline=None, length=None,
                 refined=False, two_exits=False):
        """ initialize the structure using line on its outline """
        if len(outline) < 3:
            raise ValueError("Burrow outline must be defined by at least "
                             "three points.")
        self.centerline = centerline
        self.outline = outline
        self.refined = refined
        self.two_exits = two_exits

        if length is not None:
            self.length = length

    def copy(self):
        return Burrow(copy.copy(self.outline), copy.copy(self.centerline),
                      self.length, self.refined)

    def __repr__(self):
        center = self.polygon.centroid
        flags = ['center=(%d, %d)' % (center.x, center.y),
                 'points=%d' % len(self._outline)]
        if self.length:
            flags.append('length=%d' % self.length)
        else:
            flags.append('area=%d' % self.polygon.area)
        if self.refined:
            flags.append('refined')
        if self.two_exits:
            flags.append('two_exits')

        return ('Burrow(%s)' % (', '.join(flags)))

    @property
    def centerline(self):
        return self._centerline

    @centerline.setter
    def centerline(self, points):
        if points is None:
            self._centerline = None
            self.length = 0
        else:
            self._centerline = np.array(points, np.double)
            self.length = curves.curve_length(self._centerline)
        self._cache = {}

    @property
    def outline(self):
        return self._outline

    @outline.setter
    def outline(self, points):
        """ set the outline of the burrow.
        `point_list` can be a list/array of points or a shapely LinearRing
        """
        if points is None:
            self._outline = None
        else:
            if isinstance(points, geometry.LinearRing):
                ring = points
            else:
                ring = geometry.LinearRing(points)

            # make sure that the outline is given in clockwise direction
            self._outline = np.array(points, np.double)
            if ring.is_ccw:
                self._outline = self._outline[::-1]

        self._cache = {}

    def merge(self, other):
        """ merge this burrow with another one """
        polygon = self.polygon.union(other.polygon)
        self.outline = regions.get_enclosing_outline(polygon)

        # set the centerline to the longest of the two
        if other.length > self.length:
            self.centerline = other.centerline

    @cached_property
    def outline_ring(self):
        """ return the linear ring of the burrow outline """
        return geometry.LinearRing(self.outline)

    @cached_property
    def polygon(self):
        """ return the polygon of the burrow outline """
        return geometry.Polygon(self.outline)

    @cached_property
    def linestring(self):
        if self.centerline is None:
            return geometry.LineString()
        else:
            return geometry.LineString(self.centerline)

    @cached_property
    def position(self):
        return self.polygon.representative_point()

    @cached_property
    def area(self):
        """ return the area of the burrow shape """
        return self.polygon.area

    @property
    def entry_point(self):
        if self.centerline is None:
            return None
        else:
            return self.centerline[0]

    @property
    def end_point(self):
        if self.centerline is None:
            return None
        else:
            return self.centerline[-1]

    @property
    def is_valid(self):
        return len(self._outline) > 3

    @cached_property
    def eccentricity(self):
        """ return the eccentricity of the burrow shape
        The eccentricity will be between 0 and 1, corresponding to a circle
        and a straight line, respectively.
        """
        m = cv2.moments(np.asarray(self.outline, np.uint8))
        a, b, c = m['mu20'], -m['mu11'], m['mu02']
        e1 = (a + c) + np.sqrt(4*b**2 + (a - c)**2)
        e2 = (a + c) - np.sqrt(4*b**2 + (a - c)**2)
        if e1 == 0:
            return 0
        else:
            return np.sqrt(1 - e2/e1)

    def contains(self, point):
        """ returns True if the point is inside the burrow """
        return self.polygon.contains(geometry.Point(point))

    def intersects(self, burrow_or_shape):
        """ returns True if polygon intersects the burrow """
        if isinstance(burrow_or_shape, Burrow):
            burrow_or_shape = burrow_or_shape.polygon
        try:
            return self.polygon.intersects(burrow_or_shape)
        except (shapely.geos.TopologicalError, shapely.geos.PredicateError):
            return False

    def simplify_outline(self, tolerance=0.1):
        """ simplifies the outline """
        outline = geometry.LineString(self.outline)
        tolerance *= outline.length
        outline = outline.simplify(tolerance, preserve_topology=True)
        self.outline = outline.coords

    def get_bounding_rect(self, margin=0):
        """ returns the bounding rectangle of the burrow """
        burrow_points = self.outline
        if self.centerline is not None:
            burrow_points = np.vstack((burrow_points, self.centerline))
        bounds = geometry.MultiPoint(burrow_points).bounds
        bound_rect = regions.corners_to_rect(bounds[:2], bounds[2:])
        if margin:
            bound_rect = regions.expand_rectangle(bound_rect, margin)
        return np.asarray(bound_rect, np.int)

    def extend_outline(self, extension_polygon, simplify_threshold):
        """ extends the outline of the burrow to also enclose the object given
        by polygon """
        # get the union of the burrow and the extension
        burrow = self.polygon.union(extension_polygon)

        # determine the outline of the union
        outline = regions.get_enclosing_outline(burrow)

        outline = outline.simplify(simplify_threshold*outline.length)

        self.outline = np.asarray(outline, np.int32)

    def get_mask(self, margin=0, dtype=np.uint8, ret_shift=False):
        """ builds a mask of the burrow """
        # prepare the array to store the mask into
        rect = self.get_bounding_rect(margin=margin)
        mask = np.zeros((rect[3], rect[2]), dtype)

        # draw the burrow into the mask
        outline = np.asarray(self.outline, np.int)
        offset = (-rect[0], -rect[1])
        cv2.fillPoly(mask, [outline], color=1, offset=offset)

        if ret_shift:
            return mask, (-offset[0], -offset[1])
        else:
            return mask

    def to_array(self):
        """ converts the internal representation to a single array """
        # collect the data for the first two columns
        if self.outline is None:
            data1 = np.zeros((0, 2), np.double)
        else:
            data1 = np.asarray(self.outline, np.double)

        # collect the data for the last two columns
        data2 = np.array([[self.length, self.refined],
                          [self.two_exits, np.nan]], np.double)
        if self.centerline is not None:
            data2 = np.concatenate((data2, self.centerline))

        # save the data in a convenient array
        l1, l2 = len(data1), len(data2)
        data_res = np.empty((max(l1, l2), 4), np.double)
        data_res.fill(np.nan)
        data_res[:l1, :2] = data1
        data_res[:l2, 2:] = data2
        return data_res

    @classmethod
    def create_from_array(cls, data):
        """ creates a burrow track from a single array """
        # load the data from the respective places
        data = np.asarray(data)
        data_attr = data[:2, 2:]
        data_outline = rtrim_nan(data[:, :2])
        data_centerline = rtrim_nan(data[2:, 2:])
        if len(data_centerline) == 0:
            data_centerline = None

        # create the object
        return cls(outline=data_outline,
                   centerline=data_centerline,
                   length=data_attr[0, 0],
                   refined=bool_nan(data_attr[0, 1]),
                   two_exits=bool_nan(data_attr[1, 0]))

class BurrowTrack(object):
    """ class that stores the evolution of a single burrow over time.
    """
    column_names = ('Time',) + Burrow.array_columns

    def __init__(self, time=None, burrow=None):
        self.times = [] if time is None else [time]
        self.burrows = [] if burrow is None else [burrow]

    def __repr__(self):
        if len(self.times) == 0:
            return 'BurrowTrack([])'
        elif len(self.times) == 1:
            return 'BurrowTrack(time=%d)' % (self.times[0])
        else:
            return 'BurrowTrack(times=%d..%d)' % (self.times[0], self.times[-1])

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

    def append(self, time, burrow):
        """ append a new burrow with a time code """
        self.times.append(time)
        self.burrows.append(burrow)

    def __delitem__(self, key):
        """ deletes a certain burrow from the track list """
        del self.times[key]
        del self.burrows[key]

    @property
    def last(self):
        """ return the last burrow in the track """
        return self.burrows[-1]

    @property
    def track_start(self): return self.times[0]
    @property
    def track_end(self): return self.times[-1]

    def get_burrow(self, time, ret_next_change=False):
        """ returns the burrow at a specific time.
        If ret_next_change is True, we also return the frame where the burrow
        changes next. """
        if not self.times[0] <= time <= self.times[-1]:
            raise IndexError

        idx = np.argmin(np.abs(np.asarray(self.times) - time))
        burrow = self.burrows[idx]
        if ret_next_change:
            if idx >= len(self.times) - 1:
                return burrow, time + 1
            else:
                return burrow, (self.times[idx] + self.times[idx + 1])/2
        else:
            return burrow

    def to_array(self):
        """ converts the internal representation to a single array
        useful for storing the data """
        res = []
        for time, burrow in itertools.izip(self.times, self.burrows):
            burrow_data = burrow.to_array()
            time_array = np.zeros((len(burrow_data), 1), np.int32) + time
            res.append(np.hstack((time_array, burrow_data)))
        if res:
            return np.vstack(res)
        else:
            return []

    @classmethod
    def create_from_array(cls, data):
        """ constructs an object from an array previously created by to_array() """
        burrow_track = cls()
        burrow_data = None
        time_cur = -1
        for d in data:
            if d[0] != time_cur:
                if burrow_data is not None:
                    burrow_track.append(time_cur, Burrow.create_from_array(burrow_data))
                time_cur = d[0]
                burrow_data = [d[1:]]
            else:
                burrow_data.append(d[1:])

        if burrow_data is not None:
            burrow_track.append(time_cur, Burrow.create_from_array(burrow_data))

        return burrow_track

    def save_to_hdf5(self, hdf_file, key):
        """ save the data of the current burrow to an HDF5 file """
        if key in hdf_file:
            del hdf_file[key]
        hdf_file.create_dataset(key, data=self.to_array(), track_times=True)

    @classmethod
    def create_from_hdf5(cls, hdf_file, key):
        """ creates a burrow track from data in a HDF5 file """
        return cls.create_from_array(hdf_file[key])

class BurrowTrackList(list):
    """ class that stores instances of BurrowTrack in a list """
    item_class = BurrowTrack
    storage_class = LazyHDFCollection
    hdf_attributes = {'column_names': BurrowTrack.column_names,
                      'remark':
                        'Each burrow is represented by its outline saved as a '
                        'list of points of the format (Time, X, Y, a, b), '
                        'where all points with the same Time belong to the same '
                        'burrow. The last two values (a, b) contain additional '
                        'data, e.g. the burrow length and the coordinates of '
                        'the centerline.'}

    def create_track(self, frame_id, burrow):
        """ creates a new burrow track and appends it to the list """
        self.append(BurrowTrack(frame_id, burrow))

    def find_burrows(self, frame_id, ret_next_change=False):
        """ returns a list of all burrows active in a given frame.
        If ret_next_change is True, the number of the frame where something
        will change in the returned burrows is also returned. This can be
        useful while iterating over all frames. """
        result = []
        for burrow_track in self:
            try:
                res = burrow_track.get_burrow(frame_id, ret_next_change)
            except IndexError:
                continue
            else:
                result.append(res)

        if ret_next_change:
            burrows = [res[0] for res in result]

            if burrows:
                next_change = min(res[1] for res in result)
            else:
                next_change = frame_id + 1
            return burrows, next_change
        else:
            return result

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import re

from module.plugins.internal.Crypter import Crypter
from module.plugins.internal.SimpleHoster import SimpleHoster, create_getInfo, replace_patterns, set_cookie, set_cookies
from module.utils import fixup, html_unescape

class SimpleCrypter(Crypter, SimpleHoster):
    __name__    = "SimpleCrypter"
    __type__    = "crypter"
    __version__ = "0.60"
    __status__  = "testing"

    __pattern__ = r'^unmatchable$'
    __config__  = [("use_subfolder"     , "bool", "Save package to subfolder"          , True),  #: Overrides pyload.config['general']['folder_per_package']
                   ("subfolder_per_pack", "bool", "Create a subfolder for each package", True)]

    __description__ = """Simple decrypter plugin"""
    __license__     = "GPLv3"
    __authors__     = [("Walter Purcaro", "vuolter@gmail.com")]

    """
    Following patterns should be defined by each crypter:

      LINK_PATTERN: Download link or regex to catch links in group(1)
        example: LINK_PATTERN = r'
<a href="(.+?)"'

      NAME_PATTERN: (optional) folder name or page title
        example: NAME_PATTERN = r'<title>Files of: (?P<N>[^<]+) folder</title>'

      OFFLINE_PATTERN: (optional) Checks if the page is unreachable
        example: OFFLINE_PATTERN = r'File (deleted|not found)'

      TEMP_OFFLINE_PATTERN: (optional) Checks if the page is temporarily unreachable
        example: TEMP_OFFLINE_PATTERN = r'Server maintainance'

    You can override the getLinks method if you need a more sophisticated way to extract the links.

    If the links are splitted on multiple pages you can define the PAGES_PATTERN regex:

      PAGES_PATTERN: (optional) group(1) should be the number of overall pages containing the links
        example: PAGES_PATTERN = r'Pages: (\d+)'

    and its loadPage method:

      def load_page(self, page_n):
          return the html of the page number page_n
    """

    DIRECT_LINK  = True
    LEECH_HOSTER = False

    #@TODO: Remove in 0.4.10
    def _setup(self):
        orig_name = self.__name__
        self.__name__ = (orig_name + ".py").replace("Folder.py", "").replace(".py", "")

        super(SimpleCrypter, self)._setup()

        self.__name__ = orig_name

    #@TODO: Remove in 0.4.10
    def load_account(self):
        orig_name = self.__name__
        self.__name__ = (orig_name + ".py").replace("Folder.py", "").replace(".py", "")

        super(SimpleCrypter, self).load_account()

        self.__name__ = orig_name

    def handle_direct(self, pyfile):
        for i in xrange(self.get_config("maxredirs", plugin="UserAgentSwitcher")):
            redirect = self.link or pyfile.url
            self.log_debug("Redirect #%d to: %s" % (i, redirect))

            header = self.load(redirect, just_header=True)
            if 'location' in header and header['location']:
                self.link = header['location']
            else:
                break
        else:
            self.log_error(_("Too many redirects"))

    def decrypt(self, pyfile):
        self.prepare()
        self.check_info()  #@TODO: Remove in 0.4.10

        if self.direct_dl:
            self.log_debug(_("Looking for direct download link..."))
            self.handle_direct(pyfile)

            if self.link or self.links or self.urls or self.packages:
                self.log_info(_("Direct download link detected"))
            else:
                self.log_info(_("Direct download link not found"))

        if not (self.link or self.links or self.urls or self.packages):
            self.preload()

            self.links = self.get_links() or list()

            if hasattr(self, 'PAGES_PATTERN') and hasattr(self, 'loadPage'):
                self.handle_pages(pyfile)

            self.log_debug("Package has %d links" % len(self.links))

        if self.link:
            self.urls.append(self.link)

        if self.links:
            name = folder = pyfile.name
            self.packages.append((name, self.links, folder))

    def get_links(self):
        """
        Returns the links extracted from self.html
        You should override this only if it's impossible to extract links using only the LINK_PATTERN.
        """
        return re.findall(self.LINK_PATTERN, self.html)

    def handle_pages(self, pyfile):
        try:
            pages = int(re.search(self.PAGES_PATTERN, self.html).group(1))
        except Exception:
            pages = 1

        for p in xrange(2, pages + 1):
            self.html = self.load_page(p)
            self.links += self.get_links()
<commit_msg>fix #1765<commit_after># -*- coding: utf-8 -*-

import re

from module.plugins.internal.Crypter import Crypter
from module.plugins.internal.SimpleHoster import SimpleHoster, create_getInfo, replace_patterns, set_cookie, set_cookies
from module.utils import fixup, html_unescape

class SimpleCrypter(Crypter, SimpleHoster):
    __name__    = "SimpleCrypter"
    __type__    = "crypter"
    __version__ = "0.61"
    __status__  = "testing"

    __pattern__ = r'^unmatchable$'
    __config__  = [("use_subfolder"     , "bool", "Save package to subfolder"          , True),  #: Overrides pyload.config['general']['folder_per_package']
                   ("subfolder_per_pack", "bool", "Create a subfolder for each package", True)]

    __description__ = """Simple decrypter plugin"""
    __license__     = "GPLv3"
    __authors__     = [("Walter Purcaro", "vuolter@gmail.com")]

    """
    Following patterns should be defined by each crypter:

      LINK_PATTERN: Download link or regex to catch links in group(1)
        example: LINK_PATTERN = r'
<a href="(.+?)"'

      NAME_PATTERN: (optional) folder name or page title
        example: NAME_PATTERN = r'<title>Files of: (?P<N>[^<]+) folder</title>'

      OFFLINE_PATTERN: (optional) Checks if the page is unreachable
        example: OFFLINE_PATTERN = r'File (deleted|not found)'

      TEMP_OFFLINE_PATTERN: (optional) Checks if the page is temporarily unreachable
        example: TEMP_OFFLINE_PATTERN = r'Server maintainance'

    You can override the getLinks method if you need a more sophisticated way to extract the links.

    If the links are splitted on multiple pages you can define the PAGES_PATTERN regex:

      PAGES_PATTERN: (optional) group(1) should be the number of overall pages containing the links
        example: PAGES_PATTERN = r'Pages: (\d+)'

    and its loadPage method:

      def load_page(self, page_n):
          return the html of the page number page_n
    """

    DIRECT_LINK  = True
    LEECH_HOSTER = False

    #@TODO: Remove in 0.4.10
    def _setup(self):
        orig_name = self.__name__
        self.__name__ = (orig_name + ".py").replace("Folder.py", "").replace(".py", "")

        super(SimpleCrypter, self)._setup()

        self.__name__ = orig_name

    #@TODO: Remove in 0.4.10
    def load_account(self):
        orig_name = self.__name__
        self.__name__ = (orig_name + ".py").replace("Folder.py", "").replace(".py", "")

        super(SimpleCrypter, self).load_account()

        self.__name__ = orig_name

    def handle_direct(self, pyfile):
        for i in xrange(self.get_config("maxredirs", plugin="UserAgentSwitcher")):
            redirect = self.link or pyfile.url
            self.log_debug("Redirect #%d to: %s" % (i, redirect))

            header = self.load(redirect, just_header=True)
            if 'location' in header and header['location']:
                self.link = header['location']
            else:
                break
        else:
            self.log_error(_("Too many redirects"))

    def decrypt(self, pyfile):
        self.links = []    #@TODO: Recheck in 0.4.10
        self.prepare()
        self.check_info()  #@TODO: Remove in 0.4.10

        if self.direct_dl:
            self.log_debug(_("Looking for direct download link..."))
            self.handle_direct(pyfile)

            if self.link or self.links or self.urls or self.packages:
                self.log_info(_("Direct download link detected"))
            else:
                self.log_info(_("Direct download link not found"))

        if not (self.link or self.links or self.urls or self.packages):
            self.preload()

            self.links = self.get_links() or list()

            if hasattr(self, 'PAGES_PATTERN') and hasattr(self, 'loadPage'):
                self.handle_pages(pyfile)

            self.log_debug("Package has %d links" % len(self.links))

        if self.link:
            self.urls.append(self.link)

        if self.links:
            name = folder = pyfile.name
            self.packages.append((name, self.links, folder))

    def get_links(self):
        """
        Returns the links extracted from self.html
        You should override this only if it's impossible to extract links using only the LINK_PATTERN.
        """
        return re.findall(self.LINK_PATTERN, self.html)

    def handle_pages(self, pyfile):
        try:
            pages = int(re.search(self.PAGES_PATTERN, self.html).group(1))
        except Exception:
            pages = 1

        for p in xrange(2, pages + 1):
            self.html = self.load_page(p)
            self.links += self.get_links()
<|endoftext|>


<commit_before>"""
Base class for equipment that use message-based communication.
"""
import time

from .connection import Connection
from .exceptions import MSLTimeoutError
from .constants import LF, CR

class ConnectionMessageBased(Connection):

    CR = CR
    """:class:`bytes`: The carriage-return character."""

    LF = LF
    """:class:`bytes`: The line-feed character."""

    def __init__(self, record):
        """Base class for equipment that use message-based communication.

        The :data:`~msl.equipment.record_types.ConnectionRecord.backend`
        value must be equal to :data:`~msl.equipment.constants.Backend.MSL`
        to use this class for the communication system. This is achieved by
        setting the value in the **Backend** field for a connection record
        in the :ref:`connections-database` to be ``MSL``.

        Do not instantiate this class directly. Use the
        :meth:`~.EquipmentRecord.connect` method to connect to the equipment.

        Parameters
        ----------
        record : :class:`.EquipmentRecord`
            A record from an :ref:`equipment-database`.
        """
        super(ConnectionMessageBased, self).__init__(record)

        self._encoding = 'utf-8'
        self._encoding_errors = 'strict'
        self._read_termination = ConnectionMessageBased.LF
        self._write_termination = ConnectionMessageBased.CR + ConnectionMessageBased.LF
        self._max_read_size = 2 ** 16
        self._timeout = None

    @property
    def encoding(self):
        """:class:`str`: The encoding that is used for :meth:`read` and :meth:`write` operations."""
        return self._encoding

    @encoding.setter
    def encoding(self, encoding):
        """Set the encoding to use for :meth:`read` and :meth:`write` operations."""
        _ = 'test encoding'.encode(encoding).decode(encoding)

        # re-encoding the read/write termination values ensure that the termination
        # sequence can be encoded using the new encoding
        if self._read_termination is not None:
            read_term = self._read_termination.decode(self._encoding)
        if self._write_termination is not None:
            write_term = self._write_termination.decode(self._encoding)

        self._encoding = encoding

        if self._read_termination is not None:
            self.read_termination = read_term
        if self._write_termination is not None:
            self.write_termination = write_term

    @property
    def read_termination(self):
        """:class:`bytes` or :data:`None`: The termination character sequence
        that is used for the :meth:`read` method.

        Reading stops when the equipment stops sending data or the `read_termination`
        character sequence is detected. If you set the `read_termination` to be equal
        to a variable of type :class:`str` it will automatically be encoded.
        """
        return self._read_termination

    @read_termination.setter
    def read_termination(self, termination):
        self._read_termination = self._set_termination_encoding(termination)

    @property
    def write_termination(self):
        """:class:`bytes` or :data:`None`: The termination character sequence that
        is appended to :meth:`write` messages.

        If you set the `write_termination` to be equal to a variable of type
        :class:`str` it will automatically be encoded.
        """
        return self._write_termination

    @write_termination.setter
    def write_termination(self, termination):
        self._write_termination = self._set_termination_encoding(termination)

    @property
    def max_read_size(self):
        """:class:`int`: The maximum number of bytes that can be :meth:`read`."""
        return self._max_read_size

    @max_read_size.setter
    def max_read_size(self, size):
        """The maximum number of bytes that can be :meth:`read`."""
        max_size = int(size)
        if max_size < 1:
            raise ValueError('The maximum number of bytes to read must be > 0, got {}'.format(size))
        self._max_read_size = max_size

    @property
    def timeout(self):
        """:class:`float` or :data:`None`: The timeout, in seconds, for :meth:`read` and :meth:`write` operations."""
        return self._timeout

    def raise_timeout(self, append_msg=''):
        """Raise a :exc:`~.exceptions.MSLTimeoutError`.

        Parameters
        ----------
        append_msg : :class:`str`, optional
            A message to append to the generic timeout message.
        """
        msg = 'Timeout occurred after {} seconds'.format(self.timeout)
        if append_msg:
            msg += str(append_msg)
        self.log_error('{!r} {}'.format(self, msg))
        raise MSLTimeoutError('{!r}\n{}'.format(self, msg))

    def read(self, size=None):
        """Read the response from the equipment.

        .. attention::
           The subclass must override this method.

        Parameters
        ----------
        size : :class:`int`, optional
            The number of bytes to read.

        Returns
        -------
        :class:`str`
            The response from the equipment.
        """
        raise NotImplementedError

    def write(self, msg):
        """Write a message to the equipment.

        .. attention::
           The subclass must override this method.

        Parameters
        ----------
        msg : :class:`str`
            The message to write to the equipment.

        Returns
        -------
        :class:`int`
            The number of bytes written.
        """
        raise NotImplementedError

    def query(self, msg, delay=0.0, size=None):
        """Convenience method for performing a :meth:`write` followed by a :meth:`read`.

        Parameters
        ----------
        msg : :class:`str`
            The message to write to the equipment.
        delay : :class:`float`, optional
            The time delay, in seconds, to wait between :meth:`write` and
            :meth:`read` operations.
        size : :class:`int`, optional
            The number of bytes to read.

        Returns
        -------
        :class:`str`
            The response from the equipment.
        """
        self.write(msg)
        if delay > 0.0:
            time.sleep(delay)
        return self.read(size=size)

    def _set_timeout_value(self, value):
        # convenience method for setting the timeout value
        if value is not None:
            self._timeout = float(value)
            if self._timeout == 0:
                self._timeout = None
            elif self._timeout < 0:
                raise ValueError('Not a valid timeout value: {}'.format(value))
        else:
            self._timeout = None

    def _set_termination_encoding(self, termination):
        # convenience method for setting the termination encoding
        try:
            return termination.encode(self._encoding)
        except AttributeError:
            return termination  # `termination` is already encoded

    def _encode(self, message):
        # convenience method for preparing the message for a write operation
        if isinstance(message, bytes):
            data = message
        else:
            data = message.encode(encoding=self._encoding, errors=self._encoding_errors)
        if self._write_termination is not None and not data.endswith(self._write_termination):
            data += self._write_termination
        self.log_debug('{}.write({!r})'.format(self, data))
        return data

    def _decode(self, size, message):
        # convenience method for processing the message from a read operation
        if size is None:
            self.log_debug('{}.read() -> {!r}'.format(self, message))
        else:
            self.log_debug('{}.read({}) -> {!r}'.format(self, size, message))
        return message.decode(encoding=self._encoding, errors=self._encoding_errors)
<commit_msg>create timeout.setter for ConnectionMessageBased<commit_after>"""
Base class for equipment that use message-based communication.
"""
import time

from .connection import Connection
from .exceptions import MSLTimeoutError
from .constants import LF, CR

class ConnectionMessageBased(Connection):

    CR = CR
    """:class:`bytes`: The carriage-return character."""

    LF = LF
    """:class:`bytes`: The line-feed character."""

    def __init__(self, record):
        """Base class for equipment that use message-based communication.

        The :data:`~msl.equipment.record_types.ConnectionRecord.backend`
        value must be equal to :data:`~msl.equipment.constants.Backend.MSL`
        to use this class for the communication system. This is achieved by
        setting the value in the **Backend** field for a connection record
        in the :ref:`connections-database` to be ``MSL``.

        Do not instantiate this class directly. Use the
        :meth:`~.EquipmentRecord.connect` method to connect to the equipment.

        Parameters
        ----------
        record : :class:`.EquipmentRecord`
            A record from an :ref:`equipment-database`.
        """
        super(ConnectionMessageBased, self).__init__(record)

        self._encoding = 'utf-8'
        self._encoding_errors = 'strict'
        self._read_termination = ConnectionMessageBased.LF
        self._write_termination = ConnectionMessageBased.CR + ConnectionMessageBased.LF
        self._max_read_size = 2 ** 16
        self._timeout = None

    @property
    def encoding(self):
        """:class:`str`: The encoding that is used for :meth:`read` and :meth:`write` operations."""
        return self._encoding

    @encoding.setter
    def encoding(self, encoding):
        """Set the encoding to use for :meth:`read` and :meth:`write` operations."""
        _ = 'test encoding'.encode(encoding).decode(encoding)

        # re-encoding the read/write termination values ensure that the termination
        # sequence can be encoded using the new encoding
        if self._read_termination is not None:
            read_term = self._read_termination.decode(self._encoding)
        if self._write_termination is not None:
            write_term = self._write_termination.decode(self._encoding)

        self._encoding = encoding

        if self._read_termination is not None:
            self.read_termination = read_term
        if self._write_termination is not None:
            self.write_termination = write_term

    @property
    def read_termination(self):
        """:class:`bytes` or :data:`None`: The termination character sequence
        that is used for the :meth:`read` method.

        Reading stops when the equipment stops sending data or the `read_termination`
        character sequence is detected. If you set the `read_termination` to be equal
        to a variable of type :class:`str` it will automatically be encoded.
        """
        return self._read_termination

    @read_termination.setter
    def read_termination(self, termination):
        self._read_termination = self._set_termination_encoding(termination)

    @property
    def write_termination(self):
        """:class:`bytes` or :data:`None`: The termination character sequence that
        is appended to :meth:`write` messages.

        If you set the `write_termination` to be equal to a variable of type
        :class:`str` it will automatically be encoded.
        """
        return self._write_termination

    @write_termination.setter
    def write_termination(self, termination):
        self._write_termination = self._set_termination_encoding(termination)

    @property
    def max_read_size(self):
        """:class:`int`: The maximum number of bytes that can be :meth:`read`."""
        return self._max_read_size

    @max_read_size.setter
    def max_read_size(self, size):
        """The maximum number of bytes that can be :meth:`read`."""
        max_size = int(size)
        if max_size < 1:
            raise ValueError('The maximum number of bytes to read must be > 0, got {}'.format(size))
        self._max_read_size = max_size

    @property
    def timeout(self):
        """:class:`float` or :data:`None`: The timeout, in seconds, for :meth:`read` and :meth:`write` operations."""
        return self._timeout

    @timeout.setter
    def timeout(self, value):
        self._timeout = value

    def raise_timeout(self, append_msg=''):
        """Raise a :exc:`~.exceptions.MSLTimeoutError`.

        Parameters
        ----------
        append_msg : :class:`str`, optional
            A message to append to the generic timeout message.
        """
        msg = 'Timeout occurred after {} seconds'.format(self.timeout)
        if append_msg:
            msg += str(append_msg)
        self.log_error('{!r} {}'.format(self, msg))
        raise MSLTimeoutError('{!r}\n{}'.format(self, msg))

    def read(self, size=None):
        """Read the response from the equipment.

        .. attention::
           The subclass must override this method.

        Parameters
        ----------
        size : :class:`int`, optional
            The number of bytes to read.

        Returns
        -------
        :class:`str`
            The response from the equipment.
        """
        raise NotImplementedError

    def write(self, msg):
        """Write a message to the equipment.

        .. attention::
           The subclass must override this method.

        Parameters
        ----------
        msg : :class:`str`
            The message to write to the equipment.

        Returns
        -------
        :class:`int`
            The number of bytes written.
        """
        raise NotImplementedError

    def query(self, msg, delay=0.0, size=None):
        """Convenience method for performing a :meth:`write` followed by a :meth:`read`.

        Parameters
        ----------
        msg : :class:`str`
            The message to write to the equipment.
        delay : :class:`float`, optional
            The time delay, in seconds, to wait between :meth:`write` and
            :meth:`read` operations.
        size : :class:`int`, optional
            The number of bytes to read.

        Returns
        -------
        :class:`str`
            The response from the equipment.
        """
        self.write(msg)
        if delay > 0.0:
            time.sleep(delay)
        return self.read(size=size)

    def _set_timeout_value(self, value):
        # convenience method for setting the timeout value
        if value is not None:
            self._timeout = float(value)
            if self._timeout == 0:
                self._timeout = None
            elif self._timeout < 0:
                raise ValueError('Not a valid timeout value: {}'.format(value))
        else:
            self._timeout = None

    def _set_termination_encoding(self, termination):
        # convenience method for setting the termination encoding
        try:
            return termination.encode(self._encoding)
        except AttributeError:
            return termination  # `termination` is already encoded

    def _encode(self, message):
        # convenience method for preparing the message for a write operation
        if isinstance(message, bytes):
            data = message
        else:
            data = message.encode(encoding=self._encoding, errors=self._encoding_errors)
        if self._write_termination is not None and not data.endswith(self._write_termination):
            data += self._write_termination
        self.log_debug('{}.write({!r})'.format(self, data))
        return data

    def _decode(self, size, message):
        # convenience method for processing the message from a read operation
        if size is None:
            self.log_debug('{}.read() -> {!r}'.format(self, message))
        else:
            self.log_debug('{}.read({}) -> {!r}'.format(self, size, message))
        return message.decode(encoding=self._encoding, errors=self._encoding_errors)
<|endoftext|>


<commit_before>#
# Collective Knowledge (program - features)
#
# See CK LICENSE.txt for licensing details
# See CK COPYRIGHT.txt for copyright details
#
# Developer: Grigori Fursin, Grigori.Fursin@cTuning.org, http://fursin.net
#

cfg={}  # Will be updated by CK (meta description of this module)
work={} # Will be updated by CK (temporal data)
ck=None # Will be updated by CK (initialized CK kernel)

# Local settings

##############################################################################
# Initialize module

def init(i):
    """

    Input:  {}

    Output: {
              return       - return code =  0, if successful
                                         >  0, if error
              (error)      - error text if return > 0
            }

    """
    return {'return':0}

##############################################################################
# extract program static milepost features

def extract(i):
    """
    Input:  {
              (repo_uoa)        - repository UOA
              (data_uoa)        - program UOA (can be wildcards)

              (tags)            - tags to process specific programs

              (target_repo_uoa) - repo, where to save features - if =='', use repo_uoa

            }

    Output: {
              return       - return code =  0, if successful
                                         >  0, if error
              (error)      - error text if return > 0

              dict         - final dict with key 'features'={...}
            }

    """

    import os
    import json

    o=i.get('out','')
    oo=''
    if o=='con':oo=o

    muoa=cfg['module_deps']['program']
    duoa=i.get('data_uoa','')
    ruoa=i.get('repo_uoa','')

    truoa=i.get('target_repo_uoa','')

    tags=i.get('tags','')

    rx=ck.access({'action':'search',
                  'repo_uoa':ruoa,
                  'module_uoa':muoa,
                  'data_uoa':duoa,
                  'tags':tags})
    if rx['return']>0: return rx

    lst=rx['lst']

    feat1={}

    num=0
    for q in lst:
        num+=1

        duid=q['data_uid']
        duoa=q['data_uoa']

        xtruoa=truoa
        if truoa=='':
           xtruoa=q['repo_uid']

        if o=='con':
           ck.out(str(num)+') Processing '+duoa+' ...')

        ii={'action':'pipeline',
            'module_uoa':muoa,
            'data_uoa':duoa,
            'milepost':'yes',
            'out':oo,
            'quiet':'yes',
            'skip_info_collection':'yes',
            'no_run':'yes'}
        r=ck.access(ii)
        if r['return']>0:
           if o=='con':
              ck.out('')
              ck.out('CK WARNING: pipeline failed ('+r['error']+')')
              ck.out('')
        else:
           feat=r.get('features',{}).get('program_static_milepost_features',{})

           if len(feat)>0:

              ddd={}

              found=False
              ry=ck.access({'action':'load',
                            'module_uoa':work['self_module_uid'],
                            'data_uoa':duid})
              if ry['return']==0:
                 ddd=ry['dict']
                 found=True

              feat1=ddd.get('features',{}).get('program_static_milepost_features',{})
              rz=ck.merge_dicts({'dict1':feat1, 'dict2':feat})
              if rz['return']>0: return rz
              feat1=rz['dict1']

              if 'features' not in ddd: ddd['features']={}

              ddd['features']['program_static_milepost_features']=feat1

              ii={}
              ii['action']='add'
              if found: ii['action']='update'
              ii['module_uoa']=work['self_module_uid']
              ii['data_uoa']=duoa
              ii['data_uid']=duid
              ii['repo_uoa']=xtruoa
              ii['dict']=ddd
              ii['substitute']='yes'
              ry=ck.access(ii)
              if ry['return']>0: return ry

    return {'return':0, 'dict':{'features':feat1}}

##############################################################################
# Calculate similarity between programs based on features
#
# FGG: For now just simple Euclidean distance, however it has many drawbacks
# since it doesn't take into account "importance" of different features, etc
# We plan to use more SVN, decision trees, etc (see our papers) ...

def calculate_similarity(i):
    """
    Input:  {
              features1 - MILEPOST features for program 1
              features2 - MILEPOST features for program 2
            }

    Output: {
              return       - return code =  0, if successful
                                         >  0, if error
              (error)      - error text if return > 0

              distance     - basic Euclidean distance (can be > 1)
            }

    """

    import math

    prog1=i['features1']
    prog2=i['features2']

    # Check which feature to use for normalization
    fd=cfg['milepost_features_description']
    ftn=cfg['milepost_normalization_feature']

    # Check that in vectors
    distance=None # if error

    if ftn in prog1 and ftn in prog2:

       x1=prog1.get(ftn,None)
       x2=prog2.get(ftn,None)

       if x1!=None and x1!='' and x2!=None and x2!='':
          ftnp1=float(x1)
          ftnp2=float(x2)

          distance=0.0

          for q2 in fd:
              q=fd[q2]

              ftp1=float(prog1.get(q2,0))
              ftp2=float(prog2.get(q2,0))

              if q.get('use_for_euclidean_distance','')=='yes':
                 if q.get('normalized','')!='yes':
                    ftp1/=ftnp1
                    ftp2/=ftnp2

                 distance+=(ftp1-ftp2)*(ftp1-ftp2)

          if distance!=0.0:
             distance=math.sqrt(distance)

    return {'return':0, 'distance':distance}
<commit_msg>minor improvement to feature extractor<commit_after>#
# Collective Knowledge (program - features)
#
# See CK LICENSE.txt for licensing details
# See CK COPYRIGHT.txt for copyright details
#
# Developer: Grigori Fursin, Grigori.Fursin@cTuning.org, http://fursin.net
#

cfg={}  # Will be updated by CK (meta description of this module)
work={} # Will be updated by CK (temporal data)
ck=None # Will be updated by CK (initialized CK kernel)

# Local settings

##############################################################################
# Initialize module

def init(i):
    """

    Input:  {}

    Output: {
              return       - return code =  0, if successful
                                         >  0, if error
              (error)      - error text if return > 0
            }

    """
    return {'return':0}

##############################################################################
# extract program static milepost features

def extract(i):
    """
    Input:  {
              (repo_uoa)        - repository UOA
              (data_uoa)        - program UOA (can be wildcards)

              (tags)            - tags to process specific programs

              (target_repo_uoa) - repo, where to save features - if =='', use repo_uoa

            }

    Output: {
              return       - return code =  0, if successful
                                         >  0, if error
              (error)      - error text if return > 0

              dict         - final dict with key 'features'={...}
            }

    """

    import os
    import json

    o=i.get('out','')
    oo=''
    if o=='con':oo=o

    muoa=cfg['module_deps']['program']
    duoa=i.get('data_uoa','')
    ruoa=i.get('repo_uoa','')

    truoa=i.get('target_repo_uoa','')

    tags=i.get('tags','')

    rx=ck.access({'action':'search',
                  'repo_uoa':ruoa,
                  'module_uoa':muoa,
                  'data_uoa':duoa,
                  'tags':tags})
    if rx['return']>0: return rx

    lst=rx['lst']

    feat1={}

    num=0
    for q in lst:
        num+=1

        duid=q['data_uid']
        duoa=q['data_uoa']

        xtruoa=truoa
        if truoa=='':
           xtruoa=q['repo_uid']

        if o=='con':
           ck.out(str(num)+') Processing '+duoa+' ...')

        ii={'action':'pipeline',
            'module_uoa':muoa,
            'data_uoa':duoa,
            'milepost':'yes',
            'out':oo,
            'quiet':'yes',
            'skip_info_collection':'yes',
            'no_run':'yes'}
        r=ck.access(ii)

        finish=False
        if r['return']>0:
           finish=True
           if o=='con':
              ck.out('')
              ck.out('CK WARNING: pipeline failed ('+r['error']+')')
              ck.out('')

              rx=ck.inp({'text':'Would you still like to record some extracted features (Y/n)? '})
              x=rx['string'].strip()
              if x!='n' and x!='N':
                 finish=False

        if not finish:
           feat=r.get('features',{}).get('program_static_milepost_features',{})

           if len(feat)>0:

              ddd={}

              found=False
              ry=ck.access({'action':'load',
                            'module_uoa':work['self_module_uid'],
                            'data_uoa':duid})
              if ry['return']==0:
                 ddd=ry['dict']
                 found=True

              feat1=ddd.get('features',{}).get('program_static_milepost_features',{})
              rz=ck.merge_dicts({'dict1':feat1, 'dict2':feat})
              if rz['return']>0: return rz
              feat1=rz['dict1']

              if 'features' not in ddd: ddd['features']={}

              ddd['features']['program_static_milepost_features']=feat1

              ii={}
              ii['action']='add'
              if found: ii['action']='update'
              ii['module_uoa']=work['self_module_uid']
              ii['data_uoa']=duoa
              ii['data_uid']=duid
              ii['repo_uoa']=xtruoa
              ii['dict']=ddd
              ii['substitute']='yes'
              ry=ck.access(ii)
              if ry['return']>0: return ry

    return {'return':0, 'dict':{'features':feat1}}

##############################################################################
# Calculate similarity between programs based on features
#
# FGG: For now just simple Euclidean distance, however it has many drawbacks
# since it doesn't take into account "importance" of different features, etc
# We plan to use more SVN, decision trees, etc (see our papers) ...

def calculate_similarity(i):
    """
    Input:  {
              features1 - MILEPOST features for program 1
              features2 - MILEPOST features for program 2
            }

    Output: {
              return       - return code =  0, if successful
                                         >  0, if error
              (error)      - error text if return > 0

              distance     - basic Euclidean distance (can be > 1)
            }

    """

    import math

    prog1=i['features1']
    prog2=i['features2']

    # Check which feature to use for normalization
    fd=cfg['milepost_features_description']
    ftn=cfg['milepost_normalization_feature']

    # Check that in vectors
    distance=None # if error

    if ftn in prog1 and ftn in prog2:

       x1=prog1.get(ftn,None)
       x2=prog2.get(ftn,None)

       if x1!=None and x1!='' and x2!=None and x2!='':
          ftnp1=float(x1)
          ftnp2=float(x2)

          distance=0.0

          for q2 in fd:
              q=fd[q2]

              ftp1=float(prog1.get(q2,0))
              ftp2=float(prog2.get(q2,0))

              if q.get('use_for_euclidean_distance','')=='yes':
                 if q.get('normalized','')!='yes':
                    ftp1/=ftnp1
                    ftp2/=ftnp2

                 distance+=(ftp1-ftp2)*(ftp1-ftp2)

          if distance!=0.0:
             distance=math.sqrt(distance)

    return {'return':0, 'distance':distance}
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""This program inputs the map of elementary school teachers to hub names.
It assumes the presence of a file called 'hub_map.csv' in the directory from
which the program is executed.
It also assumes 'hub_map.csv' was created by saving an Excel spreadsheet as
a "Windows Comma Separated (.csv)" file (Excel offers other CSV options that
do not work with this program).
"""

def ReadMap():
    map_d = {}   # empty dictionary
    count = 0

    try:
        open_file = open('hub_map.csv')

        for line in open_file:
            fields = line.split('|')
            ## be sure to strip the '\r\n' that Excel adds at the end of lines
            map_d.update({fields[0]:fields[-1][:-2]})
            count += 1

        print "%d hub IDs read and processed." % count

    finally:
        open_file.close()

    return map_d

def IsInClassroomHub(map_d, hub_text):
    # return True if hub text is one of the Middle School grade hubs
    if hub_text in ('6', '7', '8'):
        return True

    # return True if teacher's name starts the hub text
    for teacher in map_d.keys():
        if hub_text[:len(teacher)] == teacher:
            return True

    # return False otherwise
    return False

def IsAnyHubClassroomHub(map_d, hub_field):
    hubs = hub_field.split(';')
    for hub in hubs:
        if IsInClassroomHub(map_d, hub):
            return True

    return False

def PrintMap(map_d):
    print map_d

def main():
    map_d = ReadMap()
    PrintMap(map_d)

if __name__ == '__main__':
    main()
<commit_msg>Add SDC hub check<commit_after>#!/usr/bin/env python
"""This program inputs the map of elementary school teachers to hub names.
It assumes the presence of a file called 'hub_map.csv' in the directory from
which the program is executed.
It also assumes 'hub_map.csv' was created by saving an Excel spreadsheet as
a "Windows Comma Separated (.csv)" file (Excel offers other CSV options that
do not work with this program).
"""

def ReadMap():
    map_d = {}   # empty dictionary
    count = 0

    try:
        open_file = open('hub_map.csv')

        for line in open_file:
            fields = line.split('|')
            ## be sure to strip the '\r\n' that Excel adds at the end of lines
            map_d.update({fields[0]:fields[-1][:-2]})
            count += 1

        print "%d hub IDs read and processed." % count

    finally:
        open_file.close()

    return map_d

def IsInClassroomHub(map_d, hub_text):
    # return True if hub text is one of the Middle School grade hubs
    if hub_text in ('6th Grade', '7th Grade', '8th Grade', 'SDC'):
        return True

    # return True if teacher's name starts the hub text
    for teacher in map_d.keys():
        if hub_text[:len(teacher)] == teacher:
            return True

    # return False otherwise
    return False

def IsAnyHubClassroomHub(map_d, hub_field):
    hubs = hub_field.split(';')
    for hub in hubs:
        if IsInClassroomHub(map_d, hub):
            return True

    return False

def PrintMap(map_d):
    print map_d

def main():
    map_d = ReadMap()
    PrintMap(map_d)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>import functools
import io
import logging
import math
import os
import warnings

logger = logging.getLogger("fsspec")

class BaseCache(object):
    """Pass-though cache: doesn't keep anything, calls every time

    Acts as base class for other cachers

    Parameters
    ----------
    blocksize: int
        How far to read ahead in numbers of bytes
    fetcher: func
        Function of the form f(start, end) which gets bytes from remote as
        specified
    size: int
        How big this file is
    """

    name = "none"

    def __init__(self, blocksize, fetcher, size):
        self.blocksize = blocksize
        self.fetcher = fetcher
        self.size = size

    def _fetch(self, start, stop):
        if start is None:
            start = 0
        if stop is None:
            stop = self.size
        if start >= self.size or start >= stop:
            return b""
        return self.fetcher(start, stop)

class MMapCache(BaseCache):
    """memory-mapped sparse file cache

    Opens temporary file, which is filled blocks-wise when data is requested.
    Ensure there is enough disc space in the temporary location.

    This cache method might only work on posix
    """

    name = "mmap"

    def __init__(self, blocksize, fetcher, size, location=None, blocks=None):
        super().__init__(blocksize, fetcher, size)
        self.blocks = set() if blocks is None else blocks
        self.location = location
        self.cache = self._makefile()

    def _makefile(self):
        import mmap
        import tempfile

        if self.size == 0:
            return bytearray()

        # posix version
        if self.location is None or not os.path.exists(self.location):
            if self.location is None:
                fd = tempfile.TemporaryFile()
                self.blocks = set()
            else:
                fd = io.open(self.location, "wb+")
            fd.seek(self.size - 1)
            fd.write(b"1")
            fd.flush()
        else:
            fd = io.open(self.location, "rb+")

        return mmap.mmap(fd.fileno(), self.size)

    def _fetch(self, start, end):
        logger.debug(f"MMap cache fetching {start}-{end}")
        if start is None:
            start = 0
        if end is None:
            end = self.size
        if start >= self.size or start >= end:
            return b""
        start_block = start // self.blocksize
        end_block = end // self.blocksize
        need = [i for i in range(start_block, end_block + 1) if i not in self.blocks]
        while need:
            # TODO: not a for loop so we can consolidate blocks later to
            # make fewer fetch calls; this could be parallel
            i = need.pop(0)
            sstart = i * self.blocksize
            send = min(sstart + self.blocksize, self.size)
            logger.debug(f"MMap get block #{i} ({sstart}-{send}")
            self.cache[sstart:send] = self.fetcher(sstart, send)
            self.blocks.add(i)

        return self.cache[start:end]

    def __getstate__(self):
        state = self.__dict__.copy()
        # Remove the unpicklable entries.
        del state["cache"]
        return state

    def __setstate__(self, state):
        # Restore instance attributes
        self.__dict__.update(state)
        self.cache = self._makefile()

class ReadAheadCache(BaseCache):
    """Cache which reads only when we get beyond a block of data

    This is a much simpler version of BytesCache, and does not attempt to
    fill holes in the cache or keep fragments alive. It is best suited to
    many small reads in a sequential order (e.g., reading lines from a file).
    """

    name = "readahead"

    def __init__(self, blocksize, fetcher, size):
        super().__init__(blocksize, fetcher, size)
        self.cache = b""
        self.start = 0
        self.end = 0

    def _fetch(self, start, end):
        if start is None:
            start = 0
        if end is None or end > self.size:
            end = self.size
        if start >= self.size or start >= end:
            return b""
        l = end - start
        if start >= self.start and end <= self.end:
            # cache hit
            return self.cache[start - self.start : end - self.start]
        elif self.start <= start < self.end:
            # partial hit
            part = self.cache[start - self.start :]
            l -= len(part)
            start = self.end
        else:
            # miss
            part = b""
        end = min(self.size, end + self.blocksize)
        self.cache = self.fetcher(start, end)  # new block replaces old
        self.start = start
        self.end = self.start + len(self.cache)
        return part + self.cache[:l]

class FirstChunkCache(BaseCache):
    """Caches the first block of a file only

    This may be useful for file types where the metadata is stored in the header,
    but is randomly accessed.
    """

    name = "first"

    def __init__(self, blocksize, fetcher, size):
        super().__init__(blocksize, fetcher, size)
        self.cache = None

    def _fetch(self, start, end):
        start = start or 0
        end = end or self.size
        if start < self.blocksize:
            if self.cache is None:
                if end > self.blocksize:
                    data = self.fetcher(0, end)
                    self.cache = data[: self.blocksize]
                    return data[start:]
                self.cache = self.fetcher(0, self.blocksize)
            part = self.cache[start:end]
            if end > self.blocksize:
                part += self.fetcher(self.blocksize, end)
            return part
        else:
            return self.fetcher(start, end)

class BlockCache(BaseCache):
    """
    Cache holding memory as a set of blocks.

    Requests are only ever made `blocksize` at a time, and are
    stored in an LRU cache. The least recently accessed block is
    discarded when more than `maxblocks` are stored.

    Parameters
    ----------
    blocksize : int
        The number of bytes to store in each block.
        Requests are only ever made for `blocksize`, so this
        should balance the overhead of making a request against
        the granularity of the blocks.
    fetcher : Callable
    size : int
        The total size of the file being cached.
    maxblocks : int
        The maximum number of blocks to cache for. The maximum memory
        use for this cache is then ``blocksize * maxblocks``.
    """

    name = "blockcache"

    def __init__(self, blocksize, fetcher, size, maxblocks=32):
        super().__init__(blocksize, fetcher, size)
        self.nblocks = math.ceil(size / blocksize)
        self.maxblocks = maxblocks
        self._fetch_block_cached = functools.lru_cache(maxblocks)(self._fetch_block)

    def __repr__(self):
        return "<BlockCache blocksize={}, size={}, nblocks={}>".format(
            self.blocksize, self.size, self.nblocks
        )

    def cache_info(self):
        """
        The statistics on the block cache.

        Returns
        ----------
        NamedTuple
            Returned directly from the LRU Cache used internally.
        """
        return self._fetch_block_cached.cache_info()

    def __getstate__(self):
        state = self.__dict__
        del state["_fetch_block_cached"]
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self._fetch_block_cached = functools.lru_cache(state["maxblocks"])(
            self._fetch_block
        )

    def _fetch(self, start, end):
        if start is None:
            start = 0
        if end is None:
            end = self.size
        if start >= self.size or start >= end:
            return b""

        # byte position -> block numbers
        start_block_number = start // self.blocksize
        end_block_number = end // self.blocksize

        # these are cached, so safe to do multiple calls for the same start and end.
        for block_number in range(start_block_number, end_block_number + 1):
            self._fetch_block_cached(block_number)

        return self._read_cache(
            start,
            end,
            start_block_number=start_block_number,
            end_block_number=end_block_number,
        )

    def _fetch_block(self, block_number):
        """
        Fetch the block of data for `block_number`.
        """
        if block_number > self.nblocks:
            raise ValueError(
                "'block_number={}' is greater than the number of blocks ({})".format(
                    block_number, self.nblocks
                )
            )

        start = block_number * self.blocksize
        end = start + self.blocksize
        logger.info("BlockCache fetching block %d", block_number)
        block_contents = super()._fetch(start, end)
        return block_contents

    def _read_cache(self, start, end, start_block_number, end_block_number):
        """
        Read from our block cache.

        Parameters
        ----------
        start, end : int
            The start and end byte positions.
        start_block_number, end_block_number : int
            The start and end block numbers.
        """
        start_pos = start % self.blocksize
        end_pos = end % self.blocksize

        if start_block_number == end_block_number:
            block = self._fetch_block_cached(start_block_number)
            return block[start_pos:end_pos]

        else:
            # read from the initial
            out = []
            out.append(self._fetch_block_cached(start_block_number)[start_pos:])

            # intermediate blocks
            # Note: it'd be nice to combine these into one big request. However
            # that doesn't play nicely with our LRU cache.
            for block_number in range(start_block_number + 1, end_block_number):
                out.append(self._fetch_block_cached(block_number))

            # final block
            out.append(self._fetch_block_cached(end_block_number)[:end_pos])

            return b"".join(out)

class BytesCache(BaseCache):
    """Cache which holds data in a in-memory bytes object

    Implements read-ahead by the block size, for semi-random reads progressing
    through the file.

    Parameters
    ----------
    trim: bool
        As we read more data, whether to discard the start of the buffer when
        we are more than a blocksize ahead of it.
    """

    name = "bytes"

    def __init__(self, blocksize, fetcher, size, trim=True):
        super().__init__(blocksize, fetcher, size)
        self.cache = b""
        self.start = None
        self.end = None
        self.trim = trim

    def _fetch(self, start, end):
        # TODO: only set start/end after fetch, in case it fails?
        # is this where retry logic might go?
        if start is None:
            start = 0
        if end is None:
            end = self.size
        if start >= self.size or start >= end:
            return b""
        if (
            self.start is not None
            and start >= self.start
            and self.end is not None
            and end < self.end
        ):
            # cache hit: we have all the required data
            offset = start - self.start
            return self.cache[offset : offset + end - start]

        if self.blocksize:
            bend = min(self.size, end + self.blocksize)
        else:
            bend = end

        if bend == start or start > self.size:
            return b""

        if (self.start is None or start < self.start) and (
            self.end is None or end > self.end
        ):
            # First read, or extending both before and after
            self.cache = self.fetcher(start, bend)
            self.start = start
        elif start < self.start:
            if self.end - end > self.blocksize:
                self.cache = self.fetcher(start, bend)
                self.start = start
            else:
                new = self.fetcher(start, self.start)
                self.start = start
                self.cache = new + self.cache
        elif bend > self.end:
            if self.end > self.size:
                pass
            elif end - self.end > self.blocksize:
                self.cache = self.fetcher(start, bend)
                self.start = start
            else:
                new = self.fetcher(self.end, bend)
                self.cache = self.cache + new

        self.end = self.start + len(self.cache)
        offset = start - self.start
        out = self.cache[offset : offset + end - start]
        if self.trim:
            num = (self.end - self.start) // (self.blocksize + 1)
            if num > 1:
                self.start += self.blocksize * num
                self.cache = self.cache[self.blocksize * num :]
        return out

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

class AllBytes(BaseCache):
    """Cache entire contents of the file"""

    name = "all"

    def __init__(self, blocksize=None, fetcher=None, size=None, data=None):
        super().__init__(blocksize, fetcher, size)
        if data is None:
            data = self.fetcher(0, self.size)
        self.data = data

    def _fetch(self, start, end):
        return self.data[start:end]

class KnownPartsOfAFile(BaseCache):
    name = "parts"

    def __init__(self, blocksize, fetcher, size, data={}, **_):
        super(KnownPartsOfAFile, self).__init__(blocksize, fetcher, size)

        # simple consolidation of contiguous blocks
        for start0, stop in data.copy():
            for start, stop1 in data.copy():
                if stop == start:
                    data[(start0, stop1)] = data.pop((start0, stop)) + data.pop(
                        (start, stop1)
                    )
        self.data = data

    def _fetch(self, start, stop):
        for (loc0, loc1), data in self.data.items():
            if loc0 <= start < loc1 and loc0 <= stop <= loc1:
                off = start - loc0
                out = data[off : off + stop - start]
                # reads beyond buffer are padded with zero
                out += b"\x00" * (stop - start - len(out))
                return out

        # We only get here if there is a request outside the
        # known parts of the file. In an ideal world, this
        # should never happen.
        warnings.warn(
            f"Read is outside the known file parts: {(start, stop)}. "
            f"IO/caching performance may be poor!"
        )
        return super()._fetch(start, stop)

caches = {
    "none": BaseCache,
    "mmap": MMapCache,
    "bytes": BytesCache,
    "readahead": ReadAheadCache,
    "block": BlockCache,
    "first": FirstChunkCache,
    "all": AllBytes,
    "parts": KnownPartsOfAFile,
}
<commit_msg>fix cache test failure<commit_after>import functools
import io
import logging
import math
import os
import warnings

logger = logging.getLogger("fsspec")

class BaseCache(object):
    """Pass-though cache: doesn't keep anything, calls every time

    Acts as base class for other cachers

    Parameters
    ----------
    blocksize: int
        How far to read ahead in numbers of bytes
    fetcher: func
        Function of the form f(start, end) which gets bytes from remote as
        specified
    size: int
        How big this file is
    """

    name = "none"

    def __init__(self, blocksize, fetcher, size):
        self.blocksize = blocksize
        self.fetcher = fetcher
        self.size = size

    def _fetch(self, start, stop):
        if start is None:
            start = 0
        if stop is None:
            stop = self.size
        if start >= self.size or start >= stop:
            return b""
        return self.fetcher(start, stop)

class MMapCache(BaseCache):
    """memory-mapped sparse file cache

    Opens temporary file, which is filled blocks-wise when data is requested.
    Ensure there is enough disc space in the temporary location.

    This cache method might only work on posix
    """

    name = "mmap"

    def __init__(self, blocksize, fetcher, size, location=None, blocks=None):
        super().__init__(blocksize, fetcher, size)
        self.blocks = set() if blocks is None else blocks
        self.location = location
        self.cache = self._makefile()

    def _makefile(self):
        import mmap
        import tempfile

        if self.size == 0:
            return bytearray()

        # posix version
        if self.location is None or not os.path.exists(self.location):
            if self.location is None:
                fd = tempfile.TemporaryFile()
                self.blocks = set()
            else:
                fd = io.open(self.location, "wb+")
            fd.seek(self.size - 1)
            fd.write(b"1")
            fd.flush()
        else:
            fd = io.open(self.location, "rb+")

        return mmap.mmap(fd.fileno(), self.size)

    def _fetch(self, start, end):
        logger.debug(f"MMap cache fetching {start}-{end}")
        if start is None:
            start = 0
        if end is None:
            end = self.size
        if start >= self.size or start >= end:
            return b""
        start_block = start // self.blocksize
        end_block = end // self.blocksize
        need = [i for i in range(start_block, end_block + 1) if i not in self.blocks]
        while need:
            # TODO: not a for loop so we can consolidate blocks later to
            # make fewer fetch calls; this could be parallel
            i = need.pop(0)
            sstart = i * self.blocksize
            send = min(sstart + self.blocksize, self.size)
            logger.debug(f"MMap get block #{i} ({sstart}-{send}")
            self.cache[sstart:send] = self.fetcher(sstart, send)
            self.blocks.add(i)

        return self.cache[start:end]

    def __getstate__(self):
        state = self.__dict__.copy()
        # Remove the unpicklable entries.
        del state["cache"]
        return state

    def __setstate__(self, state):
        # Restore instance attributes
        self.__dict__.update(state)
        self.cache = self._makefile()

class ReadAheadCache(BaseCache):
    """Cache which reads only when we get beyond a block of data

    This is a much simpler version of BytesCache, and does not attempt to
    fill holes in the cache or keep fragments alive. It is best suited to
    many small reads in a sequential order (e.g., reading lines from a file).
    """

    name = "readahead"

    def __init__(self, blocksize, fetcher, size):
        super().__init__(blocksize, fetcher, size)
        self.cache = b""
        self.start = 0
        self.end = 0

    def _fetch(self, start, end):
        if start is None:
            start = 0
        if end is None or end > self.size:
            end = self.size
        if start >= self.size or start >= end:
            return b""
        l = end - start
        if start >= self.start and end <= self.end:
            # cache hit
            return self.cache[start - self.start : end - self.start]
        elif self.start <= start < self.end:
            # partial hit
            part = self.cache[start - self.start :]
            l -= len(part)
            start = self.end
        else:
            # miss
            part = b""
        end = min(self.size, end + self.blocksize)
        self.cache = self.fetcher(start, end)  # new block replaces old
        self.start = start
        self.end = self.start + len(self.cache)
        return part + self.cache[:l]

class FirstChunkCache(BaseCache):
    """Caches the first block of a file only

    This may be useful for file types where the metadata is stored in the header,
    but is randomly accessed.
    """

    name = "first"

    def __init__(self, blocksize, fetcher, size):
        super().__init__(blocksize, fetcher, size)
        self.cache = None

    def _fetch(self, start, end):
        start = start or 0
        end = end or self.size
        if start < self.blocksize:
            if self.cache is None:
                if end > self.blocksize:
                    data = self.fetcher(0, end)
                    self.cache = data[: self.blocksize]
                    return data[start:]
                self.cache = self.fetcher(0, self.blocksize)
            part = self.cache[start:end]
            if end > self.blocksize:
                part += self.fetcher(self.blocksize, end)
            return part
        else:
            return self.fetcher(start, end)

class BlockCache(BaseCache):
    """
    Cache holding memory as a set of blocks.

    Requests are only ever made `blocksize` at a time, and are
    stored in an LRU cache. The least recently accessed block is
    discarded when more than `maxblocks` are stored.

    Parameters
    ----------
    blocksize : int
        The number of bytes to store in each block.
        Requests are only ever made for `blocksize`, so this
        should balance the overhead of making a request against
        the granularity of the blocks.
    fetcher : Callable
    size : int
        The total size of the file being cached.
    maxblocks : int
        The maximum number of blocks to cache for. The maximum memory
        use for this cache is then ``blocksize * maxblocks``.
    """

    name = "blockcache"

    def __init__(self, blocksize, fetcher, size, maxblocks=32):
        super().__init__(blocksize, fetcher, size)
        self.nblocks = math.ceil(size / blocksize)
        self.maxblocks = maxblocks
        self._fetch_block_cached = functools.lru_cache(maxblocks)(self._fetch_block)

    def __repr__(self):
        return "<BlockCache blocksize={}, size={}, nblocks={}>".format(
            self.blocksize, self.size, self.nblocks
        )

    def cache_info(self):
        """
        The statistics on the block cache.

        Returns
        ----------
        NamedTuple
            Returned directly from the LRU Cache used internally.
        """
        return self._fetch_block_cached.cache_info()

    def __getstate__(self):
        state = self.__dict__
        del state["_fetch_block_cached"]
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self._fetch_block_cached = functools.lru_cache(state["maxblocks"])(
            self._fetch_block
        )

    def _fetch(self, start, end):
        if start is None:
            start = 0
        if end is None:
            end = self.size
        if start >= self.size or start >= end:
            return b""

        # byte position -> block numbers
        start_block_number = start // self.blocksize
        end_block_number = end // self.blocksize

        # these are cached, so safe to do multiple calls for the same start and end.
        for block_number in range(start_block_number, end_block_number + 1):
            self._fetch_block_cached(block_number)

        return self._read_cache(
            start,
            end,
            start_block_number=start_block_number,
            end_block_number=end_block_number,
        )

    def _fetch_block(self, block_number):
        """
        Fetch the block of data for `block_number`.
        """
        if block_number > self.nblocks:
            raise ValueError(
                "'block_number={}' is greater than the number of blocks ({})".format(
                    block_number, self.nblocks
                )
            )

        start = block_number * self.blocksize
        end = start + self.blocksize
        logger.info("BlockCache fetching block %d", block_number)
        block_contents = super()._fetch(start, end)
        return block_contents

    def _read_cache(self, start, end, start_block_number, end_block_number):
        """
        Read from our block cache.

        Parameters
        ----------
        start, end : int
            The start and end byte positions.
        start_block_number, end_block_number : int
            The start and end block numbers.
        """
        start_pos = start % self.blocksize
        end_pos = end % self.blocksize

        if start_block_number == end_block_number:
            block = self._fetch_block_cached(start_block_number)
            return block[start_pos:end_pos]

        else:
            # read from the initial
            out = []
            out.append(self._fetch_block_cached(start_block_number)[start_pos:])

            # intermediate blocks
            # Note: it'd be nice to combine these into one big request. However
            # that doesn't play nicely with our LRU cache.
            for block_number in range(start_block_number + 1, end_block_number):
                out.append(self._fetch_block_cached(block_number))

            # final block
            out.append(self._fetch_block_cached(end_block_number)[:end_pos])

            return b"".join(out)

class BytesCache(BaseCache):
    """Cache which holds data in a in-memory bytes object

    Implements read-ahead by the block size, for semi-random reads progressing
    through the file.

    Parameters
    ----------
    trim: bool
        As we read more data, whether to discard the start of the buffer when
        we are more than a blocksize ahead of it.
    """

    name = "bytes"

    def __init__(self, blocksize, fetcher, size, trim=True):
        super().__init__(blocksize, fetcher, size)
        self.cache = b""
        self.start = None
        self.end = None
        self.trim = trim

    def _fetch(self, start, end):
        # TODO: only set start/end after fetch, in case it fails?
        # is this where retry logic might go?
        if start is None:
            start = 0
        if end is None:
            end = self.size
        if start >= self.size or start >= end:
            return b""
        if (
            self.start is not None
            and start >= self.start
            and self.end is not None
            and end < self.end
        ):
            # cache hit: we have all the required data
            offset = start - self.start
            return self.cache[offset : offset + end - start]

        if self.blocksize:
            bend = min(self.size, end + self.blocksize)
        else:
            bend = end

        if bend == start or start > self.size:
            return b""

        if (self.start is None or start < self.start) and (
            self.end is None or end > self.end
        ):
            # First read, or extending both before and after
            self.cache = self.fetcher(start, bend)
            self.start = start
        elif start < self.start:
            if self.end - end > self.blocksize:
                self.cache = self.fetcher(start, bend)
                self.start = start
            else:
                new = self.fetcher(start, self.start)
                self.start = start
                self.cache = new + self.cache
        elif bend > self.end:
            if self.end > self.size:
                pass
            elif end - self.end > self.blocksize:
                self.cache = self.fetcher(start, bend)
                self.start = start
            else:
                new = self.fetcher(self.end, bend)
                self.cache = self.cache + new

        self.end = self.start + len(self.cache)
        offset = start - self.start
        out = self.cache[offset : offset + end - start]
        if self.trim:
            num = (self.end - self.start) // (self.blocksize + 1)
            if num > 1:
                self.start += self.blocksize * num
                self.cache = self.cache[self.blocksize * num :]
        return out

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

class AllBytes(BaseCache):
    """Cache entire contents of the file"""

    name = "all"

    def __init__(self, blocksize=None, fetcher=None, size=None, data=None):
        super().__init__(blocksize, fetcher, size)
        if data is None:
            data = self.fetcher(0, self.size)
        self.data = data

    def _fetch(self, start, end):
        return self.data[start:end]

class KnownPartsOfAFile(BaseCache):
    name = "parts"

    def __init__(self, blocksize, fetcher, size, data={}, **_):
        super(KnownPartsOfAFile, self).__init__(blocksize, fetcher, size)

        # simple consolidation of contiguous blocks
        for start0, stop in data.copy():
            for start, stop1 in data.copy():
                if stop == start:
                    data[(start0, stop1)] = data.pop((start0, stop)) + data.pop(
                        (start, stop1)
                    )
        self.data = data

    def _fetch(self, start, stop):
        for (loc0, loc1), data in self.data.items():
            if loc0 <= start < loc1 and loc0 <= stop <= loc1:
                off = start - loc0
                out = data[off : off + stop - start]
                # reads beyond buffer are padded with zero
                out += b"\x00" * (stop - start - len(out))
                return out

        # We only get here if there is a request outside the
        # known parts of the file. In an ideal world, this
        # should never happen
        if self.fetcher is None:
            # We cannot fetch the data, so raise an error
            raise ValueError(f"Read is outside the known file parts: {(start, stop)}. ")
        # We can fetch the data, but should warn the user
        # that this may be slow
        warnings.warn(
            f"Read is outside the known file parts: {(start, stop)}. "
            f"IO/caching performance may be poor!"
        )
        return super()._fetch(start, stop)

caches = {
    "none": BaseCache,
    "mmap": MMapCache,
    "bytes": BytesCache,
    "readahead": ReadAheadCache,
    "block": BlockCache,
    "first": FirstChunkCache,
    "all": AllBytes,
    "parts": KnownPartsOfAFile,
}
<|endoftext|>