<commit_before>from contextlib import contextmanager, closing
import csv
from datetime import datetime, timedelta
import gzip
import os
import shutil
import tempfile

import boto
import requests
from pytz import UTC
from sqlalchemy.sql import (
    and_,
    func,
    select,
)

from ichnaea.async.task import DatabaseTask
from ichnaea.data.validation import normalized_cell_dict
from ichnaea.models import (
    cell_table,
    ocid_cell_table,
    CellAreaKey,
    RADIO_TYPE,
    RADIO_TYPE_INVERSE,
)
from ichnaea.data.tasks import update_lac
from ichnaea.worker import celery
from ichnaea import util

CELL_FIELDS = [
    'radio', 'mcc', 'mnc', 'lac', 'cid', 'psc',
    'lon', 'lat', 'range', 'samples', 'changeable',
    'created', 'updated', 'averageSignal']
CELL_FIELD_INDICES = dict(
    [(e, i) for (i, e) in enumerate(CELL_FIELDS)]
)
# Map our internal names to the public export names
CELL_HEADER_DICT = dict([(field, field) for field in CELL_FIELDS])
CELL_HEADER_DICT['mnc'] = 'net'
CELL_HEADER_DICT['lac'] = 'area'
CELL_HEADER_DICT['cid'] = 'cell'
CELL_HEADER_DICT['psc'] = 'unit'

# The list of cell columns, we actually need for the export
CELL_COLUMN_NAMES = [
    'created', 'modified', 'lat', 'lon',
    'radio', 'mcc', 'mnc', 'lac', 'cid', 'psc',
    'range', 'total_measures']

CELL_COLUMN_NAME_INDICES = dict(
    [(e, i) for (i, e) in enumerate(CELL_COLUMN_NAMES)]
)
CELL_COLUMNS = []
for name in CELL_COLUMN_NAMES:
    if name in ('created', 'modified'):
        CELL_COLUMNS.append(func.unix_timestamp(getattr(cell_table.c, name)))
    else:
        CELL_COLUMNS.append(getattr(cell_table.c, name))

CELL_EXPORT_RADIO_NAMES = dict(
    [(k, v.upper()) for k, v in RADIO_TYPE_INVERSE.items()])

@contextmanager
def selfdestruct_tempdir():
    base_path = tempfile.mkdtemp()
    try:
        yield base_path
    finally:
        shutil.rmtree(base_path)

# Python 2.6 Gzipfile doesn't have __exit__
class GzipFile(gzip.GzipFile):

    # Add a default value, as this is used in the __repr__ and only
    # set in the __init__ on successfully opening the file. Sentry/raven
    # would bark on this while trying to capture the stack frame locals.
    fileobj = None

    def __enter__(self):
        if self.fileobj is None:  # pragma: no cover
            raise ValueError('I/O operation on closed GzipFile object')
        return self

    def __exit__(self, *args):
        self.close()

def make_cell_export_dict(row):
    d = {
        'changeable': 1,
        'averageSignal': '',
    }
    ix = CELL_COLUMN_NAME_INDICES

    for field in CELL_FIELDS:
        pos = ix.get(field, None)
        if pos is not None:
            d[field] = row[pos]

    # Fix up specific entry formatting
    radio = row[ix['radio']]
    if radio is None:  # pragma: no cover
        radio = -1

    psc = row[ix['psc']]
    if psc is None or psc == -1:
        psc = ''

    d['radio'] = CELL_EXPORT_RADIO_NAMES[radio]
    d['created'] = row[ix['created']]
    d['updated'] = row[ix['modified']]
    d['samples'] = row[ix['total_measures']]
    d['psc'] = psc
    return d

def make_cell_import_dict(row):

    def val(key, default):
        if key in row and row[key] != '' and row[key] is not None:
            return row[key]
        else:
            return default

    d = dict()

    d['created'] = datetime.fromtimestamp(
        int(val('created', 0))).replace(tzinfo=UTC)

    d['modified'] = datetime.fromtimestamp(
        int(val('updated', 0))).replace(tzinfo=UTC)

    d['lat'] = float(val('lat', -255))
    d['lon'] = float(val('lon', -255))

    d['radio'] = RADIO_TYPE.get(row['radio'].lower(), -1)

    for k in ['mcc', 'mnc', 'lac', 'cid', 'psc']:
        d[k] = int(val(k, -1))

    d['range'] = int(float(val('range', 0)))

    d['total_measures'] = int(val('samples', -1))
    d['changeable'] = bool(val('changeable', True))
    return normalized_cell_dict(d)

def write_stations_to_csv(session, table, columns,
                          cond, path, make_dict, fields):
    with GzipFile(path, 'wb') as f:
        w = csv.DictWriter(f, fields, extrasaction='ignore')
        limit = 10000
        offset = 0
        # Write header row
        w.writerow(CELL_HEADER_DICT)
        while True:
            q = select(columns=columns).where(cond).limit(
                limit).offset(offset).order_by(table.c.id)
            rows = session.execute(q).fetchall()
            if rows:
                w.writerows([make_dict(r) for r in rows])
                offset += limit
            else:
                break

def write_stations_to_s3(path, bucketname):
    conn = boto.connect_s3()
    bucket = conn.get_bucket(bucketname)
    k = boto.s3.key.Key(bucket)
    k.key = 'export/' + os.path.split(path)[-1]
    k.set_contents_from_filename(path, reduced_redundancy=True)

@celery.task(base=DatabaseTask, bind=True)
def export_modified_cells(self, hourly=True, bucket=None):
    if bucket is None:  # pragma: no cover
        bucket = self.app.s3_settings['assets_bucket']
    now = util.utcnow()

    if hourly:
        end_time = now.replace(minute=0, second=0)
        file_time = end_time
        file_type = 'diff'
        start_time = end_time - timedelta(hours=1)
        cond = and_(cell_table.c.modified >= start_time,
                    cell_table.c.modified < end_time,
                    cell_table.c.lat.isnot(None))
    else:
        file_time = now.replace(hour=0, minute=0, second=0)
        file_type = 'full'
        cond = cell_table.c.lat.isnot(None)

    filename = 'MLS-%s-cell-export-' % file_type
    filename = filename + file_time.strftime('%Y-%m-%dT%H0000.csv.gz')
    try:
        with selfdestruct_tempdir() as d:
            path = os.path.join(d, filename)
            with self.db_session() as session:
                write_stations_to_csv(session, cell_table, CELL_COLUMNS, cond,
                                      path, make_cell_export_dict, CELL_FIELDS)
            write_stations_to_s3(path, bucket)
    except Exception as exc:  # pragma: no cover
        self.heka_client.raven('error')
        raise self.retry(exc=exc)

def import_stations(session, filename, fields):
    with GzipFile(filename, 'rb') as zip_file:
        csv_reader = csv.DictReader(zip_file, fields)
        batch = 10000
        rows = []
        lacs = set()
        ins = ocid_cell_table.insert(
            on_duplicate=((
                'modified = values(modified), '
                'total_measures = values(total_measures), '
                'lat = values(lat), '
                'lon = values(lon), '
                'psc = values(psc), '
                '`range` = values(`range`)')))

        for row in csv_reader:
            # skip any header row
            if csv_reader.line_num == 1 and 'radio' in row.values():
                continue

            data = make_cell_import_dict(row)
            if data is not None:
                rows.append(data)
                lacs.add(CellAreaKey(
                    radio=data['radio'],
                    mcc=data['mcc'],
                    mnc=data['mnc'],
                    lac=data['lac']))

            if len(rows) == batch:  # pragma: no cover
                session.execute(ins, rows)
                session.commit()
                rows = []

        if rows:
            session.execute(ins, rows)
            session.commit()

        for lac in lacs:
            update_lac.delay(
                lac.radio,
                lac.mcc,
                lac.mnc,
                lac.lac,
                cell_model_key='ocid_cell',
                cell_area_model_key='ocid_cell_area')

@celery.task(base=DatabaseTask, bind=True)
def import_ocid_cells(self, filename=None, session=None):
    try:
        with self.db_session() as dbsession:
            if session is None:  # pragma: no cover
                session = dbsession
            import_stations(session,
                            filename,
                            CELL_FIELDS)
    except Exception as exc:  # pragma: no cover
        self.heka_client.raven('error')
        raise self.retry(exc=exc)

@celery.task(base=DatabaseTask, bind=True)
def import_latest_ocid_cells(self, diff=True, filename=None, session=None):
    url = self.app.ocid_settings['ocid_url']
    apikey = self.app.ocid_settings['ocid_apikey']
    if filename is None:
        if diff:
            prev_hour = util.utcnow() - timedelta(hours=1)
            filename = prev_hour.strftime('cell_towers_diff-%Y%m%d%H.csv.gz')
        else:  # pragma: no cover
            filename = 'cell_towers.csv.gz'
    try:
        with closing(requests.get(url,
                                  params={'apiKey': apikey,
                                          'filename': filename},
                                  stream=True)) as r:
            with selfdestruct_tempdir() as d:
                path = os.path.join(d, filename)
                with open(path, 'wb') as f:
                    for chunk in r.iter_content(chunk_size=2 ** 20):
                        f.write(chunk)
                        f.flush()

                with self.db_session() as dbsession:
                    if session is None:  # pragma: no cover
                        session = dbsession
                    import_stations(session,
                                    path,
                                    CELL_FIELDS)

    except Exception as exc:  # pragma: no cover
        self.heka_client.raven('error')
        raise self.retry(exc=exc)
<commit_msg>Cheat coverage tests.<commit_after>from contextlib import contextmanager, closing
import csv
from datetime import datetime, timedelta
import gzip
import os
import shutil
import tempfile

import boto
import requests
from pytz import UTC
from sqlalchemy.sql import (
    and_,
    func,
    select,
)

from ichnaea.async.task import DatabaseTask
from ichnaea.data.validation import normalized_cell_dict
from ichnaea.models import (
    cell_table,
    ocid_cell_table,
    CellAreaKey,
    RADIO_TYPE,
    RADIO_TYPE_INVERSE,
)
from ichnaea.data.tasks import update_lac
from ichnaea.worker import celery
from ichnaea import util

CELL_FIELDS = [
    'radio', 'mcc', 'mnc', 'lac', 'cid', 'psc',
    'lon', 'lat', 'range', 'samples', 'changeable',
    'created', 'updated', 'averageSignal']
CELL_FIELD_INDICES = dict(
    [(e, i) for (i, e) in enumerate(CELL_FIELDS)]
)
# Map our internal names to the public export names
CELL_HEADER_DICT = dict([(field, field) for field in CELL_FIELDS])
CELL_HEADER_DICT['mnc'] = 'net'
CELL_HEADER_DICT['lac'] = 'area'
CELL_HEADER_DICT['cid'] = 'cell'
CELL_HEADER_DICT['psc'] = 'unit'

# The list of cell columns, we actually need for the export
CELL_COLUMN_NAMES = [
    'created', 'modified', 'lat', 'lon',
    'radio', 'mcc', 'mnc', 'lac', 'cid', 'psc',
    'range', 'total_measures']

CELL_COLUMN_NAME_INDICES = dict(
    [(e, i) for (i, e) in enumerate(CELL_COLUMN_NAMES)]
)
CELL_COLUMNS = []
for name in CELL_COLUMN_NAMES:
    if name in ('created', 'modified'):
        CELL_COLUMNS.append(func.unix_timestamp(getattr(cell_table.c, name)))
    else:
        CELL_COLUMNS.append(getattr(cell_table.c, name))

CELL_EXPORT_RADIO_NAMES = dict(
    [(k, v.upper()) for k, v in RADIO_TYPE_INVERSE.items()])

@contextmanager
def selfdestruct_tempdir():
    base_path = tempfile.mkdtemp()
    try:
        yield base_path
    finally:
        shutil.rmtree(base_path)

# Python 2.6 Gzipfile doesn't have __exit__
class GzipFile(gzip.GzipFile):

    # Add a default value, as this is used in the __repr__ and only
    # set in the __init__ on successfully opening the file. Sentry/raven
    # would bark on this while trying to capture the stack frame locals.
    fileobj = None

    def __enter__(self):
        if self.fileobj is None:  # pragma: no cover
            raise ValueError('I/O operation on closed GzipFile object')
        return self

    def __exit__(self, *args):
        self.close()

def make_cell_export_dict(row):
    d = {
        'changeable': 1,
        'averageSignal': '',
    }
    ix = CELL_COLUMN_NAME_INDICES

    for field in CELL_FIELDS:
        pos = ix.get(field, None)
        if pos is not None:
            d[field] = row[pos]

    # Fix up specific entry formatting
    radio = row[ix['radio']]
    if radio is None:  # pragma: no cover
        radio = -1

    psc = row[ix['psc']]
    if psc is None or psc == -1:
        psc = ''

    d['radio'] = CELL_EXPORT_RADIO_NAMES[radio]
    d['created'] = row[ix['created']]
    d['updated'] = row[ix['modified']]
    d['samples'] = row[ix['total_measures']]
    d['psc'] = psc
    return d

def make_cell_import_dict(row):

    def val(key, default):
        if key in row and row[key] != '' and row[key] is not None:
            return row[key]
        else:
            return default

    d = dict()

    d['created'] = datetime.fromtimestamp(
        int(val('created', 0))).replace(tzinfo=UTC)

    d['modified'] = datetime.fromtimestamp(
        int(val('updated', 0))).replace(tzinfo=UTC)

    d['lat'] = float(val('lat', -255))
    d['lon'] = float(val('lon', -255))

    d['radio'] = RADIO_TYPE.get(row['radio'].lower(), -1)

    for k in ['mcc', 'mnc', 'lac', 'cid', 'psc']:
        d[k] = int(val(k, -1))

    d['range'] = int(float(val('range', 0)))

    d['total_measures'] = int(val('samples', -1))
    d['changeable'] = bool(val('changeable', True))
    return normalized_cell_dict(d)

def write_stations_to_csv(session, table, columns,
                          cond, path, make_dict, fields):
    with GzipFile(path, 'wb') as f:
        w = csv.DictWriter(f, fields, extrasaction='ignore')
        limit = 10000
        offset = 0
        # Write header row
        w.writerow(CELL_HEADER_DICT)
        while True:
            q = select(columns=columns).where(cond).limit(
                limit).offset(offset).order_by(table.c.id)
            rows = session.execute(q).fetchall()
            if rows:
                w.writerows([make_dict(r) for r in rows])
                offset += limit
            else:
                break

def write_stations_to_s3(path, bucketname):
    conn = boto.connect_s3()
    bucket = conn.get_bucket(bucketname)
    k = boto.s3.key.Key(bucket)
    k.key = 'export/' + os.path.split(path)[-1]
    k.set_contents_from_filename(path, reduced_redundancy=True)

@celery.task(base=DatabaseTask, bind=True)
def export_modified_cells(self, hourly=True, bucket=None):
    if bucket is None:  # pragma: no cover
        bucket = self.app.s3_settings['assets_bucket']
    now = util.utcnow()

    if hourly:
        end_time = now.replace(minute=0, second=0)
        file_time = end_time
        file_type = 'diff'
        start_time = end_time - timedelta(hours=1)
        cond = and_(cell_table.c.modified >= start_time,
                    cell_table.c.modified < end_time,
                    cell_table.c.lat.isnot(None))
    else:
        file_time = now.replace(hour=0, minute=0, second=0)
        file_type = 'full'
        cond = cell_table.c.lat.isnot(None)

    filename = 'MLS-%s-cell-export-' % file_type
    filename = filename + file_time.strftime('%Y-%m-%dT%H0000.csv.gz')
    try:
        with selfdestruct_tempdir() as d:
            path = os.path.join(d, filename)
            with self.db_session() as session:
                write_stations_to_csv(session, cell_table, CELL_COLUMNS, cond,
                                      path, make_cell_export_dict, CELL_FIELDS)
            write_stations_to_s3(path, bucket)
    except Exception as exc:  # pragma: no cover
        self.heka_client.raven('error')
        raise self.retry(exc=exc)

def import_stations(session, filename, fields):
    with GzipFile(filename, 'rb') as zip_file:
        csv_reader = csv.DictReader(zip_file, fields)
        batch = 10000
        rows = []
        lacs = set()
        ins = ocid_cell_table.insert(
            on_duplicate=((
                'modified = values(modified), '
                'total_measures = values(total_measures), '
                'lat = values(lat), '
                'lon = values(lon), '
                'psc = values(psc), '
                '`range` = values(`range`)')))

        for row in csv_reader:
            # skip any header row
            if csv_reader.line_num == 1 and \
               'radio' in row.values():  # pragma: no cover
                continue

            data = make_cell_import_dict(row)
            if data is not None:
                rows.append(data)
                lacs.add(CellAreaKey(
                    radio=data['radio'],
                    mcc=data['mcc'],
                    mnc=data['mnc'],
                    lac=data['lac']))

            if len(rows) == batch:  # pragma: no cover
                session.execute(ins, rows)
                session.commit()
                rows = []

        if rows:
            session.execute(ins, rows)
            session.commit()

        for lac in lacs:
            update_lac.delay(
                lac.radio,
                lac.mcc,
                lac.mnc,
                lac.lac,
                cell_model_key='ocid_cell',
                cell_area_model_key='ocid_cell_area')

@celery.task(base=DatabaseTask, bind=True)
def import_ocid_cells(self, filename=None, session=None):
    try:
        with self.db_session() as dbsession:
            if session is None:  # pragma: no cover
                session = dbsession
            import_stations(session,
                            filename,
                            CELL_FIELDS)
    except Exception as exc:  # pragma: no cover
        self.heka_client.raven('error')
        raise self.retry(exc=exc)

@celery.task(base=DatabaseTask, bind=True)
def import_latest_ocid_cells(self, diff=True, filename=None, session=None):
    url = self.app.ocid_settings['ocid_url']
    apikey = self.app.ocid_settings['ocid_apikey']
    if filename is None:
        if diff:
            prev_hour = util.utcnow() - timedelta(hours=1)
            filename = prev_hour.strftime('cell_towers_diff-%Y%m%d%H.csv.gz')
        else:  # pragma: no cover
            filename = 'cell_towers.csv.gz'
    try:
        with closing(requests.get(url,
                                  params={'apiKey': apikey,
                                          'filename': filename},
                                  stream=True)) as r:
            with selfdestruct_tempdir() as d:
                path = os.path.join(d, filename)
                with open(path, 'wb') as f:
                    for chunk in r.iter_content(chunk_size=2 ** 20):
                        f.write(chunk)
                        f.flush()

                with self.db_session() as dbsession:
                    if session is None:  # pragma: no cover
                        session = dbsession
                    import_stations(session,
                                    path,
                                    CELL_FIELDS)

    except Exception as exc:  # pragma: no cover
        self.heka_client.raven('error')
        raise self.retry(exc=exc)
<|endoftext|>


<commit_before>from unittest import TestCase

from wallace.dataset import Dataset, DatasetVariable

class DatasetTest(TestCase):
    def setUp(self):
        data_matrix = [
                [0,"detective", "book"],
                [1, "pablo", "escobar"],
                [2, "african", "queen"],
                [3, "applause", "movie"]
                ]
        headers = ["number", "string", "another_string"]
        self.dataset = Dataset(data_matrix, headers)

    def test_data_types_and_categories_when_initializing_dataset(self):
        self.assertEqual(3, len(self.dataset.data_types))
        self.assertEqual("integer", self.dataset.data_types[0].data_type)
        self.assertEqual("string", self.dataset.data_types[1].data_type)
        self.assertEqual("string", self.dataset.data_types[2].data_type)

        self.assertEqual(4, len(self.dataset.data_types[1].categories))
        self.assertTrue(self.dataset.data_types[1].in_categories("detective"))
        self.assertTrue(self.dataset.data_types[1].in_categories("pablo"))
        self.assertTrue(self.dataset.data_types[1].in_categories("african"))
        self.assertTrue(self.dataset.data_types[1].in_categories("applause"))

        self.assertEqual(4, len(self.dataset.data_types[2].categories))
        self.assertTrue(self.dataset.data_types[2].in_categories("book"))
        self.assertTrue(self.dataset.data_types[2].in_categories("escobar"))
        self.assertTrue(self.dataset.data_types[2].in_categories("queen"))
        self.assertTrue(self.dataset.data_types[2].in_categories("movie"))

    def test_getting_independent_variables(self):
        dependent_variable_index = DatasetVariable(0)
        dependent_variable_header = DatasetVariable("number")

        independent_variables = self.dataset.get_independent_variables(dependent_variable_index)
        self.assertEqual(2, len(independent_variables))
        for variable in independent_variables:
            self.assertIn(variable.variable, ["string", "another_string"])

        independent_variables = self.dataset.get_independent_variables(dependent_variable_header)
        self.assertEqual(2, len(independent_variables))
        for variable in independent_variables:
            self.assertIn(variable.variable, ["string", "another_string"])

    def test_crossfold_partitions(self):
        past_test_datasets = []
        for training_dataset, test_dataset in self.dataset.crossfold_partitions(4):
            self.assertEqual(3, training_dataset.num_rows)
            self.assertEqual(1, test_dataset.num_rows)
            self.assertNotIn(test_dataset.get(0, 0), past_test_datasets)
            self.assertListEqual(["number", "string", "another_string"], training_dataset.headers)
            self.assertListEqual(["number", "string", "another_string"], test_dataset.headers)

            past_test_datasets.append(test_dataset.get(0, 0))

    def test_get_filtered_matrix(self):
        variables = [DatasetVariable(0), DatasetVariable("string")]

        filtered_matrix = self.dataset.get_filtered_matrix(variables)
        self.assertEqual(4, len(filtered_matrix))
        for i in xrange(len(filtered_matrix)):
            self.assertEqual(2, len(filtered_matrix[i]))

        self.assertListEqual([0, "detective"], filtered_matrix[0])
        self.assertListEqual([1, "pablo"], filtered_matrix[1])
        self.assertListEqual([2, "african"], filtered_matrix[2])
        self.assertListEqual([3, "applause"], filtered_matrix[3])

    def test_getting_filtered_matrix_with_no_variables(self):
        variables = []

        with self.assertRaises(ValueError):
            filtered_matrix = self.dataset.get_filtered_matrix(variables)
<commit_msg>Adding extra crossfold partitioning tests.<commit_after>from unittest import TestCase

from wallace.dataset import Dataset, DatasetVariable

class DatasetTest(TestCase):
    def setUp(self):
        data_matrix = [
                [0,"detective", "book"],
                [1, "pablo", "escobar"],
                [2, "african", "queen"],
                [3, "applause", "movie"]
                ]
        headers = ["number", "string", "another_string"]
        self.dataset = Dataset(data_matrix, headers)

    def test_data_types_and_categories_when_initializing_dataset(self):
        self.assertEqual(3, len(self.dataset.data_types))
        self.assertEqual("integer", self.dataset.data_types[0].data_type)
        self.assertEqual("string", self.dataset.data_types[1].data_type)
        self.assertEqual("string", self.dataset.data_types[2].data_type)

        self.assertEqual(4, len(self.dataset.data_types[1].categories))
        self.assertTrue(self.dataset.data_types[1].in_categories("detective"))
        self.assertTrue(self.dataset.data_types[1].in_categories("pablo"))
        self.assertTrue(self.dataset.data_types[1].in_categories("african"))
        self.assertTrue(self.dataset.data_types[1].in_categories("applause"))

        self.assertEqual(4, len(self.dataset.data_types[2].categories))
        self.assertTrue(self.dataset.data_types[2].in_categories("book"))
        self.assertTrue(self.dataset.data_types[2].in_categories("escobar"))
        self.assertTrue(self.dataset.data_types[2].in_categories("queen"))
        self.assertTrue(self.dataset.data_types[2].in_categories("movie"))

    def test_getting_independent_variables(self):
        dependent_variable_index = DatasetVariable(0)
        dependent_variable_header = DatasetVariable("number")

        independent_variables = self.dataset.get_independent_variables(dependent_variable_index)
        self.assertEqual(2, len(independent_variables))
        for variable in independent_variables:
            self.assertIn(variable.variable, ["string", "another_string"])

        independent_variables = self.dataset.get_independent_variables(dependent_variable_header)
        self.assertEqual(2, len(independent_variables))
        for variable in independent_variables:
            self.assertIn(variable.variable, ["string", "another_string"])

    def test_crossfold_partitions(self):
        past_test_datasets = []
        for training_dataset, test_dataset in self.dataset.crossfold_partitions(4):
            self.assertEqual(3, training_dataset.num_rows)
            self.assertEqual(1, test_dataset.num_rows)
            self.assertNotIn(test_dataset.get(0, 0), past_test_datasets)
            self.assertListEqual(["number", "string", "another_string"], training_dataset.headers)
            self.assertListEqual(["number", "string", "another_string"], test_dataset.headers)

            past_test_datasets.append(test_dataset.get(0, 0))

        for training_dataset, test_dataset in self.dataset.crossfold_partitions(2):
            self.assertEqual(2, training_dataset.num_rows)
            self.assertEqual(2, test_dataset.num_rows)
            self.assertListEqual(["number", "string", "another_string"], training_dataset.headers)
            self.assertListEqual(["number", "string", "another_string"], test_dataset.headers)

    def test_get_filtered_matrix(self):
        variables = [DatasetVariable(0), DatasetVariable("string")]

        filtered_matrix = self.dataset.get_filtered_matrix(variables)
        self.assertEqual(4, len(filtered_matrix))
        for i in xrange(len(filtered_matrix)):
            self.assertEqual(2, len(filtered_matrix[i]))

        self.assertListEqual([0, "detective"], filtered_matrix[0])
        self.assertListEqual([1, "pablo"], filtered_matrix[1])
        self.assertListEqual([2, "african"], filtered_matrix[2])
        self.assertListEqual([3, "applause"], filtered_matrix[3])

    def test_getting_filtered_matrix_with_no_variables(self):
        variables = []

        with self.assertRaises(ValueError):
            filtered_matrix = self.dataset.get_filtered_matrix(variables)
<|endoftext|>


<commit_before>"""
Custom storage for django with Mosso Cloud Files backend.
Created by Rich Leland <rich@richleland.com>.
"""
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from django.core.files import File
from django.core.files.storage import Storage
from django.utils.text import get_valid_filename

try:
    import cloudfiles
    from cloudfiles.errors import NoSuchObject
except ImportError:
    raise ImproperlyConfigured("Could not load cloudfiles dependency. See "
                               "http://www.mosso.com/cloudfiles.jsp.")

# TODO: implement TTL into cloudfiles methods
CLOUDFILES_TTL = getattr(settings, 'CLOUDFILES_TTL', 600)

def cloudfiles_upload_to(self, filename):
    """
    Simple, custom upload_to because Cloud Files doesn't support
    nested containers (directories).

    Actually found this out from @minter:
    @richleland The Cloud Files APIs do support pseudo-subdirectories, by
    creating zero-byte files with type application/directory.

    May implement in a future version.
    """
    return get_valid_filename(filename)

class CloudFilesStorage(Storage):
    """
    Custom storage for Mosso Cloud Files.
    """
    default_quick_listdir = True

    def __init__(self, username=None, api_key=None, container=None,
                 connection_kwargs=None):
        """
        Initialize the settings for the connection and container.
        """
        self.username = username or settings.CLOUDFILES_USERNAME
        self.api_key = api_key or settings.CLOUDFILES_API_KEY
        self.container_name = container or settings.CLOUDFILES_CONTAINER
        self.connection_kwargs = connection_kwargs or {}

    def __getstate__(self):
        """
        Return a picklable representation of the storage.
        """
        return dict(username=self.username,
                    api_key=self.api_key,
                    container_name=self.container_name,
                    connection_kwargs=self.connection_kwargs)

    def _get_connection(self):
        if not hasattr(self, '_connection'):
            self._connection = cloudfiles.get_connection(self.username,
                                    self.api_key, **self.connection_kwargs)
        return self._connection

    def _set_connection(self, value):
        self._connection = value

    connection = property(_get_connection, _set_connection)

    def _get_container(self):
        if not hasattr(self, '_container'):
            self.container = self.connection.get_container(
                                                        self.container_name)
        return self._container

    def _set_container(self, container):
        """
        Set the container, making it publicly available (on Limelight CDN) if
        it is not already.
        """
        if not container.is_public():
            container.make_public()
        if hasattr(self, '_container_public_uri'):
            delattr(self, '_container_public_uri')
        self._container = container

    container = property(_get_container, _set_container)

    def _get_container_url(self):
        if not hasattr(self, '_container_public_uri'):
            self._container_public_uri = self.container.public_uri()
        return self._container_public_uri

    container_url = property(_get_container_url)

    def _get_cloud_obj(self, name):
        """
        Helper function to get retrieve the requested Cloud Files Object.
        """
        return self.container.get_object(name)

    def _open(self, name, mode='rb'):
        """
        Return the CloudFilesStorageFile.
        """
        return CloudFilesStorageFile(storage=self, name=name)

    def _save(self, name, content):
        """
        Use the Cloud Files service to write ``content`` to a remote file
        (called ``name``).
        """

        cloud_obj = self.container.create_object(name)
        cloud_obj.size = content.size

        content.open()
        # If the content type is available, pass it in directly rather than
        # getting the cloud object to try to guess.
        if hasattr(content.file, 'content_type'):
            cloud_obj.content_type = content.file.content_type
        cloud_obj.send(content)
        content.close()
        return name

    def delete(self, name):
        """
        Deletes the specified file from the storage system.
        """
        # If the file exists, delete it.
        if self.exists(name):
            self.container.delete_object(name)

    def exists(self, name):
        """
        Returns True if a file referenced by the given name already exists in
        the storage system, or False if the name is available for a new file.
        """
        try:
            self._get_cloud_obj(name)
            return True
        except NoSuchObject:
            return False

    def listdir(self, path):
        """
        Lists the contents of the specified path, returning a 2-tuple; the
        first being an empty list of directories (not available for quick-
        listing), the second being a list of filenames.

        If the list of directories is required, use the full_listdir method.
        """
        files = []
        if path and not path.endswith('/'):
            path = '%s/' % path
        path_len = len(path)
        for name in self.container.list_objects(path=path):
            files.append(name[path_len:])
        return ([], files)

    def full_listdir(self, path):
        """
        Lists the contents of the specified path, returning a 2-tuple of lists;
        the first item being directories, the second item being files.

        On large containers, this may be a slow operation for root containers
        because every single object must be returned (cloudfiles does not
        provide an explicit way of listing directories).
        """
        dirs = set()
        files = []
        if path and not path.endswith('/'):
            path = '%s/' % path
        path_len = len(path)
        for name in self.container.list_objects(prefix=path):
            name = name[path_len:]
            slash = name[1:-1].find('/') + 1
            if slash:
                dirs.add(name[:slash])
            elif name:
                files.append(name)
        dirs = list(dirs)
        dirs.sort()
        return (dirs, files)

    def size(self, name):
        """
        Returns the total size, in bytes, of the file specified by name.
        """
        return self._get_cloud_obj(name).size

    def url(self, name):
        """
        Returns an absolute URL where the file's contents can be accessed
        directly by a web browser.
        """
        return '%s/%s' % (self.container_url, name)

class CloudFilesStorageFile(File):
    closed = False

    def __init__(self, storage, name, *args, **kwargs):
        self._storage = storage
        super(CloudFilesStorageFile, self).__init__(file=None, name=name,
                                                    *args, **kwargs)

    def _get_size(self):
        if not hasattr(self, '_size'):
            self._size = self._storage.size(self.name)
        return self._size

    def _set_size(self, size):
        self._size = size

    size = property(_get_size, _set_size)

    def _get_file(self):
        if not hasattr(self, '_file'):
            self._file = self._storage._get_cloud_obj(self.name)
        return self._file

    def _set_file(self, value):
        if value is None:
            if hasattr(self, '_file'):
                del self._file
        else:
            self._file = value

    file = property(_get_file, _set_file)

    def read(self, num_bytes=None):
        data = self.file.read(size=num_bytes or -1, offset=self._pos)
        self._pos += len(data)
        return data

    def open(self, *args, **kwargs):
        """
        Open the cloud file object.
        """
        self.file
        self._pos = 0

    def close(self, *args, **kwargs):
        self._pos = 0

    @property
    def closed(self):
        return not hasattr(self, '_file')

    def seek(self, pos):
        self._pos = pos
<commit_msg>Added ThreadSafeCloudFilesStorage, a thread safe subclass of CloudFilesStorage<commit_after>"""
Custom storage for django with Mosso Cloud Files backend.
Created by Rich Leland <rich@richleland.com>.
"""
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from django.core.files import File
from django.core.files.storage import Storage
from django.utils.text import get_valid_filename

try:
    import cloudfiles
    from cloudfiles.errors import NoSuchObject
except ImportError:
    raise ImproperlyConfigured("Could not load cloudfiles dependency. See "
                               "http://www.mosso.com/cloudfiles.jsp.")

# TODO: implement TTL into cloudfiles methods
CLOUDFILES_TTL = getattr(settings, 'CLOUDFILES_TTL', 600)

def cloudfiles_upload_to(self, filename):
    """
    Simple, custom upload_to because Cloud Files doesn't support
    nested containers (directories).

    Actually found this out from @minter:
    @richleland The Cloud Files APIs do support pseudo-subdirectories, by
    creating zero-byte files with type application/directory.

    May implement in a future version.
    """
    return get_valid_filename(filename)

class CloudFilesStorage(Storage):
    """
    Custom storage for Mosso Cloud Files.
    """
    default_quick_listdir = True

    def __init__(self, username=None, api_key=None, container=None,
                 connection_kwargs=None):
        """
        Initialize the settings for the connection and container.
        """
        self.username = username or settings.CLOUDFILES_USERNAME
        self.api_key = api_key or settings.CLOUDFILES_API_KEY
        self.container_name = container or settings.CLOUDFILES_CONTAINER
        self.connection_kwargs = connection_kwargs or {}

    def __getstate__(self):
        """
        Return a picklable representation of the storage.
        """
        return dict(username=self.username,
                    api_key=self.api_key,
                    container_name=self.container_name,
                    connection_kwargs=self.connection_kwargs)

    def _get_connection(self):
        if not hasattr(self, '_connection'):
            self._connection = cloudfiles.get_connection(self.username,
                                    self.api_key, **self.connection_kwargs)
        return self._connection

    def _set_connection(self, value):
        self._connection = value

    connection = property(_get_connection, _set_connection)

    def _get_container(self):
        if not hasattr(self, '_container'):
            self.container = self.connection.get_container(
                                                        self.container_name)
        return self._container

    def _set_container(self, container):
        """
        Set the container, making it publicly available (on Limelight CDN) if
        it is not already.
        """
        if not container.is_public():
            container.make_public()
        if hasattr(self, '_container_public_uri'):
            delattr(self, '_container_public_uri')
        self._container = container

    container = property(_get_container, _set_container)

    def _get_container_url(self):
        if not hasattr(self, '_container_public_uri'):
            self._container_public_uri = self.container.public_uri()
        return self._container_public_uri

    container_url = property(_get_container_url)

    def _get_cloud_obj(self, name):
        """
        Helper function to get retrieve the requested Cloud Files Object.
        """
        return self.container.get_object(name)

    def _open(self, name, mode='rb'):
        """
        Return the CloudFilesStorageFile.
        """
        return CloudFilesStorageFile(storage=self, name=name)

    def _save(self, name, content):
        """
        Use the Cloud Files service to write ``content`` to a remote file
        (called ``name``).
        """

        cloud_obj = self.container.create_object(name)
        cloud_obj.size = content.size

        content.open()
        # If the content type is available, pass it in directly rather than
        # getting the cloud object to try to guess.
        if hasattr(content.file, 'content_type'):
            cloud_obj.content_type = content.file.content_type
        cloud_obj.send(content)
        content.close()
        return name

    def delete(self, name):
        """
        Deletes the specified file from the storage system.
        """
        # If the file exists, delete it.
        if self.exists(name):
            self.container.delete_object(name)

    def exists(self, name):
        """
        Returns True if a file referenced by the given name already exists in
        the storage system, or False if the name is available for a new file.
        """
        try:
            self._get_cloud_obj(name)
            return True
        except NoSuchObject:
            return False

    def listdir(self, path):
        """
        Lists the contents of the specified path, returning a 2-tuple; the
        first being an empty list of directories (not available for quick-
        listing), the second being a list of filenames.

        If the list of directories is required, use the full_listdir method.
        """
        files = []
        if path and not path.endswith('/'):
            path = '%s/' % path
        path_len = len(path)
        for name in self.container.list_objects(path=path):
            files.append(name[path_len:])
        return ([], files)

    def full_listdir(self, path):
        """
        Lists the contents of the specified path, returning a 2-tuple of lists;
        the first item being directories, the second item being files.

        On large containers, this may be a slow operation for root containers
        because every single object must be returned (cloudfiles does not
        provide an explicit way of listing directories).
        """
        dirs = set()
        files = []
        if path and not path.endswith('/'):
            path = '%s/' % path
        path_len = len(path)
        for name in self.container.list_objects(prefix=path):
            name = name[path_len:]
            slash = name[1:-1].find('/') + 1
            if slash:
                dirs.add(name[:slash])
            elif name:
                files.append(name)
        dirs = list(dirs)
        dirs.sort()
        return (dirs, files)

    def size(self, name):
        """
        Returns the total size, in bytes, of the file specified by name.
        """
        return self._get_cloud_obj(name).size

    def url(self, name):
        """
        Returns an absolute URL where the file's contents can be accessed
        directly by a web browser.
        """
        return '%s/%s' % (self.container_url, name)

class CloudFilesStorageFile(File):
    closed = False

    def __init__(self, storage, name, *args, **kwargs):
        self._storage = storage
        super(CloudFilesStorageFile, self).__init__(file=None, name=name,
                                                    *args, **kwargs)

    def _get_size(self):
        if not hasattr(self, '_size'):
            self._size = self._storage.size(self.name)
        return self._size

    def _set_size(self, size):
        self._size = size

    size = property(_get_size, _set_size)

    def _get_file(self):
        if not hasattr(self, '_file'):
            self._file = self._storage._get_cloud_obj(self.name)
        return self._file

    def _set_file(self, value):
        if value is None:
            if hasattr(self, '_file'):
                del self._file
        else:
            self._file = value

    file = property(_get_file, _set_file)

    def read(self, num_bytes=None):
        data = self.file.read(size=num_bytes or -1, offset=self._pos)
        self._pos += len(data)
        return data

    def open(self, *args, **kwargs):
        """
        Open the cloud file object.
        """
        self.file
        self._pos = 0

    def close(self, *args, **kwargs):
        self._pos = 0

    @property
    def closed(self):
        return not hasattr(self, '_file')

    def seek(self, pos):
        self._pos = pos

class ThreadSafeCloudFilesStorage(CloudFilesStorage):
    """
    Extends CloudFilesStorage to make it thread safer.

    As long as you don't pass container or cloud objects
    between threads, you'll be thread safe.

    Uses one cloudfiles connection per thread.
    """

    def __init__(self, *args, **kwargs):
        super(ThreadSafeCloudFilesStorage, self).__init__(*args, **kwargs)

        import threading
        self.local_cache = threading.local()

    def _get_connection(self):
        if not hasattr(self.local_cache, 'connection'):
            connection = cloudfiles.get_connection(self.username,
                                    self.api_key, **self.connection_kwargs)
            self.local_cache.connection = connection

        return self.local_cache.connection

    connection = property(_get_connection, CloudFilesStorage._set_connection)

    def _get_container(self):
        if not hasattr(self.local_cache, 'container'):
            container = self.connection.get_container(self.container_name)
            self.local_cache.container = container

        return self.local_cache.container

    container = property(_get_container, CloudFilesStorage._set_container)

<|endoftext|>


<commit_before>"""Zip Tool.

@see: Cake Build System (http://sourceforge.net/projects/cake-build)
@copyright: Copyright (c) 2010 Lewis Baker, Stuart McMahon.
@license: Licensed under the MIT license.
"""

from cake.library import Tool, getPathAndTask
import cake.filesys
import zipfile
import os
import os.path
import time

def _isZipDirectory(zipInfo):

  return (zipInfo.external_attr & 0x00000010L) != 0L # FILE_ATTRIBUTE_DIRECTORY

def _shouldExtractFile(engine, absTargetFile, zipTime, onlyNewer):

  if engine.forceBuild:
    return "rebuild has been forced"

  if not onlyNewer:
    return "onlyNewer is set to False" # Always rebuild

  try:
    mtime = os.stat(absTargetFile).st_mtime
  except EnvironmentError:
    return "it doesn't exist"

  if zipTime != mtime:
    # Assume the zip and the extract are the same file.
    return "it has been changed"

  return None

def _extractFile(engine, zipFile, zipPath, zipInfo, targetDir, absTargetDir, onlyNewer):
  """Extract the ZipInfo object to a physical file at targetDir.
  """
  targetFile = os.path.join(targetDir, zipInfo.filename)
  absTargetFile = os.path.join(absTargetDir, zipInfo.filename)

  if _isZipDirectory(zipInfo):
    # The zip info corresponds to a directory.
    cake.filesys.makeDirs(absTargetFile)
  else:
    # The zip info corresponds to a file.
    year, month, day, hour, minute, second = zipInfo.date_time
    zipTime = time.mktime(time.struct_time((year, month, day, hour, minute, second, 0, 0, 0)))

    reasonToBuild = _shouldExtractFile(engine, absTargetFile, zipTime, onlyNewer)
    if reasonToBuild is None:
      return # Target is up to date

    engine.logger.outputDebug(
      "reason",
      "Extracting '" + targetFile + "' because " + reasonToBuild + ".\n",
      )

    engine.logger.outputInfo("Extracting %s\n" % targetFile)

    try:
      cake.filesys.writeFile(absTargetFile, zipFile.read(zipInfo.filename))
    except Exception, e:
      engine.raiseError(
        "Failed to extract file %s from zip %s: %s\n" % (
          zipInfo.filename,
          zipPath,
          str(e),
          ),
        )

    # Set the file modification time to match the zip time
    os.utime(absTargetFile, (zipTime, zipTime))

def _writeFile(configuration, file, sourcePath, targetZipPath, targetFilePath):
  absSourcePath = configuration.abspath(sourcePath)
  sourceFilePath = os.path.join(sourcePath, targetFilePath)
  absSourceFilePath = os.path.join(absSourcePath, targetFilePath)
  targetFilePath = targetFilePath.replace("\\", "/") # Zips use forward slashes
  utcTime = time.gmtime(os.stat(absSourceFilePath).st_mtime)

  configuration.engine.logger.outputInfo("Adding %s to %s\n" % (sourceFilePath, targetZipPath))

  if os.path.isdir(absSourceFilePath):
    if not targetFilePath.endswith("/"):
      targetFilePath += "/" # Trailing slash denotes directory for some zip packages

    zi = zipfile.ZipInfo(targetFilePath, utcTime[0:6])
    zi.compress_type = zipfile.ZIP_DEFLATED
    zi.external_attr = 0x00000010L # FILE_ATTRIBUTE_DIRECTORY
    file.writestr(zi, "")
  else:
    f = open(absSourceFilePath, "rb")
    try:
      data = f.read()
    finally:
      f.close()

    zi = zipfile.ZipInfo(targetFilePath, utcTime[0:6])
    zi.compress_type = zipfile.ZIP_DEFLATED
    zi.external_attr = 0x00000020L # FILE_ATTRIBUTE_ARCHIVE
    file.writestr(zi, data)

def _walkTree(path):
  """Recursively walk a directory tree.
  """
  for dirPath, dirNames, fileNames in os.walk(path):
    for name in dirNames:
      yield os.path.join(dirPath, name)

    for name in fileNames:
      yield os.path.join(dirPath, name)

def _shouldCompress(
  configuration,
  sourcePath,
  targetPath,
  toZip,
  onlyNewer,
  removeStale,
  ):

  if configuration.engine.forceBuild:
    return None, "rebuild has been forced"

  if not onlyNewer:
    return None, "onlyNewer is set to False" # Always rebuild

  absSourcePath = configuration.abspath(sourcePath)
  absTargetPath = configuration.abspath(targetPath)

  # Try to open an existing zip file
  try:
    file = zipfile.ZipFile(absTargetPath, "r")
    try:
      zipInfos = file.infolist()
    finally:
      file.close()

    # Build a list of files/dirs in the current zip
    fromZip = {}
    for zipInfo in zipInfos:
      path = os.path.normpath(os.path.normcase(zipInfo.filename))
      fromZip[path] = zipInfo
  except EnvironmentError:
    # File doesn't exist or is invalid
    return None, "'" + targetPath + "' doesn't exist"

  if onlyNewer:
    for casedPath, originalPath in toZip.iteritems():
      zipInfo = fromZip.get(casedPath, None)

      # Not interested in modified directories
      if zipInfo is not None and not _isZipDirectory(zipInfo):
        absSourceFilePath = os.path.join(absSourcePath, originalPath)
        utcTime = time.gmtime(os.stat(absSourceFilePath).st_mtime)
        zipTime = utcTime[0:5] + (
          utcTime[5] & 0xFE, # Zip only saves 2 second resolution
          )
        if zipTime != zipInfo.date_time:
          sourceFilePath = os.path.join(sourcePath, originalPath)
          # We must recreate the zip to update files
          return None, "'" + sourceFilePath + "' has been changed"

  if removeStale:
    for path, zipInfo in fromZip.iteritems():
      if path not in toZip:
        # We must recreate the zip to remove files
        sourceFilePath = os.path.join(sourcePath, zipInfo.filename)
        return None, "'" + sourceFilePath + "' has been removed"

  toAppend = []
  reasonToBuild = None
  for casedPath, originalPath in toZip.iteritems():
    if casedPath not in fromZip:
      toAppend.append(originalPath)
      if reasonToBuild is None:
        sourceFilePath = os.path.join(sourcePath, originalPath)
        reasonToBuild = "'" + sourceFilePath + "' is not in zip"

  return toAppend, reasonToBuild

def _getFilesToCompress(configuration, sourcePath, includeMatch, excludeMatch):

  absSourcePath = configuration.abspath(sourcePath)

  toZip = {}
  if os.path.isdir(absSourcePath):
    firstChar = len(absSourcePath)+1
    for path in _walkTree(absSourcePath):
      path = path[firstChar:] # Strip the search dir name
      if includeMatch is not None and not includeMatch(path):
        continue
      if excludeMatch is not None and excludeMatch(path):
        continue
      toZip[os.path.normcase(path)] = path
  else:
    toZip[os.path.normcase(path)] = path

  return toZip

class ZipTool(Tool):

  def extract(
    self,
    targetDir,
    source,
    onlyNewer=True,
    removeStale=False,
    includeMatch=None,
    excludeMatch=None,
    ):
    """Extract all files in a Zip to the specified path.

    @param targetDir: The directory to extract files to.
    @type targetDir: string
    @param source: Path to the zip file to extract files from.
    @type source: string
    @param onlyNewer: Only extract files that are newer than those in
    the target directory.
    @type onlyNewer: bool
    @param removeStale: Remove files and directories in the target
    directory that no longer exist in the zip.
    @type removeStale: bool
    @param includeMatch: A callable used to decide whether to include
    certain files in the extraction. This could be a python callable that
    returns True to include the file or False to exclude it, or a regular
    expression function such as re.compile().match or re.match.
    @type includeMatch: any callable
    @param excludeMatch: A callable used to decide whether to exclude certain
    files from being extracted. This could be a python callable that
    returns True to exclude the file or False to include it, or a regular
    expression function such as re.compile().match or re.match.
    @type excludeMatch: any callable

    @return: A task that will complete when the extraction has finished.
    @rtype: L{Task}
    """
    if not isinstance(targetDir, basestring):
      raise TypeError("targetDir must be a string")

    sourcePath, sourceTask = getPathAndTask(source)

    engine = self.engine
    configuration = self.configuration

    def doIt():
      absTargetDir = configuration.abspath(targetDir)
      file = zipfile.ZipFile(configuration.abspath(sourcePath), "r")
      try:
        zipInfos = file.infolist()

        if includeMatch is not None:
          newZipInfos = []
          for zipInfo in zipInfos:
            if includeMatch(zipInfo.filename):
              newZipInfos.append(zipInfo)
          zipInfos = newZipInfos

        if excludeMatch is not None:
          newZipInfos = []
          for zipInfo in zipInfos:
            if not excludeMatch(zipInfo.filename):
              newZipInfos.append(zipInfo)
          zipInfos = newZipInfos

        if removeStale:
          zipFiles = set()
          for zipInfo in zipInfos:
            zipFiles.add(os.path.normcase(os.path.normpath(zipInfo.filename)))

          firstChar = len(absTargetDir) + 1
          for absPath in _walkTree(absTargetDir):
            path = absPath[firstChar:] # Strip the search dir name
            if os.path.normcase(path) not in zipFiles:
              engine.logger.outputInfo(
                "Deleting %s\n" % os.path.join(targetDir, path),
                )
              if os.path.isdir(absPath):
                cake.filesys.removeTree(absPath)
              else:
                cake.filesys.remove(absPath)

        for zipinfo in zipInfos:
          _extractFile(engine, file, sourcePath, zipinfo, targetDir, absTargetDir, onlyNewer)
      finally:
        file.close()

    task = engine.createTask(doIt)
    task.startAfter(sourceTask)
    return task

  def compress(
    self,
    target,
    source,
    onlyNewer=True,
    removeStale=True,
    includeMatch=None,
    excludeMatch=None,
    ):
    """Extract all files in a Zip to the specified path.

    @param target: Path to the zip file to add files to.
    @type target: string
    @param source: Path to the source file or directory to add.
    @type source: string
    @param onlyNewer: Only add files that are newer than those in
    the zip file. Otherwise all files are re-added every time.
    @type onlyNewer: bool
    @param removeStale: Remove files and directories in the zip
    file that no longer exist in the source directory.
    @type removeStale: bool
    @param includeMatch: A callable used to decide whether to include
    certain files in the zip file. This could be a python callable that
    returns True to include the file or False to exclude it, or a regular
    expression function such as re.compile().match or re.match.
    @type includeMatch: any callable
    @param excludeMatch: A callable used to decide whether to exclude certain
    files from the zip file. This could be a python callable that
    returns True to exclude the file or False to include it, or a regular
    expression function such as re.compile().match or re.match.
    @type excludeMatch: any callable

    @return: A task that will complete when the compression has finished.
    @rtype: L{Task}
    """
    if not isinstance(target, basestring):
      raise TypeError("target must be a string")

    engine = self.engine
    configuration = self.configuration

    sourcePath, sourceTask = getPathAndTask(source)

    def doIt():

      # Build a list of files/dirs to zip
      toZip = _getFilesToCompress(configuration, sourcePath, includeMatch, excludeMatch)

      # Figure out if we need to rebuild/append
      toAppend, reasonToBuild = _shouldCompress(
        configuration,
        sourcePath,
        target,
        toZip,
        onlyNewer,
        removeStale,
        )
      if reasonToBuild is None:
        return # Target is up to date

      engine.logger.outputDebug(
        "reason",
        "Rebuilding '" + target + "' because " + reasonToBuild + ".\n",
        )

      absTargetPath = configuration.abspath(target)
      absTargetTmpPath = absTargetPath + ".tmp"
      if toAppend is None:
        # Recreate zip
        cake.filesys.makeDirs(os.path.dirname(absTargetTmpPath))
        f = open(absTargetTmpPath, "wb")
        try:
          zipFile = zipfile.ZipFile(f, "w")
          for originalPath in toZip.itervalues():
            _writeFile(configuration, zipFile, sourcePath, target, originalPath)
          zipFile.close()
        finally:
          f.close()
        cake.filesys.renameFile(absTargetTmpPath, absTargetPath)
      else:
        # Append to existing zip
        cake.filesys.renameFile(absTargetPath, absTargetTmpPath)
        f = open(absTargetTmpPath, "r+b")
        try:
          zipFile = zipfile.ZipFile(f, "a")
          for originalPath in toAppend:
            _writeFile(configuration, zipFile, sourcePath, target, originalPath)
          zipFile.close()
        finally:
          f.close()
        cake.filesys.renameFile(absTargetTmpPath, absTargetPath)

    task = engine.createTask(doIt)
    task.startAfter(sourceTask)
    return task
    <commit_msg>- Fix for ZipTool trying deleting files after it had corrupted the filenames by stripping too much off. Now I am normalising slashes so I know there will be no slash on the end of a path when calculating the 'firstChar'.<commit_after>"""Zip Tool.

@see: Cake Build System (http://sourceforge.net/projects/cake-build)
@copyright: Copyright (c) 2010 Lewis Baker, Stuart McMahon.
@license: Licensed under the MIT license.
"""

from cake.library import Tool, getPathAndTask
import cake.filesys
import zipfile
import os
import os.path
import time

def _isZipDirectory(zipInfo):

  return (zipInfo.external_attr & 0x00000010L) != 0L # FILE_ATTRIBUTE_DIRECTORY

def _shouldExtractFile(engine, absTargetFile, zipTime, onlyNewer):

  if engine.forceBuild:
    return "rebuild has been forced"

  if not onlyNewer:
    return "onlyNewer is set to False" # Always rebuild

  try:
    mtime = os.stat(absTargetFile).st_mtime
  except EnvironmentError:
    return "it doesn't exist"

  if zipTime != mtime:
    # Assume the zip and the extract are the same file.
    return "it has been changed"

  return None

def _extractFile(engine, zipFile, zipPath, zipInfo, targetDir, absTargetDir, onlyNewer):
  """Extract the ZipInfo object to a physical file at targetDir.
  """
  targetFile = os.path.join(targetDir, zipInfo.filename)
  absTargetFile = os.path.join(absTargetDir, zipInfo.filename)

  if _isZipDirectory(zipInfo):
    # The zip info corresponds to a directory.
    cake.filesys.makeDirs(absTargetFile)
  else:
    # The zip info corresponds to a file.
    year, month, day, hour, minute, second = zipInfo.date_time
    zipTime = time.mktime(time.struct_time((year, month, day, hour, minute, second, 0, 0, 0)))

    reasonToBuild = _shouldExtractFile(engine, absTargetFile, zipTime, onlyNewer)
    if reasonToBuild is None:
      return # Target is up to date

    engine.logger.outputDebug(
      "reason",
      "Extracting '" + targetFile + "' because " + reasonToBuild + ".\n",
      )

    engine.logger.outputInfo("Extracting %s\n" % targetFile)

    try:
      cake.filesys.writeFile(absTargetFile, zipFile.read(zipInfo.filename))
    except Exception, e:
      engine.raiseError(
        "Failed to extract file %s from zip %s: %s\n" % (
          zipInfo.filename,
          zipPath,
          str(e),
          ),
        )

    # Set the file modification time to match the zip time
    os.utime(absTargetFile, (zipTime, zipTime))

def _writeFile(configuration, file, sourcePath, targetZipPath, targetFilePath):
  absSourcePath = configuration.abspath(sourcePath)
  sourceFilePath = os.path.join(sourcePath, targetFilePath)
  absSourceFilePath = os.path.join(absSourcePath, targetFilePath)
  targetFilePath = targetFilePath.replace("\\", "/") # Zips use forward slashes
  utcTime = time.gmtime(os.stat(absSourceFilePath).st_mtime)

  configuration.engine.logger.outputInfo("Adding %s to %s\n" % (sourceFilePath, targetZipPath))

  if os.path.isdir(absSourceFilePath):
    if not targetFilePath.endswith("/"):
      targetFilePath += "/" # Trailing slash denotes directory for some zip packages

    zi = zipfile.ZipInfo(targetFilePath, utcTime[0:6])
    zi.compress_type = zipfile.ZIP_DEFLATED
    zi.external_attr = 0x00000010L # FILE_ATTRIBUTE_DIRECTORY
    file.writestr(zi, "")
  else:
    f = open(absSourceFilePath, "rb")
    try:
      data = f.read()
    finally:
      f.close()

    zi = zipfile.ZipInfo(targetFilePath, utcTime[0:6])
    zi.compress_type = zipfile.ZIP_DEFLATED
    zi.external_attr = 0x00000020L # FILE_ATTRIBUTE_ARCHIVE
    file.writestr(zi, data)

def _walkTree(path):
  """Recursively walk a directory tree.
  """
  for dirPath, dirNames, fileNames in os.walk(path):
    for name in dirNames:
      yield os.path.join(dirPath, name)

    for name in fileNames:
      yield os.path.join(dirPath, name)

def _shouldCompress(
  configuration,
  sourcePath,
  targetPath,
  toZip,
  onlyNewer,
  removeStale,
  ):

  if configuration.engine.forceBuild:
    return None, "rebuild has been forced"

  if not onlyNewer:
    return None, "onlyNewer is set to False" # Always rebuild

  absSourcePath = configuration.abspath(sourcePath)
  absTargetPath = configuration.abspath(targetPath)

  # Try to open an existing zip file
  try:
    file = zipfile.ZipFile(absTargetPath, "r")
    try:
      zipInfos = file.infolist()
    finally:
      file.close()

    # Build a list of files/dirs in the current zip
    fromZip = {}
    for zipInfo in zipInfos:
      path = os.path.normpath(os.path.normcase(zipInfo.filename))
      fromZip[path] = zipInfo
  except EnvironmentError:
    # File doesn't exist or is invalid
    return None, "'" + targetPath + "' doesn't exist"

  if onlyNewer:
    for casedPath, originalPath in toZip.iteritems():
      zipInfo = fromZip.get(casedPath, None)

      # Not interested in modified directories
      if zipInfo is not None and not _isZipDirectory(zipInfo):
        absSourceFilePath = os.path.join(absSourcePath, originalPath)
        utcTime = time.gmtime(os.stat(absSourceFilePath).st_mtime)
        zipTime = utcTime[0:5] + (
          utcTime[5] & 0xFE, # Zip only saves 2 second resolution
          )
        if zipTime != zipInfo.date_time:
          sourceFilePath = os.path.join(sourcePath, originalPath)
          # We must recreate the zip to update files
          return None, "'" + sourceFilePath + "' has been changed"

  if removeStale:
    for path, zipInfo in fromZip.iteritems():
      if path not in toZip:
        # We must recreate the zip to remove files
        sourceFilePath = os.path.join(sourcePath, zipInfo.filename)
        return None, "'" + sourceFilePath + "' has been removed"

  toAppend = []
  reasonToBuild = None
  for casedPath, originalPath in toZip.iteritems():
    if casedPath not in fromZip:
      toAppend.append(originalPath)
      if reasonToBuild is None:
        sourceFilePath = os.path.join(sourcePath, originalPath)
        reasonToBuild = "'" + sourceFilePath + "' is not in zip"

  return toAppend, reasonToBuild

def _getFilesToCompress(configuration, sourcePath, includeMatch, excludeMatch):

  absSourcePath = configuration.abspath(sourcePath)

  toZip = {}
  if os.path.isdir(absSourcePath):
    # Remove any trailing slash
    searchDir = os.path.normpath(absSourcePath)
    firstChar = len(searchDir) + 1
    for path in _walkTree(searchDir):
      path = path[firstChar:] # Strip the search dir name
      if includeMatch is not None and not includeMatch(path):
        continue
      if excludeMatch is not None and excludeMatch(path):
        continue
      toZip[os.path.normcase(path)] = path
  else:
    toZip[os.path.normcase(path)] = path

  return toZip

class ZipTool(Tool):

  def extract(
    self,
    targetDir,
    source,
    onlyNewer=True,
    removeStale=False,
    includeMatch=None,
    excludeMatch=None,
    ):
    """Extract all files in a Zip to the specified path.

    @param targetDir: The directory to extract files to.
    @type targetDir: string
    @param source: Path to the zip file to extract files from.
    @type source: string
    @param onlyNewer: Only extract files that are newer than those in
    the target directory.
    @type onlyNewer: bool
    @param removeStale: Remove files and directories in the target
    directory that no longer exist in the zip.
    @type removeStale: bool
    @param includeMatch: A callable used to decide whether to include
    certain files in the extraction. This could be a python callable that
    returns True to include the file or False to exclude it, or a regular
    expression function such as re.compile().match or re.match.
    @type includeMatch: any callable
    @param excludeMatch: A callable used to decide whether to exclude certain
    files from being extracted. This could be a python callable that
    returns True to exclude the file or False to include it, or a regular
    expression function such as re.compile().match or re.match.
    @type excludeMatch: any callable

    @return: A task that will complete when the extraction has finished.
    @rtype: L{Task}
    """
    if not isinstance(targetDir, basestring):
      raise TypeError("targetDir must be a string")

    sourcePath, sourceTask = getPathAndTask(source)

    engine = self.engine
    configuration = self.configuration

    def doIt():
      absTargetDir = configuration.abspath(targetDir)
      file = zipfile.ZipFile(configuration.abspath(sourcePath), "r")
      try:
        zipInfos = file.infolist()

        if includeMatch is not None:
          newZipInfos = []
          for zipInfo in zipInfos:
            if includeMatch(zipInfo.filename):
              newZipInfos.append(zipInfo)
          zipInfos = newZipInfos

        if excludeMatch is not None:
          newZipInfos = []
          for zipInfo in zipInfos:
            if not excludeMatch(zipInfo.filename):
              newZipInfos.append(zipInfo)
          zipInfos = newZipInfos

        if removeStale:
          zipFiles = set()
          for zipInfo in zipInfos:
            zipFiles.add(os.path.normcase(os.path.normpath(zipInfo.filename)))

          searchDir = os.path.normpath(absTargetDir)
          firstChar = len(searchDir) + 1
          for absPath in _walkTree(searchDir):
            path = absPath[firstChar:] # Strip the search dir name
            if os.path.normcase(path) not in zipFiles:
              engine.logger.outputInfo(
                "Deleting %s\n" % os.path.join(targetDir, path),
                )
              if os.path.isdir(absPath):
                cake.filesys.removeTree(absPath)
              else:
                cake.filesys.remove(absPath)

        for zipinfo in zipInfos:
          _extractFile(engine, file, sourcePath, zipinfo, targetDir, absTargetDir, onlyNewer)
      finally:
        file.close()

    task = engine.createTask(doIt)
    task.startAfter(sourceTask)
    return task

  def compress(
    self,
    target,
    source,
    onlyNewer=True,
    removeStale=True,
    includeMatch=None,
    excludeMatch=None,
    ):
    """Extract all files in a Zip to the specified path.

    @param target: Path to the zip file to add files to.
    @type target: string
    @param source: Path to the source file or directory to add.
    @type source: string
    @param onlyNewer: Only add files that are newer than those in
    the zip file. Otherwise all files are re-added every time.
    @type onlyNewer: bool
    @param removeStale: Remove files and directories in the zip
    file that no longer exist in the source directory.
    @type removeStale: bool
    @param includeMatch: A callable used to decide whether to include
    certain files in the zip file. This could be a python callable that
    returns True to include the file or False to exclude it, or a regular
    expression function such as re.compile().match or re.match.
    @type includeMatch: any callable
    @param excludeMatch: A callable used to decide whether to exclude certain
    files from the zip file. This could be a python callable that
    returns True to exclude the file or False to include it, or a regular
    expression function such as re.compile().match or re.match.
    @type excludeMatch: any callable

    @return: A task that will complete when the compression has finished.
    @rtype: L{Task}
    """
    if not isinstance(target, basestring):
      raise TypeError("target must be a string")

    engine = self.engine
    configuration = self.configuration

    sourcePath, sourceTask = getPathAndTask(source)

    def doIt():

      # Build a list of files/dirs to zip
      toZip = _getFilesToCompress(configuration, sourcePath, includeMatch, excludeMatch)

      # Figure out if we need to rebuild/append
      toAppend, reasonToBuild = _shouldCompress(
        configuration,
        sourcePath,
        target,
        toZip,
        onlyNewer,
        removeStale,
        )
      if reasonToBuild is None:
        return # Target is up to date

      engine.logger.outputDebug(
        "reason",
        "Rebuilding '" + target + "' because " + reasonToBuild + ".\n",
        )

      absTargetPath = configuration.abspath(target)
      absTargetTmpPath = absTargetPath + ".tmp"
      if toAppend is None:
        # Recreate zip
        cake.filesys.makeDirs(os.path.dirname(absTargetTmpPath))
        f = open(absTargetTmpPath, "wb")
        try:
          zipFile = zipfile.ZipFile(f, "w")
          for originalPath in toZip.itervalues():
            _writeFile(configuration, zipFile, sourcePath, target, originalPath)
          zipFile.close()
        finally:
          f.close()
        cake.filesys.renameFile(absTargetTmpPath, absTargetPath)
      else:
        # Append to existing zip
        cake.filesys.renameFile(absTargetPath, absTargetTmpPath)
        f = open(absTargetTmpPath, "r+b")
        try:
          zipFile = zipfile.ZipFile(f, "a")
          for originalPath in toAppend:
            _writeFile(configuration, zipFile, sourcePath, target, originalPath)
          zipFile.close()
        finally:
          f.close()
        cake.filesys.renameFile(absTargetTmpPath, absTargetPath)

    task = engine.createTask(doIt)
    task.startAfter(sourceTask)
    return task
    <|endoftext|>