<commit_before># -*- coding: utf-8 -*-
#
# Copyright 2012-2015 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import abc
import logging
import operator
import os
import subprocess
import tempfile

from luigi import six

import luigi
import luigi.contrib.hadoop
from luigi.target import FileAlreadyExists, FileSystemTarget
from luigi.task import flatten

if six.PY3:
    unicode = str

logger = logging.getLogger('luigi-interface')

class HiveCommandError(RuntimeError):

    def __init__(self, message, out=None, err=None):
        super(HiveCommandError, self).__init__(message, out, err)
        self.message = message
        self.out = out
        self.err = err

def load_hive_cmd():
    return luigi.configuration.get_config().get('hive', 'command', 'hive')

def get_hive_syntax():
    return luigi.configuration.get_config().get('hive', 'release', 'cdh4')

def run_hive(args, check_return_code=True):
    """
    Runs the `hive` from the command line, passing in the given args, and
    returning stdout.

    With the apache release of Hive, so of the table existence checks
    (which are done using DESCRIBE do not exit with a return code of 0
    so we need an option to ignore the return code and just return stdout for parsing
    """
    cmd = [load_hive_cmd()] + args
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = p.communicate()
    if check_return_code and p.returncode != 0:
        raise HiveCommandError("Hive command: {0} failed with error code: {1}".format(" ".join(cmd), p.returncode),
                               stdout, stderr)
    return stdout

def run_hive_cmd(hivecmd, check_return_code=True):
    """
    Runs the given hive query and returns stdout.
    """
    return run_hive(['-e', hivecmd], check_return_code)

def run_hive_script(script):
    """
    Runs the contents of the given script in hive and returns stdout.
    """
    if not os.path.isfile(script):
        raise RuntimeError("Hive script: {0} does not exist.".format(script))
    return run_hive(['-f', script])

@six.add_metaclass(abc.ABCMeta)
class HiveClient(object):  # interface

    @abc.abstractmethod
    def table_location(self, table, database='default', partition=None):
        """
        Returns location of db.table (or db.table.partition). partition is a dict of partition key to
        value.
        """
        pass

    @abc.abstractmethod
    def table_schema(self, table, database='default'):
        """
        Returns list of [(name, type)] for each column in database.table.
        """
        pass

    @abc.abstractmethod
    def table_exists(self, table, database='default', partition=None):
        """
        Returns true if db.table (or db.table.partition) exists. partition is a dict of partition key to
        value.
        """
        pass

    @abc.abstractmethod
    def partition_spec(self, partition):
        """ Turn a dict into a string partition specification """
        pass

class HiveCommandClient(HiveClient):
    """
    Uses `hive` invocations to find information.
    """

    def table_location(self, table, database='default', partition=None):
        cmd = "use {0}; describe formatted {1}".format(database, table)
        if partition is not None:
            cmd += " PARTITION ({0})".format(self.partition_spec(partition))

        stdout = run_hive_cmd(cmd)

        for line in stdout.split("\n"):
            if "Location:" in line:
                return line.split("\t")[1]

    def table_exists(self, table, database='default', partition=None):
        if partition is None:
            stdout = run_hive_cmd('use {0}; show tables like "{1}";'.format(database, table))

            return stdout and table.lower() in stdout
        else:
            stdout = run_hive_cmd("""use %s; show partitions %s partition
                                (%s)""" % (database, table, self.partition_spec(partition)))

            if stdout:
                return True
            else:
                return False

    def table_schema(self, table, database='default'):
        describe = run_hive_cmd("use {0}; describe {1}".format(database, table))
        if not describe or "does not exist" in describe:
            return None
        return [tuple([x.strip() for x in line.strip().split("\t")]) for line in describe.strip().split("\n")]

    def partition_spec(self, partition):
        """
        Turns a dict into the a Hive partition specification string.
        """
        return ','.join(["`{0}`='{1}'".format(k, v) for (k, v) in
                         sorted(six.iteritems(partition), key=operator.itemgetter(0))])

class ApacheHiveCommandClient(HiveCommandClient):
    """
    A subclass for the HiveCommandClient to (in some cases) ignore the return code from
    the hive command so that we can just parse the output.
    """

    def table_schema(self, table, database='default'):
        describe = run_hive_cmd("use {0}; describe {1}".format(database, table), False)
        if not describe or "Table not found" in describe:
            return None
        return [tuple([x.strip() for x in line.strip().split("\t")]) for line in describe.strip().split("\n")]

class MetastoreClient(HiveClient):

    def table_location(self, table, database='default', partition=None):
        with HiveThriftContext() as client:
            if partition is not None:
                try:
                    partition_str = self.partition_spec(partition)
                    thrift_table = client.get_partition_by_name(database, table, partition_str)
                except NoSuchObjectException as e:
                    return ''
            else:
                thrift_table = client.get_table(database, table)
            return thrift_table.sd.location

    def table_exists(self, table, database='default', partition=None):
        with HiveThriftContext() as client:
            if partition is None:
                return table in client.get_all_tables(database)
            else:
                return partition in self._existing_partitions(table, database, client)

    def _existing_partitions(self, table, database, client):
        def _parse_partition_string(partition_string):
            partition_def = {}
            for part in partition_string.split("/"):
                name, value = part.split("=")
                partition_def[name] = value
            return partition_def

        # -1 is max_parts, the # of partition names to return (-1 = unlimited)
        partition_strings = client.get_partition_names(database, table, -1)
        return [_parse_partition_string(existing_partition) for existing_partition in partition_strings]

    def table_schema(self, table, database='default'):
        with HiveThriftContext() as client:
            return [(field_schema.name, field_schema.type) for field_schema in client.get_schema(database, table)]

    def partition_spec(self, partition):
        return "/".join("%s=%s" % (k, v) for (k, v) in sorted(six.iteritems(partition), key=operator.itemgetter(0)))

class HiveThriftContext(object):
    """
    Context manager for hive metastore client.
    """

    def __enter__(self):
        try:
            from thrift.transport import TSocket
            from thrift.transport import TTransport
            from thrift.protocol import TBinaryProtocol
            # Note that this will only work with a CDH release.
            # This uses the thrift bindings generated by the ThriftHiveMetastore service in Beeswax.
            # If using the Apache release of Hive this import will fail.
            from hive_metastore import ThriftHiveMetastore
            config = luigi.configuration.get_config()
            host = config.get('hive', 'metastore_host')
            port = config.getint('hive', 'metastore_port')
            transport = TSocket.TSocket(host, port)
            transport = TTransport.TBufferedTransport(transport)
            protocol = TBinaryProtocol.TBinaryProtocol(transport)
            transport.open()
            self.transport = transport
            return ThriftHiveMetastore.Client(protocol)
        except ImportError as e:
            raise Exception('Could not import Hive thrift library:' + str(e))

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.transport.close()

def get_default_client():
    syntax = get_hive_syntax()
    if syntax == "apache":
        return ApacheHiveCommandClient()
    elif syntax == "metastore":
        return MetastoreClient()
    else:
        return HiveCommandClient()

client = get_default_client()

class HiveQueryTask(luigi.contrib.hadoop.BaseHadoopJobTask):
    """
    Task to run a hive query.
    """

    # by default, we let hive figure these out.
    n_reduce_tasks = None
    bytes_per_reducer = None
    reducers_max = None

    @abc.abstractmethod
    def query(self):
        """ Text of query to run in hive """
        raise RuntimeError("Must implement query!")

    def hiverc(self):
        """
        Location of an rc file to run before the query
        if hiverc-location key is specified in luigi.cfg, will default to the value there
        otherwise returns None.

        Returning a list of rc files will load all of them in order.
        """
        return luigi.configuration.get_config().get('hive', 'hiverc-location', default=None)

    def hiveconfs(self):
        """
        Returns an dict of key=value settings to be passed along
        to the hive command line via --hiveconf. By default, sets
        mapred.job.name to task_id and if not None, sets:

        * mapred.reduce.tasks (n_reduce_tasks)
        * mapred.fairscheduler.pool (pool) or mapred.job.queue.name (pool)
        * hive.exec.reducers.bytes.per.reducer (bytes_per_reducer)
        * hive.exec.reducers.max (reducers_max)
        """
        jcs = {}
        jcs['mapred.job.name'] = self.task_id
        if self.n_reduce_tasks is not None:
            jcs['mapred.reduce.tasks'] = self.n_reduce_tasks
        if self.pool is not None:
            # Supporting two schedulers: fair (default) and capacity using the same option
            scheduler_type = luigi.configuration.get_config().get('hadoop', 'scheduler', 'fair')
            if scheduler_type == 'fair':
                jcs['mapred.fairscheduler.pool'] = self.pool
            elif scheduler_type == 'capacity':
                jcs['mapred.job.queue.name'] = self.pool
        if self.bytes_per_reducer is not None:
            jcs['hive.exec.reducers.bytes.per.reducer'] = self.bytes_per_reducer
        if self.reducers_max is not None:
            jcs['hive.exec.reducers.max'] = self.reducers_max
        return jcs

    def job_runner(self):
        return HiveQueryRunner()

class HiveQueryRunner(luigi.contrib.hadoop.JobRunner):
    """
    Runs a HiveQueryTask by shelling out to hive.
    """

    def prepare_outputs(self, job):
        """
        Called before job is started.

        If output is a `FileSystemTarget`, create parent directories so the hive command won't fail
        """
        outputs = flatten(job.output())
        for o in outputs:
            if isinstance(o, FileSystemTarget):
                parent_dir = os.path.dirname(o.path)
                if parent_dir and not o.fs.exists(parent_dir):
                    logger.info("Creating parent directory %r", parent_dir)
                    try:
                        # there is a possible race condition
                        # which needs to be handled here
                        o.fs.mkdir(parent_dir)
                    except FileAlreadyExists:
                        pass

    def run_job(self, job, tracking_url_callback=None):
        self.prepare_outputs(job)
        with tempfile.NamedTemporaryFile() as f:
            query = job.query()
            if isinstance(query, unicode):
                query = query.encode('utf8')
            f.write(query)
            f.flush()
            arglist = [load_hive_cmd(), '-f', f.name]
            hiverc = job.hiverc()
            if hiverc:
                if isinstance(hiverc, str):
                    hiverc = [hiverc]
                for rcfile in hiverc:
                    arglist += ['-i', rcfile]
            if job.hiveconfs():
                for k, v in six.iteritems(job.hiveconfs()):
                    arglist += ['--hiveconf', '{0}={1}'.format(k, v)]

            logger.info(arglist)
            return luigi.contrib.hadoop.run_and_track_hadoop_job(arglist, tracking_url_callback)

class HiveTableTarget(luigi.Target):
    """
    exists returns true if the table exists.
    """

    def __init__(self, table, database='default', client=None):
        self.database = database
        self.table = table
        self.hive_cmd = load_hive_cmd()
        if client is None:
            client = get_default_client()
        self.client = client

    def exists(self):
        logger.debug("Checking Hive table '%s.%s' exists", self.database, self.table)
        return self.client.table_exists(self.table, self.database)

    @property
    def path(self):
        """
        Returns the path to this table in HDFS.
        """
        location = self.client.table_location(self.table, self.database)
        if not location:
            raise Exception("Couldn't find location for table: {0}".format(str(self)))
        return location

    def open(self, mode):
        return NotImplementedError("open() is not supported for HiveTableTarget")

class HivePartitionTarget(luigi.Target):
    """
    exists returns true if the table's partition exists.
    """

    def __init__(self, table, partition, database='default', fail_missing_table=True, client=None):
        self.database = database
        self.table = table
        self.partition = partition
        if client is None:
            client = get_default_client()
        self.client = client

        self.fail_missing_table = fail_missing_table

    def exists(self):
        try:
            logger.debug("Checking Hive table '{d}.{t}' for partition {p}".format(d=self.database, t=self.table, p=str(self.partition)))
            return self.client.table_exists(self.table, self.database, self.partition)
        except HiveCommandError:
            if self.fail_missing_table:
                raise
            else:
                if self.client.table_exists(self.table, self.database):
                    # a real error occurred
                    raise
                else:
                    # oh the table just doesn't exist
                    return False

    @property
    def path(self):
        """
        Returns the path for this HiveTablePartitionTarget's data.
        """
        location = self.client.table_location(self.table, self.database, self.partition)
        if not location:
            raise Exception("Couldn't find location for table: {0}".format(str(self)))
        return location

    def open(self, mode):
        return NotImplementedError("open() is not supported for HivePartitionTarget")

class ExternalHiveTask(luigi.ExternalTask):
    """
    External task that depends on a Hive table/partition.
    """

    database = luigi.Parameter(default='default')
    table = luigi.Parameter()
    # since this is an external task and will never be initialized from the CLI, partition can be any python object, in this case a dictionary
    partition = luigi.Parameter(default=None, description='Python dictionary specifying the target partition e.g. {"date": "2013-01-25"}')

    def output(self):
        if self.partition is not None:
            assert self.partition, "partition required"
            return HivePartitionTarget(table=self.table,
                                       partition=self.partition,
                                       database=self.database)
        else:
            return HiveTableTarget(self.table, self.database)
<commit_msg>Fix module path to exception<commit_after># -*- coding: utf-8 -*-
#
# Copyright 2012-2015 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import abc
import logging
import operator
import os
import subprocess
import tempfile

from luigi import six

import luigi
import luigi.contrib.hadoop
from luigi.target import FileAlreadyExists, FileSystemTarget
from luigi.task import flatten

if six.PY3:
    unicode = str

logger = logging.getLogger('luigi-interface')

class HiveCommandError(RuntimeError):

    def __init__(self, message, out=None, err=None):
        super(HiveCommandError, self).__init__(message, out, err)
        self.message = message
        self.out = out
        self.err = err

def load_hive_cmd():
    return luigi.configuration.get_config().get('hive', 'command', 'hive')

def get_hive_syntax():
    return luigi.configuration.get_config().get('hive', 'release', 'cdh4')

def run_hive(args, check_return_code=True):
    """
    Runs the `hive` from the command line, passing in the given args, and
    returning stdout.

    With the apache release of Hive, so of the table existence checks
    (which are done using DESCRIBE do not exit with a return code of 0
    so we need an option to ignore the return code and just return stdout for parsing
    """
    cmd = [load_hive_cmd()] + args
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = p.communicate()
    if check_return_code and p.returncode != 0:
        raise HiveCommandError("Hive command: {0} failed with error code: {1}".format(" ".join(cmd), p.returncode),
                               stdout, stderr)
    return stdout

def run_hive_cmd(hivecmd, check_return_code=True):
    """
    Runs the given hive query and returns stdout.
    """
    return run_hive(['-e', hivecmd], check_return_code)

def run_hive_script(script):
    """
    Runs the contents of the given script in hive and returns stdout.
    """
    if not os.path.isfile(script):
        raise RuntimeError("Hive script: {0} does not exist.".format(script))
    return run_hive(['-f', script])

@six.add_metaclass(abc.ABCMeta)
class HiveClient(object):  # interface

    @abc.abstractmethod
    def table_location(self, table, database='default', partition=None):
        """
        Returns location of db.table (or db.table.partition). partition is a dict of partition key to
        value.
        """
        pass

    @abc.abstractmethod
    def table_schema(self, table, database='default'):
        """
        Returns list of [(name, type)] for each column in database.table.
        """
        pass

    @abc.abstractmethod
    def table_exists(self, table, database='default', partition=None):
        """
        Returns true if db.table (or db.table.partition) exists. partition is a dict of partition key to
        value.
        """
        pass

    @abc.abstractmethod
    def partition_spec(self, partition):
        """ Turn a dict into a string partition specification """
        pass

class HiveCommandClient(HiveClient):
    """
    Uses `hive` invocations to find information.
    """

    def table_location(self, table, database='default', partition=None):
        cmd = "use {0}; describe formatted {1}".format(database, table)
        if partition is not None:
            cmd += " PARTITION ({0})".format(self.partition_spec(partition))

        stdout = run_hive_cmd(cmd)

        for line in stdout.split("\n"):
            if "Location:" in line:
                return line.split("\t")[1]

    def table_exists(self, table, database='default', partition=None):
        if partition is None:
            stdout = run_hive_cmd('use {0}; show tables like "{1}";'.format(database, table))

            return stdout and table.lower() in stdout
        else:
            stdout = run_hive_cmd("""use %s; show partitions %s partition
                                (%s)""" % (database, table, self.partition_spec(partition)))

            if stdout:
                return True
            else:
                return False

    def table_schema(self, table, database='default'):
        describe = run_hive_cmd("use {0}; describe {1}".format(database, table))
        if not describe or "does not exist" in describe:
            return None
        return [tuple([x.strip() for x in line.strip().split("\t")]) for line in describe.strip().split("\n")]

    def partition_spec(self, partition):
        """
        Turns a dict into the a Hive partition specification string.
        """
        return ','.join(["`{0}`='{1}'".format(k, v) for (k, v) in
                         sorted(six.iteritems(partition), key=operator.itemgetter(0))])

class ApacheHiveCommandClient(HiveCommandClient):
    """
    A subclass for the HiveCommandClient to (in some cases) ignore the return code from
    the hive command so that we can just parse the output.
    """

    def table_schema(self, table, database='default'):
        describe = run_hive_cmd("use {0}; describe {1}".format(database, table), False)
        if not describe or "Table not found" in describe:
            return None
        return [tuple([x.strip() for x in line.strip().split("\t")]) for line in describe.strip().split("\n")]

class MetastoreClient(HiveClient):

    def table_location(self, table, database='default', partition=None):
        with HiveThriftContext() as client:
            if partition is not None:
                try:
                    partition_str = self.partition_spec(partition)
                    thrift_table = client.get_partition_by_name(database, table, partition_str)
                except hive_metastore.ttypes.NoSuchObjectException as e:
                    return ''
            else:
                thrift_table = client.get_table(database, table)
            return thrift_table.sd.location

    def table_exists(self, table, database='default', partition=None):
        with HiveThriftContext() as client:
            if partition is None:
                return table in client.get_all_tables(database)
            else:
                return partition in self._existing_partitions(table, database, client)

    def _existing_partitions(self, table, database, client):
        def _parse_partition_string(partition_string):
            partition_def = {}
            for part in partition_string.split("/"):
                name, value = part.split("=")
                partition_def[name] = value
            return partition_def

        # -1 is max_parts, the # of partition names to return (-1 = unlimited)
        partition_strings = client.get_partition_names(database, table, -1)
        return [_parse_partition_string(existing_partition) for existing_partition in partition_strings]

    def table_schema(self, table, database='default'):
        with HiveThriftContext() as client:
            return [(field_schema.name, field_schema.type) for field_schema in client.get_schema(database, table)]

    def partition_spec(self, partition):
        return "/".join("%s=%s" % (k, v) for (k, v) in sorted(six.iteritems(partition), key=operator.itemgetter(0)))

class HiveThriftContext(object):
    """
    Context manager for hive metastore client.
    """

    def __enter__(self):
        try:
            from thrift.transport import TSocket
            from thrift.transport import TTransport
            from thrift.protocol import TBinaryProtocol
            # Note that this will only work with a CDH release.
            # This uses the thrift bindings generated by the ThriftHiveMetastore service in Beeswax.
            # If using the Apache release of Hive this import will fail.
            from hive_metastore import ThriftHiveMetastore
            config = luigi.configuration.get_config()
            host = config.get('hive', 'metastore_host')
            port = config.getint('hive', 'metastore_port')
            transport = TSocket.TSocket(host, port)
            transport = TTransport.TBufferedTransport(transport)
            protocol = TBinaryProtocol.TBinaryProtocol(transport)
            transport.open()
            self.transport = transport
            return ThriftHiveMetastore.Client(protocol)
        except ImportError as e:
            raise Exception('Could not import Hive thrift library:' + str(e))

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.transport.close()

def get_default_client():
    syntax = get_hive_syntax()
    if syntax == "apache":
        return ApacheHiveCommandClient()
    elif syntax == "metastore":
        return MetastoreClient()
    else:
        return HiveCommandClient()

client = get_default_client()

class HiveQueryTask(luigi.contrib.hadoop.BaseHadoopJobTask):
    """
    Task to run a hive query.
    """

    # by default, we let hive figure these out.
    n_reduce_tasks = None
    bytes_per_reducer = None
    reducers_max = None

    @abc.abstractmethod
    def query(self):
        """ Text of query to run in hive """
        raise RuntimeError("Must implement query!")

    def hiverc(self):
        """
        Location of an rc file to run before the query
        if hiverc-location key is specified in luigi.cfg, will default to the value there
        otherwise returns None.

        Returning a list of rc files will load all of them in order.
        """
        return luigi.configuration.get_config().get('hive', 'hiverc-location', default=None)

    def hiveconfs(self):
        """
        Returns an dict of key=value settings to be passed along
        to the hive command line via --hiveconf. By default, sets
        mapred.job.name to task_id and if not None, sets:

        * mapred.reduce.tasks (n_reduce_tasks)
        * mapred.fairscheduler.pool (pool) or mapred.job.queue.name (pool)
        * hive.exec.reducers.bytes.per.reducer (bytes_per_reducer)
        * hive.exec.reducers.max (reducers_max)
        """
        jcs = {}
        jcs['mapred.job.name'] = self.task_id
        if self.n_reduce_tasks is not None:
            jcs['mapred.reduce.tasks'] = self.n_reduce_tasks
        if self.pool is not None:
            # Supporting two schedulers: fair (default) and capacity using the same option
            scheduler_type = luigi.configuration.get_config().get('hadoop', 'scheduler', 'fair')
            if scheduler_type == 'fair':
                jcs['mapred.fairscheduler.pool'] = self.pool
            elif scheduler_type == 'capacity':
                jcs['mapred.job.queue.name'] = self.pool
        if self.bytes_per_reducer is not None:
            jcs['hive.exec.reducers.bytes.per.reducer'] = self.bytes_per_reducer
        if self.reducers_max is not None:
            jcs['hive.exec.reducers.max'] = self.reducers_max
        return jcs

    def job_runner(self):
        return HiveQueryRunner()

class HiveQueryRunner(luigi.contrib.hadoop.JobRunner):
    """
    Runs a HiveQueryTask by shelling out to hive.
    """

    def prepare_outputs(self, job):
        """
        Called before job is started.

        If output is a `FileSystemTarget`, create parent directories so the hive command won't fail
        """
        outputs = flatten(job.output())
        for o in outputs:
            if isinstance(o, FileSystemTarget):
                parent_dir = os.path.dirname(o.path)
                if parent_dir and not o.fs.exists(parent_dir):
                    logger.info("Creating parent directory %r", parent_dir)
                    try:
                        # there is a possible race condition
                        # which needs to be handled here
                        o.fs.mkdir(parent_dir)
                    except FileAlreadyExists:
                        pass

    def run_job(self, job, tracking_url_callback=None):
        self.prepare_outputs(job)
        with tempfile.NamedTemporaryFile() as f:
            query = job.query()
            if isinstance(query, unicode):
                query = query.encode('utf8')
            f.write(query)
            f.flush()
            arglist = [load_hive_cmd(), '-f', f.name]
            hiverc = job.hiverc()
            if hiverc:
                if isinstance(hiverc, str):
                    hiverc = [hiverc]
                for rcfile in hiverc:
                    arglist += ['-i', rcfile]
            if job.hiveconfs():
                for k, v in six.iteritems(job.hiveconfs()):
                    arglist += ['--hiveconf', '{0}={1}'.format(k, v)]

            logger.info(arglist)
            return luigi.contrib.hadoop.run_and_track_hadoop_job(arglist, tracking_url_callback)

class HiveTableTarget(luigi.Target):
    """
    exists returns true if the table exists.
    """

    def __init__(self, table, database='default', client=None):
        self.database = database
        self.table = table
        self.hive_cmd = load_hive_cmd()
        if client is None:
            client = get_default_client()
        self.client = client

    def exists(self):
        logger.debug("Checking Hive table '%s.%s' exists", self.database, self.table)
        return self.client.table_exists(self.table, self.database)

    @property
    def path(self):
        """
        Returns the path to this table in HDFS.
        """
        location = self.client.table_location(self.table, self.database)
        if not location:
            raise Exception("Couldn't find location for table: {0}".format(str(self)))
        return location

    def open(self, mode):
        return NotImplementedError("open() is not supported for HiveTableTarget")

class HivePartitionTarget(luigi.Target):
    """
    exists returns true if the table's partition exists.
    """

    def __init__(self, table, partition, database='default', fail_missing_table=True, client=None):
        self.database = database
        self.table = table
        self.partition = partition
        if client is None:
            client = get_default_client()
        self.client = client

        self.fail_missing_table = fail_missing_table

    def exists(self):
        try:
            logger.debug("Checking Hive table '{d}.{t}' for partition {p}".format(d=self.database, t=self.table, p=str(self.partition)))
            return self.client.table_exists(self.table, self.database, self.partition)
        except HiveCommandError:
            if self.fail_missing_table:
                raise
            else:
                if self.client.table_exists(self.table, self.database):
                    # a real error occurred
                    raise
                else:
                    # oh the table just doesn't exist
                    return False

    @property
    def path(self):
        """
        Returns the path for this HiveTablePartitionTarget's data.
        """
        location = self.client.table_location(self.table, self.database, self.partition)
        if not location:
            raise Exception("Couldn't find location for table: {0}".format(str(self)))
        return location

    def open(self, mode):
        return NotImplementedError("open() is not supported for HivePartitionTarget")

class ExternalHiveTask(luigi.ExternalTask):
    """
    External task that depends on a Hive table/partition.
    """

    database = luigi.Parameter(default='default')
    table = luigi.Parameter()
    # since this is an external task and will never be initialized from the CLI, partition can be any python object, in this case a dictionary
    partition = luigi.Parameter(default=None, description='Python dictionary specifying the target partition e.g. {"date": "2013-01-25"}')

    def output(self):
        if self.partition is not None:
            assert self.partition, "partition required"
            return HivePartitionTarget(table=self.table,
                                       partition=self.partition,
                                       database=self.database)
        else:
            return HiveTableTarget(self.table, self.database)
<|endoftext|>


<commit_before>import utilities
import subprocess
import datetime

def create_gain_year_count(tile_id):

    print "Processing:", tile_id

    # start time
    start = datetime.datetime.now()

    # Names of the loss, gain and tree cover density tiles
    loss = '{}.tif'.format(tile_id)
    gain = 'Hansen_GFC2015_gain_{}.tif'.format(tile_id)
    tcd = 'Hansen_GFC2014_treecover2000_{}.tif'.format(tile_id)

    # Location to upload files to
    upload_dir = 's3://gfw2-data/climate/carbon_model/forest_age/20180830/'

    print 'Loss tile is', loss
    print 'Gain tile is', gain
    print 'tcd tile is', tcd

    # # Creates four separate rasters for the four tree cover loss/gain combinations for pixels. Then merges the rasters.
    # # In all rasters, 0 is NoData value.
    # # Pixels with loss only
    # print "Creating raster of growth years for loss-only pixels"
    # #gdal_calc.py -A 00N_050W.tif -B Hansen_GFC2015_gain_00N_050W.tif --calc="(A>0)*(B==0)*(A-1)" --outfile=loss_only.tif --NoDataValue=0 --overwrite
    # loss_calc = '--calc=(A>0)*(B==0)*(A-1)'
    # loss_outfile1 = 'growth_years_loss_only_{}.tif'.format(tile_id)
    # loss_outfile2 = '--outfile={}'.format(loss_outfile1)
    # cmd = ['gdal_calc.py', '-A', loss, '-B', gain, loss_calc, loss_outfile2, '--NoDataValue=0', '--overwrite', '--co', 'COMPRESS=LZW']
    # subprocess.check_call(cmd)
    #
    # # Pixels with gain only
    # print "Creating raster of growth years for gain-only pixels"
    # #gdal_calc.py -A 00N_050W.tif -B Hansen_GFC2015_gain_00N_050W.tif --calc="(A==0)*(B==1)*6" --outfile=gain_only.tif --NoDataValue=0 --overwrite
    # gain_calc = '--calc=(A==0)*(B==1)*6'
    # gain_outfile1 = 'growth_years_gain_only_{}.tif'.format(tile_id)
    # gain_outfile2 = '--outfile={}'.format(gain_outfile1)
    # cmd = ['gdal_calc.py', '-A', loss, '-B', gain, gain_calc, gain_outfile2, '--NoDataValue=0', '--overwrite', '--co', 'COMPRESS=LZW']
    # subprocess.check_call(cmd)
    #
    # # Pixels with neither loss nor gain but in areas with tree cover density >0
    # print "Creating raster of growth years for no change pixels"
    # # gdal_calc.py -A 00N_050W.tif -B Hansen_GFC2015_gain_00N_050W.tif -C Hansen_GFC2014_treecover2000_00N_050W.tif --calc "(A==0)*(B==0)*(C>0)*15" --outfile=no_change.tif --NoDataValue=0 --overwrite
    # no_change_calc = '--calc=(A==0)*(B==0)*(C>0)*15'
    # no_change_outfile1 = 'growth_years_no_change_{}.tif'.format(tile_id)
    # no_change_outfile2 = '--outfile={}'.format(no_change_outfile1)
    # cmd = ['gdal_calc.py', '-A', loss, '-B', gain, '-C', tcd,  no_change_calc, no_change_outfile2, '--NoDataValue=0', '--overwrite', '--co', 'COMPRESS=LZW']
    # subprocess.check_call(cmd)
    #
    # # Pixels with both loss and gain
    # print "Creating raster of growth years for loss and gain pixels"
    # # gdal_calc.py -A 00N_050W.tif -B Hansen_GFC2015_gain_00N_050W.tif --calc="((A>0)*(B==1)*((A-1)+(16-A)/2))" --outfile=gain_and_loss.tif --NoDataValue=0 --overwrite
    # loss_and_gain_calc = '--calc=((A>0)*(B==1)*((A-1)+(16-A)/2))'
    # loss_and_gain_outfile1 = 'growth_years_loss_and_gain_{}.tif'.format(tile_id)
    # loss_and_gain_outfile2 = '--outfile={}'.format(loss_and_gain_outfile1)
    # cmd = ['gdal_calc.py', '-A', loss, '-B', gain, loss_and_gain_calc, loss_and_gain_outfile2, '--NoDataValue=0', '--overwrite', '--co', 'COMPRESS=LZW']
    # subprocess.check_call(cmd)

    print "Merging loss, gain, no change, and loss/gain pixels into single raster"
    age_outfile = 'growth_years_total_{}.tif'.format(tile_id)
    cmd = ['gdal_merge.py', '-o', age_outfile, loss_outfile1, gain_outfile1, no_change_outfile1, loss_and_gain_outfile1, '-co', 'COMPRESS=LZW', '-a_nodata', '0']
    subprocess.check_call(cmd)

    utilities.upload_final(upload_dir, tile_id)

    end = datetime.datetime.now()
    elapsed_time = end-start

    print "Processing time for tile", tile_id, ":", elapsed_time

<commit_msg>Gain year calculations work on single tile. Now trying to multiprocess. The issue with gdal_calc in multiprocessing was that the --calc and --outfile arguments need to be supplied as variables, not as items in the cmd list. Trying to get that to work for all four gdal_calc calls now.<commit_after>import utilities
import subprocess
import datetime

def create_gain_year_count(tile_id):

    print "Processing:", tile_id

    # start time
    start = datetime.datetime.now()

    # Names of the loss, gain and tree cover density tiles
    loss = '{}.tif'.format(tile_id)
    gain = 'Hansen_GFC2015_gain_{}.tif'.format(tile_id)
    tcd = 'Hansen_GFC2014_treecover2000_{}.tif'.format(tile_id)

    # Location to upload files to
    upload_dir = 's3://gfw2-data/climate/carbon_model/forest_age/20180830/'

    print 'Loss tile is', loss
    print 'Gain tile is', gain
    print 'tcd tile is', tcd

    # # Creates four separate rasters for the four tree cover loss/gain combinations for pixels. Then merges the rasters.
    # # In all rasters, 0 is NoData value.
    # # Pixels with loss only
    # print "Creating raster of growth years for loss-only pixels"
    # #gdal_calc.py -A 00N_050W.tif -B Hansen_GFC2015_gain_00N_050W.tif --calc="(A>0)*(B==0)*(A-1)" --outfile=loss_only.tif --NoDataValue=0 --overwrite
    # loss_calc = '--calc=(A>0)*(B==0)*(A-1)'
    loss_outfile1 = 'growth_years_loss_only_{}.tif'.format(tile_id)
    # loss_outfile2 = '--outfile={}'.format(loss_outfile1)
    # cmd = ['gdal_calc.py', '-A', loss, '-B', gain, loss_calc, loss_outfile2, '--NoDataValue=0', '--overwrite', '--co', 'COMPRESS=LZW']
    # subprocess.check_call(cmd)
    #
    # # Pixels with gain only
    # print "Creating raster of growth years for gain-only pixels"
    # #gdal_calc.py -A 00N_050W.tif -B Hansen_GFC2015_gain_00N_050W.tif --calc="(A==0)*(B==1)*6" --outfile=gain_only.tif --NoDataValue=0 --overwrite
    # gain_calc = '--calc=(A==0)*(B==1)*6'
    gain_outfile1 = 'growth_years_gain_only_{}.tif'.format(tile_id)
    # gain_outfile2 = '--outfile={}'.format(gain_outfile1)
    # cmd = ['gdal_calc.py', '-A', loss, '-B', gain, gain_calc, gain_outfile2, '--NoDataValue=0', '--overwrite', '--co', 'COMPRESS=LZW']
    # subprocess.check_call(cmd)
    #
    # # Pixels with neither loss nor gain but in areas with tree cover density >0
    # print "Creating raster of growth years for no change pixels"
    # # gdal_calc.py -A 00N_050W.tif -B Hansen_GFC2015_gain_00N_050W.tif -C Hansen_GFC2014_treecover2000_00N_050W.tif --calc "(A==0)*(B==0)*(C>0)*15" --outfile=no_change.tif --NoDataValue=0 --overwrite
    # no_change_calc = '--calc=(A==0)*(B==0)*(C>0)*15'
    no_change_outfile1 = 'growth_years_no_change_{}.tif'.format(tile_id)
    # no_change_outfile2 = '--outfile={}'.format(no_change_outfile1)
    # cmd = ['gdal_calc.py', '-A', loss, '-B', gain, '-C', tcd,  no_change_calc, no_change_outfile2, '--NoDataValue=0', '--overwrite', '--co', 'COMPRESS=LZW']
    # subprocess.check_call(cmd)
    #
    # # Pixels with both loss and gain
    # print "Creating raster of growth years for loss and gain pixels"
    # # gdal_calc.py -A 00N_050W.tif -B Hansen_GFC2015_gain_00N_050W.tif --calc="((A>0)*(B==1)*((A-1)+(16-A)/2))" --outfile=gain_and_loss.tif --NoDataValue=0 --overwrite
    # loss_and_gain_calc = '--calc=((A>0)*(B==1)*((A-1)+(16-A)/2))'
    loss_and_gain_outfile1 = 'growth_years_loss_and_gain_{}.tif'.format(tile_id)
    # loss_and_gain_outfile2 = '--outfile={}'.format(loss_and_gain_outfile1)
    # cmd = ['gdal_calc.py', '-A', loss, '-B', gain, loss_and_gain_calc, loss_and_gain_outfile2, '--NoDataValue=0', '--overwrite', '--co', 'COMPRESS=LZW']
    # subprocess.check_call(cmd)

    print "Merging loss, gain, no change, and loss/gain pixels into single raster"
    age_outfile = 'growth_years_total_{}.tif'.format(tile_id)
    cmd = ['gdal_merge.py', '-o', age_outfile, loss_outfile1, gain_outfile1, no_change_outfile1, loss_and_gain_outfile1, '-co', 'COMPRESS=LZW', '-a_nodata', '0']
    subprocess.check_call(cmd)

    utilities.upload_final(upload_dir, tile_id)

    end = datetime.datetime.now()
    elapsed_time = end-start

    print "Processing time for tile", tile_id, ":", elapsed_time

<|endoftext|>


<commit_before>import datetime

from django.conf import settings

from django_conneg.views import HTMLView

from humfrey.utils.views import CachedView

# Only create FeedView class if feedparser and pytz are importable.
# This will make feedparser and pytz optional dependencies if one
# doesn't want to use FeedView.
try:
    import feedparser
    import pytz
except ImportError, e:
    pass
else:
    class FeedView(HTMLView, CachedView):
        template_name = 'index'
        rss_url = None

        def initial_context(self, request):
            try:
                feed = feedparser.parse(self.rss_url)
                for entry in feed.entries:
                    entry.updated_datetime = datetime.datetime(*entry.updated_parsed[:6] + (pytz.utc,)) \
                                                 .astimezone(pytz.timezone(settings.TIME_ZONE))
            except Exception, e:
                feed = None
            return {
                'feed': feed,
            }

        def get(self, request):
            context = self.initial_context(request)
            return self.render(request, context, self.template_name)

class SimpleView(HTMLView, CachedView):
    context = {}
    template_name = None

    def get(self, request):
        return self.render(request, self.context.copy(), self.template_name)
<commit_msg>Added a CannedQueryView to humfrey.misc.views.<commit_after>import datetime

import rdflib
from django.conf import settings

from django_conneg.views import HTMLView, ContentNegotiatedView

from humfrey.utils.views import CachedView, EndpointView

# Only create FeedView class if feedparser and pytz are importable.
# This will make feedparser and pytz optional dependencies if one
# doesn't want to use FeedView.
try:
    import feedparser
    import pytz
except ImportError, e:
    pass
else:
    class FeedView(HTMLView, CachedView):
        template_name = 'index'
        rss_url = None

        def initial_context(self, request):
            try:
                feed = feedparser.parse(self.rss_url)
                for entry in feed.entries:
                    entry.updated_datetime = datetime.datetime(*entry.updated_parsed[:6] + (pytz.utc,)) \
                                                 .astimezone(pytz.timezone(settings.TIME_ZONE))
            except Exception, e:
                feed = None
            return {
                'feed': feed,
            }

        def get(self, request):
            context = self.initial_context(request)
            return self.render(request, context, self.template_name)

class SimpleView(HTMLView, CachedView):
    context = {}
    template_name = None

    def get(self, request):
        return self.render(request, self.context.copy(), self.template_name)

class CannedQueryView(CachedView, EndpointView):
    query = None
    template_name = None

    def get_query(self, request, *args, **kwargs):
        return self.query

    def get_locations(self, request, *args, **kwargs):
        """
        Override to provide the canonical and format-specific locations for this resource.

        Example return value: ('http://example.com/data', 'http://example.org/data.rss')
        """
        return None, None

    def get(self, request, *args, **kwargs):
        self.base_location, self.content_location = self.get_locations(request, *args, **kwargs)
        query = self.get_query(request, *args, **kwargs)
        result = self.endpoint.query(query)
        if isinstance(result, list):
            context = {'results': result}
        elif isinstance(result, bool):
            context = {'result': result}
        elif isinstance(result, rdflib.ConjunctiveGraph):
            context = {'graph': result}

        if self.content_location:
            context['additional_headers'] = {'Content-location': self.content_location}

        if 'format' in kwargs:
            self.render_to_format(request, context, self.template_name, kwargs['format'])
        else:
            return self.render(request, context, self.template_name)
<|endoftext|>


<commit_before>__author__ = "Christian Kongsgaard"

# -------------------------------------------------------------------------------------------------------------------- #
# IMPORTS

# Modules:
import os

# RiBuild Modules:
import delphin_6_automation.database_interactions.db_templates.delphin_entry as delphin_db
import delphin_6_automation.database_interactions.db_templates.material_entry as material_db
from delphin_6_automation.file_parsing import material_parser

# -------------------------------------------------------------------------------------------------------------------- #
# DATABASE INTERACTIONS

def find_material_ids(project_materials: list) -> list:
    """
    Find ids of given material entries based on material name and material unique id.

    :param project_materials: List tuples with material file names and unique material ids
    :type project_materials: list
    :return: list with material entries
    :rtype: list
    """

    material_entries = []
    for material_pair in project_materials:
        material_entries.append(material_db.Material.objects(material_name=material_pair[0],
                                                             material_id=material_pair[1]).first())
    if material_entries:
        return material_entries
    else:
        raise FileNotFoundError('Materials does not exist in database. '
                                'Please upload material files before uploading Delphin Projects')

def list_project_materials(delphin_document: delphin_db.Delphin) -> list:
    """
    Returns a list with the materials in a project entry.

    :param delphin_document: Delphin entry
    :return: List tuples with material file names and unique material ids
    """

    materials = dict(delphin_document.dp6_file)['DelphinProject']['Materials']['MaterialReference']

    material_list = [(material['#text'].split('/')[-1].split('_')[0],
                      int(material['#text'].split('/')[-1].split('_')[-1][:-3]))
                     for material in materials]

    return material_list

def convert_and_upload_file(user_path_input):

    material_dict_lst = []

    if user_path_input.endswith(".m6"):
        upload_material_file(user_path_input)

    else:
        for root, dirs, files in os.walk(user_path_input):
            for file, _ in zip(files, root):
                if file.endswith(".m6"):
                    upload_material_file(file)

    return material_dict_lst

def upload_material_file(material_path: str) -> delphin_db.Delphin.id:
    """
    Uploads a Delphin file to a database.rst.

    :param material_path: Path to a Delphin 6 material project file
    :return: Database entry id
    """

    entry = material_db.Material()
    entry.material_data = material_parser.material_file_to_dict(material_path)
    entry.material_name = os.path.split(material_path)[-1].split('_')[0]
    entry.material_id = int(os.path.split(material_path)[-1].split('_')[1][:-3])
    entry.save()

    return entry.id

def change_material_location(delphin_id: str, path: str) -> str:
    """
    Changes the location of the material database location for the Delphin Project file.

    :param delphin_id: ID of entry
    :type delphin_id: str
    :param path: Path to change it to
    :type path: str
    :return: ID of entry
    :rtype: str
    """

    delphin_document = delphin_db.Delphin.objects(id=delphin_id).first()
    delphin_dict = dict(delphin_document.dp6_file)
    delphin_dict['DelphinProject']['DirectoryPlaceholders']['Placeholder'][1]['#text'] = path
    delphin_document.update(set__dp6_file=delphin_dict)

    return delphin_document.id

def download_materials(delphin_id: str, path: str) -> bool:
    """
    Downloads the materials of a Delphin Project

    :param delphin_id: Delphin entry ID
    :type delphin_id: str
    :param path: Path to save to
    :type path: str
    :return: True
    :rtype: bool
    """

    materials_list = delphin_db.Delphin.objects(id=delphin_id).first().materials
    change_material_location(delphin_id, path)

    if not os.path.isdir(path):
        os.mkdir(path)

    for material in materials_list:
        material_parser.dict_to_m6(material, path)

    return True
<commit_msg>catched if material was missing from database<commit_after>__author__ = "Christian Kongsgaard"

# -------------------------------------------------------------------------------------------------------------------- #
# IMPORTS

# Modules:
import os

# RiBuild Modules:
import delphin_6_automation.database_interactions.db_templates.delphin_entry as delphin_db
import delphin_6_automation.database_interactions.db_templates.material_entry as material_db
from delphin_6_automation.file_parsing import material_parser

# -------------------------------------------------------------------------------------------------------------------- #
# DATABASE INTERACTIONS

def find_material_ids(project_materials: list) -> list:
    """
    Find ids of given material entries based on material name and material unique id.

    :param project_materials: List tuples with material file names and unique material ids
    :type project_materials: list
    :return: list with material entries
    :rtype: list
    """

    material_entries = []
    for material_pair in project_materials:
        material_entry = material_db.Material.objects(material_name=material_pair[0],
                                                      material_id=material_pair[1]).first()

        if material_entry:
            material_entries.append(material_entry)
        else:
            raise FileNotFoundError(f'Material: {material_pair[0]} with material ID: {material_pair[1]} '
                                    f'does not exist in database.\n'
                                    f'Please upload material files before uploading Delphin Projects\n')

    return material_entries

def list_project_materials(delphin_document: delphin_db.Delphin) -> list:
    """
    Returns a list with the materials in a project entry.

    :param delphin_document: Delphin entry
    :return: List tuples with material file names and unique material ids
    """

    materials = dict(delphin_document.dp6_file)['DelphinProject']['Materials']['MaterialReference']

    material_list = [(material['#text'].split('/')[-1].split('_')[0],
                      int(material['#text'].split('/')[-1].split('_')[-1][:-3]))
                     for material in materials]

    return material_list

def convert_and_upload_file(user_path_input):

    material_dict_lst = []

    if user_path_input.endswith(".m6"):
        upload_material_file(user_path_input)

    else:
        for root, dirs, files in os.walk(user_path_input):
            for file, _ in zip(files, root):
                if file.endswith(".m6"):
                    upload_material_file(file)

    return material_dict_lst

def upload_material_file(material_path: str) -> delphin_db.Delphin.id:
    """
    Uploads a Delphin file to a database.rst.

    :param material_path: Path to a Delphin 6 material project file
    :return: Database entry id
    """

    entry = material_db.Material()
    entry.material_data = material_parser.material_file_to_dict(material_path)
    entry.material_name = os.path.split(material_path)[-1].split('_')[0]
    entry.material_id = int(os.path.split(material_path)[-1].split('_')[1][:-3])
    entry.save()

    return entry.id

def change_material_location(delphin_id: str, path: str) -> str:
    """
    Changes the location of the material database location for the Delphin Project file.

    :param delphin_id: ID of entry
    :type delphin_id: str
    :param path: Path to change it to
    :type path: str
    :return: ID of entry
    :rtype: str
    """

    delphin_document = delphin_db.Delphin.objects(id=delphin_id).first()
    delphin_dict = dict(delphin_document.dp6_file)
    delphin_dict['DelphinProject']['DirectoryPlaceholders']['Placeholder'][1]['#text'] = path
    delphin_document.update(set__dp6_file=delphin_dict)

    return delphin_document.id

def download_materials(delphin_id: str, path: str) -> bool:
    """
    Downloads the materials of a Delphin Project

    :param delphin_id: Delphin entry ID
    :type delphin_id: str
    :param path: Path to save to
    :type path: str
    :return: True
    :rtype: bool
    """

    materials_list = delphin_db.Delphin.objects(id=delphin_id).first().materials
    change_material_location(delphin_id, path)

    if not os.path.isdir(path):
        os.mkdir(path)

    for material in materials_list:
        material_parser.dict_to_m6(material, path)

    return True
<|endoftext|>


<commit_before>
import time
import cStringIO
import struct

from twisted.internet import reactor

from logs import lg

from lib import udp

#------------------------------------------------------------------------------

"""
Datagram Format

[Data] packet

bytes:
  0        software version number
  1        command identifier, see ``lib.udp`` module
  2-5      stream_id
  6-9      total data size to be transferred, peer must know when to stop receiving
  10-13    block_id, outgoing blocks are counted from 1
  from 14  payload data

"""

UDP_DATAGRAM_SIZE = 508
BLOCK_SIZE = UDP_DATAGRAM_SIZE - 14

BLOCKS_PER_ACK = 8

MAX_BUFFER_SIZE = 64*1024
BUFFER_SIZE = BLOCK_SIZE * BLOCKS_PER_ACK # BLOCK_SIZE * int(float(BLOCKS_PER_ACK)*0.8) - 20% extra space in ack packet

RTT_MIN_LIMIT = 0.0001
RTT_MAX_LIMIT = 0.5

#------------------------------------------------------------------------------

class UDPStream():
    def __init__(self, stream_id, consumer, producer):
        self.stream_id = stream_id
        self.consumer = consumer
        self.producer = producer
        self.consumer.stream = self
        # self.send_data_packet_func = send_data_packet_func
        # self.send_ack_packet_func = send_ack_packet_func
        # self.received_raw_data_callback = received_raw_data_callback
        # self.sent_raw_data_callback = sent_raw_data_callback
        self.output_buffer_size = 0
        self.output_blocks = {}
        self.output_blocks_acks = []
        self.output_block_id = 0
        self.input_blocks = {}
        self.input_block_id = 0
        self.blocks_to_ack = set()
        self.bytes_in = 0
        self.bytes_sent = 0
        self.bytes_acked = 0
        self.resend_bytes = 0
        self.last_ack_moment = 0
        self.rtt_avarage = RTT_MIN_LIMIT
        self.rtt_acks_counter = 1.0
        self.resend_task = None
        self.resend_inactivity_counter = 0
        self.resend_counter = 0
        self.limit_send_bytes_per_sec = 1 * 125000 # 1 Mbps limit ~ 122 KB/s
        self.creation_time = time.time()
        lg.out(18, 'udp_stream.__init__ %d' % self.stream_id)

    def __del__(self):
        lg.out(18, 'udp_stream.__del__ %d' % self.stream_id)

    def close(self):
        print 'udp_stream.close %d in:%d out:%d acked:%d resend:%d' % (
            self.stream_id, self.bytes_in, self.bytes_sent, self.bytes_acked, self.resend_bytes)
        self.stop_resending()
        self.consumer.stream = None
        self.consumer = None
        self.producer = None
        self.blocks_to_ack.clear()
        self.output_blocks.clear()
        # self.send_data_packet_func = None
        # self.send_ack_packet_func = None
        # self.received_raw_data_callback = None
        # self.sent_raw_data_callback = None

    def block_received(self, inpt):
        if self.consumer:
            block_id = struct.unpack('i', inpt.read(4))[0]
            data = inpt.read()
            self.input_blocks[block_id] = data
            self.bytes_in += len(data)
            self.blocks_to_ack.add(block_id)
            eof_state = False
            # print 'block', block_id, self.bytes_in, block_id % BLOCKS_PER_ACK
            if block_id == self.input_block_id + 1:
                newdata = []
                # print 'newdata',
                while True:
                    next_block_id = self.input_block_id + 1
                    try:
                        newdata.append(self.input_blocks.pop(next_block_id))
                    except:
                        break
                    self.input_block_id = next_block_id
                    # print next_block_id,
                num_blocks = len(newdata)
                newdata = ''.join(newdata)
                eof_state = self.consumer.on_received_raw_data(newdata)
                # print 'received %d bytes in %d blocks, eof=%r' % (len(newdata), num_blocks, eof_state)
            # want to send the first ack asap
            if time.time() - self.last_ack_moment > RTT_MAX_LIMIT \
                or block_id % BLOCKS_PER_ACK == 1 \
                or eof_state:
                    # self.send_ack()
                    self.resend()

    def ack_received(self, inpt):
        if self.consumer:
            acks = []
            eof = False
            raw_bytes = ''
            while True:
                raw_bytes = inpt.read(4)
                if not raw_bytes:
                    break
                block_id = struct.unpack('i', raw_bytes)[0]
                acks.append(block_id)
                try:
                    outblock = self.output_blocks.pop(block_id)
                except KeyError:
                    # lg.out(10, 'udp_stream.ack_received WARNING block %d not found' % (block_id))
                    continue
                block_size = len(outblock[0])
                self.output_buffer_size -= block_size
                self.bytes_acked += block_size
                relative_time = time.time() - self.creation_time
                last_ack_rtt = relative_time - outblock[1]
                self.rtt_avarage += last_ack_rtt
                self.rtt_acks_counter += 1.0
                if self.rtt_acks_counter > 1000:
                    self.rtt_acks_counter = 5.0
                    self.rtt_avarage = (self.rtt_avarage / 1000.0) * 5.0
                self.consumer.on_sent_raw_data(block_size)
                eof = self.consumer and self.consumer.size == self.bytes_acked
            if len(acks) > 0:
                print 'ack', len(acks), 'blocks,     more:', len(self.output_blocks.keys()),
                print 'rtt:', self.rtt_avarage/self.rtt_acks_counter, 'eof:', eof,
                print 'acked:', self.bytes_acked, 'sent:', self.bytes_sent, 'resend:', self.resend_bytes
                self.resend()
                return
            print 'STOP IT NOW!!!!, ZERO ACK!!!! SEEMS FINE.!!!'
            self.stop_resending()
            sum_not_acked_blocks = sum(map(lambda block: len(block[0]),
                                           self.output_blocks.values()))
            self.output_blocks.clear()
            self.output_buffer_size = 0
            self.bytes_acked += sum_not_acked_blocks
            relative_time = time.time() - self.creation_time
            self.consumer.on_zero_ack(sum_not_acked_blocks)
            return

    def write(self, data):
        if self.output_buffer_size + len(data) > MAX_BUFFER_SIZE:
            raise BufferOverflow('buffer size is %d' % self.output_buffer_size)
        # print 'write', len(data)
        outp = cStringIO.StringIO(data)
        while True:
            piece = outp.read(BLOCK_SIZE)
            if not piece:
                break
            self.output_block_id += 1
            self.output_blocks[self.output_block_id] = (piece, -1)
            self.output_buffer_size += len(piece)
        outp.close()
        self.resend()

    def send_blocks(self):
        if self.consumer:
            relative_time = time.time() - self.creation_time
            current_rate = self.limit_send_bytes_per_sec
            rtt_current = self.rtt_avarage / self.rtt_acks_counter
            if relative_time > 0.0:
                current_rate = self.bytes_sent / relative_time
            if current_rate > self.limit_send_bytes_per_sec:
                return
            resend_time_limit = 8 * BLOCKS_PER_ACK * rtt_current
            new_blocks_counter = 0
            for block_id in self.output_blocks.keys():
                piece, time_sent = self.output_blocks[block_id]
                data_size = len(piece)
                if time_sent >= 0:
                    dt = relative_time - time_sent
                    if dt > resend_time_limit:
                        self.resend_bytes += data_size
                        # print 're -'s,
                    else:
                        # print 'skip', block_id, dt, self.last_ack_rtt
                        continue
                else:
                    pass
                    # print 'go', block_id
                time_sent = relative_time
                self.output_blocks[block_id] = (piece, time_sent)
                output = ''.join((
                    struct.pack('i', block_id),
                    piece))
                # DEBUG
                # import random
                # if random.randint(0, 9) >= 1:
                    # 10 % percent lost
                    # self.producer.do_send_data(self.stream_id, self.consumer, output)
                self.producer.do_send_data(self.stream_id, self.consumer, output)
                self.bytes_sent += data_size
                new_blocks_counter += 1
            if new_blocks_counter > 0:
                print 'send blocks', new_blocks_counter, len(self.output_blocks.keys()),
                print self.bytes_sent, self.bytes_acked, self.resend_bytes, rtt_current, resend_time_limit

    def send_ack(self):
        if self.consumer:
            ack_data = ''.join(map(lambda bid: struct.pack('i', bid), self.blocks_to_ack))
            self.producer.do_send_ack(self.stream_id, self.consumer, ack_data)
            self.output_blocks_acks += list(self.blocks_to_ack)
            print 'send ack', len(self.output_blocks_acks), len(self.blocks_to_ack), self.bytes_in
            self.blocks_to_ack.clear()
            self.last_ack_moment = time.time()

    def resend(self):
        self.resend_counter += 1
        activitiy = len(self.output_blocks.keys()) # + len(self.blocks_to_ack)
        if activitiy > 0:
            # print 'resend out:%s acks:%s' % (len(self.output_blocks.keys()), len(self.blocks_to_ack))
            self.resend_inactivity_counter = 0
        else:
            self.resend_inactivity_counter += 1
        rtt_current = self.rtt_avarage / self.rtt_acks_counter
        next_resend = max(min(rtt_current, RTT_MAX_LIMIT), RTT_MIN_LIMIT)
        if len(self.blocks_to_ack) > 0:
            if time.time() - self.last_ack_moment > next_resend:
                self.send_ack()
        if len(self.output_blocks):
            self.send_blocks()
#        if self.resend_inactivity_counter > 5:
#            # if self.resend_inactivity_counter % 10 == 1:
#            #     print 'drop resend out:%s acks:%s' % (
#            #         len(self.output_blocks.keys()), len(self.blocks_to_ack))
#            next_resend = RTT_MAX_LIMIT * 4.0
#        if self.resend_inactivity_counter > 50:
#            next_resend = RTT_MAX_LIMIT * 16.0
        if self.resend_counter % 500 == 1:
            print 'resend out:%d acks:%d' % (len(self.output_blocks.keys()), len(self.blocks_to_ack)),
            print 'rtt=%r, next=%r, iterations=%d' % (rtt_current, next_resend, self.resend_counter)
        next_resend *= self.resend_inactivity_counter
        if self.resend_task is None:
            self.resend_task = reactor.callLater(next_resend, self.resend)
            return
        if self.resend_task.called:
            self.resend_task = reactor.callLater(next_resend, self.resend)
            return
        if self.resend_task.cancelled:
            self.resend_task = None
            return

    def stop_resending(self):
        if self.resend_task:
            if self.resend_task.active():
                self.resend_task.cancel()
                self.resend_task = None

#------------------------------------------------------------------------------

class BufferOverflow(Exception):
    pass

<commit_msg>28.07.2014 at 12:27:27<commit_after>
import time
import cStringIO
import struct

from twisted.internet import reactor

from logs import lg

from lib import udp

#------------------------------------------------------------------------------

"""
Datagram Format

[Data] packet

bytes:
  0        software version number
  1        command identifier, see ``lib.udp`` module
  2-5      stream_id
  6-9      total data size to be transferred, peer must know when to stop receiving
  10-13    block_id, outgoing blocks are counted from 1
  from 14  payload data

"""

UDP_DATAGRAM_SIZE = 508
BLOCK_SIZE = UDP_DATAGRAM_SIZE - 14

BLOCKS_PER_ACK = 8

MAX_BUFFER_SIZE = 64*1024
BUFFER_SIZE = BLOCK_SIZE * BLOCKS_PER_ACK # BLOCK_SIZE * int(float(BLOCKS_PER_ACK)*0.8) - 20% extra space in ack packet

RTT_MIN_LIMIT = 0.0001
RTT_MAX_LIMIT = 0.5

#------------------------------------------------------------------------------

class UDPStream():
    def __init__(self, stream_id, consumer, producer):
        self.stream_id = stream_id
        self.consumer = consumer
        self.producer = producer
        self.consumer.stream = self
        # self.send_data_packet_func = send_data_packet_func
        # self.send_ack_packet_func = send_ack_packet_func
        # self.received_raw_data_callback = received_raw_data_callback
        # self.sent_raw_data_callback = sent_raw_data_callback
        self.output_buffer_size = 0
        self.output_blocks = {}
        self.output_blocks_acks = []
        self.output_block_id = 0
        self.input_blocks = {}
        self.input_block_id = 0
        self.blocks_to_ack = set()
        self.bytes_in = 0
        self.bytes_sent = 0
        self.bytes_acked = 0
        self.resend_bytes = 0
        self.last_ack_moment = 0
        self.rtt_avarage = RTT_MIN_LIMIT
        self.rtt_acks_counter = 1.0
        self.resend_task = None
        self.resend_inactivity_counter = 0
        self.resend_counter = 0
        self.limit_send_bytes_per_sec = 1 * 125000 # 1 Mbps limit ~ 122 KB/s
        self.creation_time = time.time()
        lg.out(18, 'udp_stream.__init__ %d' % self.stream_id)

    def __del__(self):
        lg.out(18, 'udp_stream.__del__ %d' % self.stream_id)

    def close(self):
        print 'udp_stream.close %d in:%d out:%d acked:%d resend:%d' % (
            self.stream_id, self.bytes_in, self.bytes_sent, self.bytes_acked, self.resend_bytes)
        self.stop_resending()
        self.consumer.stream = None
        self.consumer = None
        self.producer = None
        self.blocks_to_ack.clear()
        self.output_blocks.clear()
        # self.send_data_packet_func = None
        # self.send_ack_packet_func = None
        # self.received_raw_data_callback = None
        # self.sent_raw_data_callback = None

    def block_received(self, inpt):
        if self.consumer:
            block_id = struct.unpack('i', inpt.read(4))[0]
            data = inpt.read()
            self.input_blocks[block_id] = data
            self.bytes_in += len(data)
            self.blocks_to_ack.add(block_id)
            eof_state = False
            # print 'block', block_id, self.bytes_in, block_id % BLOCKS_PER_ACK
            if block_id == self.input_block_id + 1:
                newdata = []
                # print 'newdata',
                while True:
                    next_block_id = self.input_block_id + 1
                    try:
                        newdata.append(self.input_blocks.pop(next_block_id))
                    except:
                        break
                    self.input_block_id = next_block_id
                    # print next_block_id,
                num_blocks = len(newdata)
                newdata = ''.join(newdata)
                eof_state = self.consumer.on_received_raw_data(newdata)
                # print 'received %d bytes in %d blocks, eof=%r' % (len(newdata), num_blocks, eof_state)
            # want to send the first ack asap
            if time.time() - self.last_ack_moment > RTT_MAX_LIMIT \
                or block_id % BLOCKS_PER_ACK == 1 \
                or eof_state:
                    # self.send_ack()
                    self.resend()

    def ack_received(self, inpt):
        if self.consumer:
            acks = []
            eof = False
            raw_bytes = ''
            while True:
                raw_bytes = inpt.read(4)
                if not raw_bytes:
                    break
                block_id = struct.unpack('i', raw_bytes)[0]
                acks.append(block_id)
                try:
                    outblock = self.output_blocks.pop(block_id)
                except KeyError:
                    # lg.out(10, 'udp_stream.ack_received WARNING block %d not found' % (block_id))
                    continue
                block_size = len(outblock[0])
                self.output_buffer_size -= block_size
                self.bytes_acked += block_size
                relative_time = time.time() - self.creation_time
                last_ack_rtt = relative_time - outblock[1]
                self.rtt_avarage += last_ack_rtt
                self.rtt_acks_counter += 1.0
                if self.rtt_acks_counter > 1000:
                    self.rtt_acks_counter = 5.0
                    self.rtt_avarage = (self.rtt_avarage / 1000.0) * 5.0
                self.consumer.on_sent_raw_data(block_size)
                eof = self.consumer and self.consumer.size == self.bytes_acked
            if len(acks) > 0:
                print 'ack', len(acks), 'blocks,     more:', len(self.output_blocks.keys()),
                print 'rtt:', self.rtt_avarage/self.rtt_acks_counter, 'eof:', eof,
                print 'acked:', self.bytes_acked, 'sent:', self.bytes_sent, 'resend:', self.resend_bytes
                self.resend()
                return
            print 'STOP IT NOW!!!!, ZERO ACK!!!! SEEMS FINE.!!!'
            self.stop_resending()
            sum_not_acked_blocks = sum(map(lambda block: len(block[0]),
                                           self.output_blocks.values()))
            self.output_blocks.clear()
            self.output_buffer_size = 0
            self.bytes_acked += sum_not_acked_blocks
            relative_time = time.time() - self.creation_time
            self.consumer.on_zero_ack(sum_not_acked_blocks)
            return

    def write(self, data):
        if self.output_buffer_size + len(data) > MAX_BUFFER_SIZE:
            raise BufferOverflow('buffer size is %d' % self.output_buffer_size)
        # print 'write', len(data)
        outp = cStringIO.StringIO(data)
        while True:
            piece = outp.read(BLOCK_SIZE)
            if not piece:
                break
            self.output_block_id += 1
            self.output_blocks[self.output_block_id] = (piece, -1)
            self.output_buffer_size += len(piece)
        outp.close()
        self.resend()

    def send_blocks(self):
        if self.consumer:
            relative_time = time.time() - self.creation_time
            current_rate = self.limit_send_bytes_per_sec
            rtt_current = self.rtt_avarage / self.rtt_acks_counter
            if relative_time > 0.0:
                current_rate = self.bytes_sent / relative_time
            if current_rate > self.limit_send_bytes_per_sec:
                return
            resend_time_limit = 4 * BLOCKS_PER_ACK * rtt_current
            new_blocks_counter = 0
            for block_id in self.output_blocks.keys():
                piece, time_sent = self.output_blocks[block_id]
                data_size = len(piece)
                if time_sent >= 0:
                    dt = relative_time - time_sent
                    if dt > resend_time_limit:
                        self.resend_bytes += data_size
                        # print 're -'s,
                    else:
                        # print 'skip', block_id, dt, self.last_ack_rtt
                        continue
                else:
                    pass
                    # print 'go', block_id
                time_sent = relative_time
                self.output_blocks[block_id] = (piece, time_sent)
                output = ''.join((
                    struct.pack('i', block_id),
                    piece))
                # DEBUG
                # import random
                # if random.randint(0, 9) >= 1:
                    # 10 % percent lost
                    # self.producer.do_send_data(self.stream_id, self.consumer, output)
                self.producer.do_send_data(self.stream_id, self.consumer, output)
                self.bytes_sent += data_size
                new_blocks_counter += 1
            if new_blocks_counter > 0:
                print 'send blocks', new_blocks_counter, len(self.output_blocks.keys()),
                print self.bytes_sent, self.bytes_acked, self.resend_bytes, rtt_current, resend_time_limit

    def send_ack(self):
        if self.consumer:
            ack_data = ''.join(map(lambda bid: struct.pack('i', bid), self.blocks_to_ack))
            self.producer.do_send_ack(self.stream_id, self.consumer, ack_data)
            self.output_blocks_acks += list(self.blocks_to_ack)
            print 'send ack', len(self.output_blocks_acks), len(self.blocks_to_ack), self.bytes_in
            self.blocks_to_ack.clear()
            self.last_ack_moment = time.time()

    def resend(self):
        self.resend_counter += 1
        activitiy = len(self.output_blocks.keys()) # + len(self.blocks_to_ack)
        if activitiy > 0:
            # print 'resend out:%s acks:%s' % (len(self.output_blocks.keys()), len(self.blocks_to_ack))
            self.resend_inactivity_counter = 0
        else:
            self.resend_inactivity_counter += 1
        rtt_current = self.rtt_avarage / self.rtt_acks_counter
        next_resend = max(min(rtt_current, RTT_MAX_LIMIT), RTT_MIN_LIMIT)
        if len(self.blocks_to_ack) > 0:
            if time.time() - self.last_ack_moment > next_resend:
                self.send_ack()
        if len(self.output_blocks):
            self.send_blocks()
#        if self.resend_inactivity_counter > 5:
#            # if self.resend_inactivity_counter % 10 == 1:
#            #     print 'drop resend out:%s acks:%s' % (
#            #         len(self.output_blocks.keys()), len(self.blocks_to_ack))
#            next_resend = RTT_MAX_LIMIT * 4.0
#        if self.resend_inactivity_counter > 50:
#            next_resend = RTT_MAX_LIMIT * 16.0
        if self.resend_counter % 500 == 1:
            print 'resend out:%d acks:%d' % (len(self.output_blocks.keys()), len(self.blocks_to_ack)),
            print 'rtt=%r, next=%r, iterations=%d' % (rtt_current, next_resend, self.resend_counter)
        next_resend *= self.resend_inactivity_counter
        if self.resend_task is None:
            self.resend_task = reactor.callLater(next_resend, self.resend)
            return
        if self.resend_task.called:
            self.resend_task = reactor.callLater(next_resend, self.resend)
            return
        if self.resend_task.cancelled:
            self.resend_task = None
            return

    def stop_resending(self):
        if self.resend_task:
            if self.resend_task.active():
                self.resend_task.cancel()
                self.resend_task = None

#------------------------------------------------------------------------------

class BufferOverflow(Exception):
    pass

<|endoftext|>


<commit_before>"""

Model Tests
-----------

These tests focus on the storage models themselves

"""

import unittest

from peewee import SqliteDatabase

from nhlstats.models import db_proxy, League, Season, SeasonType

db_proxy.initialize(SqliteDatabase(':memory:'))

class ModelTestCase(unittest.TestCase):
    MODELS = []

    def setUp(self):
        for model in self.MODELS:
            model.create_table()

    def tearDown(self):
        for model in self.MODELS:
            model.drop_table()

class TestModelLeague(ModelTestCase):
    MODELS = [League]

    def test_create(self):
        League.create(
            name='National Hockey League',
            abbreviation='NHL'
        )
        self.assertEqual(League.select().count(), 1)
        league = League.select().first()
        self.assertEqual(league.name, 'National Hockey League')
        self.assertEqual(league.abbreviation, 'NHL')

class TestModelSeasonType(ModelTestCase):
    MODELS = [League, SeasonType]

    def test_create(self):
        league = League.create(
            name='National Hockey League',
            abbreviation='NHL'
        )
        season_type = SeasonType.create(
            league=league,
            name='Regular'
        )
        self.assertEqual(league.season_types.count(), 1)
        self.assertIn(season_type, league.season_types)

class TestModelSeason(ModelTestCase):
    MODELS = [League, SeasonType, Season]

    def test_create(self):
        league = League.create(
            name='National Hockey League',
            abbreviation='NHL'
        )
        season_type = SeasonType.create(
            league=league,
            name='Regular'
        )
        season = season = Season.create(
            league=league,
            year='2014-15',
            type=season_type
        )
        self.assertEqual(league.seasons.count(), 1)
        self.assertIn(season, league.seasons)
<commit_msg>make the tests a little more DRY<commit_after>"""

Model Tests
-----------

These tests focus on the storage models themselves

"""

import unittest

from peewee import SqliteDatabase

from nhlstats.models import db_proxy, League, Season, SeasonType

db_proxy.initialize(SqliteDatabase(':memory:'))

class ModelTestCase(unittest.TestCase):
    MODELS = []

    def setUp(self):
        for model in self.MODELS:
            model.create_table()

    def tearDown(self):
        for model in self.MODELS:
            model.drop_table()

    def create_league(self, **kwargs):
        defaults = {'name': 'National Hockey League', 'abbreviation': 'NHL'}
        defaults.update(kwargs)
        return League.create(**defaults)

    def create_season_type(self, **kwargs):
        defaults = {'name': 'Regular'}
        defaults.update(kwargs)
        return SeasonType.create(**defaults)

    def create_season(self, **kwargs):
        defaults = {'year': '2014-15'}
        defaults.update(kwargs)
        return Season.create(**defaults)

class TestModelLeague(ModelTestCase):
    MODELS = [League]

    def test_create(self):
        self.assertEqual(League.select().count(), 0)
        league = self.create_league()
        self.assertEqual(League.select().count(), 1)
        league = League.select().first()
        self.assertEqual(league.name, 'National Hockey League')
        self.assertEqual(league.abbreviation, 'NHL')

class TestModelSeasonType(ModelTestCase):
    MODELS = [League, SeasonType]

    def test_create(self):
        league = self.create_league()
        season_type = self.create_season_type(league=league)
        self.assertEqual(league.season_types.count(), 1)
        self.assertIn(season_type, league.season_types)

class TestModelSeason(ModelTestCase):
    MODELS = [League, SeasonType, Season]

    def test_create(self):
        league = self.create_league()
        season_type = self.create_season_type(league=league)
        season = self.create_season(league=league, type=season_type)
        self.assertEqual(league.seasons.count(), 1)
        self.assertIn(season, league.seasons)
<|endoftext|>


<commit_before>import sublime, sublime_plugin, sys

# https://gist.github.com/idleberg/fca633438329cc5ae327
SCOPES = [
    [ "source.c++", ".source.cpp" ],
    [ "source.java-props", ".source.java-properties" ],
    [ "source.objc++", ".source.objcpp" ],
    [ "source.php", ".source.html.php" ],
    [ "source.scss", ".source.css.scss" ],
    [ "source.todo", ".text.todo" ],
    [ "source.markdown", ".source.gfm" ]
]

# Automatic conversion, based on scope
class AutomizrCommand(sublime_plugin.TextCommand):

    def run(self, edit):

        scope = self.view.scope_name(self.view.sel()[0].a)

        if "source.json" in scope:
            print("Atomizr: JSON detected, trying to convert")
            self.view.run_command('subl_to_atom')
        elif "source.coffee" in scope:
            print("Atomizr: CoffeeScript detected, trying to convert")
            self.view.run_command('atom_to_subl')
        elif "text.xml" in scope:
            print("Atomizr: XML detected, trying to convert")
            self.view.run_command('subl_snip_to_atom')
        elif "text.plain" in scope:
            sublime.error_message("Atomizr: Automatic conversion requires the 'Better CoffeeScript' package to be installed")
        else:
            sublime.error_message("Atomizr: Unsupported scope, aborting")

# Converts Sublime Text completions into Atom snippets
class SublToAtomCommand(sublime_plugin.TextCommand):

    def run(self, edit):
        import cson, json

        # read data from view
        selection = self.view.substr(sublime.Region(0, self.view.size()))

        # interprete and validate data
        try:
            data = json.loads(selection)
        except ValueError:
            sublime.error_message("Atomizr: Invalid JSON")
            return

        # but is it a Sublime Text completion?
        try:
            # get scope, convert if necessary
            for subl, atom in SCOPES:
                if data['scope'] == subl:
                    scope = atom
                    break
                else:
                    scope = "." + data['scope']

            completions = data['completions']
        except:
            sublime.error_message("Atomizr: Not a Sublime Text completions file")
            return

        array = {}

        for item in completions:
            try:
                array[item['trigger']] = {'prefix': item['trigger'], 'body':  item['contents']}
            except KeyError:
                pass

        atom = {scope: (array)}

        # write converted data to view
        selection = sublime.Region(0, self.view.size())
        self.view.replace(edit, selection, cson.dumps(atom, sort_keys=True, indent=4, separators=(',', ': ')))

        # set syntax to CSON, requires Better CoffeeScript package
        self.view.set_syntax_file("Packages/Better CoffeeScript/CoffeeScript.tmLanguage")

# Converts Sublime Text completions into Atom snippets
class SublSnipToAtomCommand(sublime_plugin.TextCommand):

    def run(self, edit):
        import cson, xmltodict

        # read data from view
        selection = self.view.substr(sublime.Region(0, self.view.size()))

        # xmltodict.parse(selection)
        # interprete and validate data
        try:
            xml = xmltodict.parse(selection)
        except:
            sublime.error_message("Atomizr: Invalid XML, aborting conversion")
            return

        body = xml['snippet']['content']
        scope = xml['snippet']['scope']
        prefix = xml['snippet']['tabTrigger']

        # <description> is optional
        try:
            xml['snippet']['description']
        except:
            description = prefix

        atom = {scope: { description: { "prefix": prefix, "body": body} } }

        # write converted data to view
        selection = sublime.Region(0, self.view.size())
        self.view.replace(edit, selection, cson.dumps(atom, sort_keys=True, indent=2, separators=(',', ': ')))

        # set syntax to CSON, requires Better CoffeeScript package
        self.view.set_syntax_file("Packages/Better CoffeeScript/CoffeeScript.tmLanguage")

# Converts Atom snippets into Sublime Text completions
class AtomToSublCommand(sublime_plugin.TextCommand):

    def run(self, edit):
        import cson, json

        # read data from view
        selection = self.view.substr(sublime.Region(0, self.view.size()))

        # interprete and validate data
        try:
            data = cson.loads(selection)
        except:
            sublime.error_message("Atomizr: Invalid CSON, aborting conversion")
            return

        completions = []

        # but is it an Atom snippet?
        try:
            # get scope, convert if necessary
            for key in data.keys():
                for subl, atom in SCOPES:
                    if key == atom:
                        scope = subl
                        break
                    else:
                        scope = key.lstrip(".")

                for item in (data[key]):
                    completions.append( {"trigger": data[key][item]["prefix"], "contents": data[key][item]["body"]} )
        except:
            sublime.error_message("Atomizr: Not an Atom snippet file")
            return

        subl = {"scope": scope, "completions": completions}

        selection = sublime.Region(0, self.view.size())
        self.view.replace(edit, selection, json.dumps(subl, sort_keys=False, indent=2, separators=(',', ': ')))

        # set syntax to JSON
        if sublime.version() >= "3103":
            self.view.set_syntax_file('Packages/JavaScript/JSON.sublime-syntax')
        else:
            self.view.set_syntax_file('Packages/JavaScript/JSON.tmLanguage')

# Converts Atom snippets (CSON into JSON)
class AtomToAtomCommand(sublime_plugin.TextCommand):

    def run(self, edit):
        import cson, json

        # read data from view
        selection = self.view.substr(sublime.Region(0, self.view.size()))

        # interprete and validate data
        try:
            data = cson.loads(selection)
        except:
            sublime.error_message("Atomizr: Invalid CSON, aborting conversion")
            return

        # write converted data to view
        selection = sublime.Region(0, self.view.size())
        self.view.replace(edit, selection, json.dumps(data, sort_keys=False, indent=2, separators=(',', ': ')))
<commit_msg>bump to cson 0.5<commit_after>import sublime, sublime_plugin, sys

# https://gist.github.com/idleberg/fca633438329cc5ae327
SCOPES = [
    [ "source.c++", ".source.cpp" ],
    [ "source.java-props", ".source.java-properties" ],
    [ "source.objc++", ".source.objcpp" ],
    [ "source.php", ".source.html.php" ],
    [ "source.scss", ".source.css.scss" ],
    [ "source.todo", ".text.todo" ],
    [ "source.markdown", ".source.gfm" ]
]

# Automatic conversion, based on scope
class AutomizrCommand(sublime_plugin.TextCommand):

    def run(self, edit):

        scope = self.view.scope_name(self.view.sel()[0].a)

        if "source.json" in scope:
            print("Atomizr: JSON detected, trying to convert")
            self.view.run_command('subl_to_atom')
        elif "source.coffee" in scope:
            print("Atomizr: CoffeeScript detected, trying to convert")
            self.view.run_command('atom_to_subl')
        elif "text.xml" in scope:
            print("Atomizr: XML detected, trying to convert")
            self.view.run_command('subl_snip_to_atom')
        elif "text.plain" in scope:
            sublime.error_message("Atomizr: Automatic conversion requires the 'Better CoffeeScript' package to be installed")
        else:
            sublime.error_message("Atomizr: Unsupported scope, aborting")

# Converts Sublime Text completions into Atom snippets
class SublToAtomCommand(sublime_plugin.TextCommand):

    def run(self, edit):
        import cson, json

        # read data from view
        selection = self.view.substr(sublime.Region(0, self.view.size()))

        # interprete and validate data
        try:
            data = json.loads(selection)
        except ValueError:
            sublime.error_message("Atomizr: Invalid JSON")
            return

        # but is it a Sublime Text completion?
        try:
            # get scope, convert if necessary
            for subl, atom in SCOPES:
                if data['scope'] == subl:
                    scope = atom
                    break
                else:
                    scope = "." + data['scope']

            completions = data['completions']
        except:
            sublime.error_message("Atomizr: Not a Sublime Text completions file")
            return

        array = {}

        for item in completions:
            try:
                array[item['trigger']] = {'prefix': item['trigger'], 'body':  item['contents']}
            except KeyError:
                pass

        atom = {scope: (array)}

        # write converted data to view
        selection = sublime.Region(0, self.view.size())
        self.view.replace(edit, selection, cson.dumps(atom, sort_keys=True, indent=4))

        # set syntax to CSON, requires Better CoffeeScript package
        self.view.set_syntax_file("Packages/Better CoffeeScript/CoffeeScript.tmLanguage")

# Converts Sublime Text completions into Atom snippets
class SublSnipToAtomCommand(sublime_plugin.TextCommand):

    def run(self, edit):
        import cson, xmltodict

        # read data from view
        selection = self.view.substr(sublime.Region(0, self.view.size()))

        # xmltodict.parse(selection)
        # interprete and validate data
        try:
            xml = xmltodict.parse(selection)
        except:
            sublime.error_message("Atomizr: Invalid XML, aborting conversion")
            return

        body = xml['snippet']['content']
        scope = xml['snippet']['scope']
        prefix = xml['snippet']['tabTrigger']

        # <description> is optional
        try:
            xml['snippet']['description']
        except:
            description = prefix

        atom = {scope: { description: { "prefix": prefix, "body": body} } }

        # write converted data to view
        selection = sublime.Region(0, self.view.size())
        self.view.replace(edit, selection, cson.dumps(atom, sort_keys=True, indent=2))

        # set syntax to CSON, requires Better CoffeeScript package
        self.view.set_syntax_file("Packages/Better CoffeeScript/CoffeeScript.tmLanguage")

# Converts Atom snippets into Sublime Text completions
class AtomToSublCommand(sublime_plugin.TextCommand):

    def run(self, edit):
        import cson, json

        # read data from view
        selection = self.view.substr(sublime.Region(0, self.view.size()))

        # interprete and validate data
        try:
            data = cson.loads(selection)
        except:
            sublime.error_message("Atomizr: Invalid CSON, aborting conversion")
            return

        completions = []

        # but is it an Atom snippet?
        try:
            # get scope, convert if necessary
            for key in data.keys():
                for subl, atom in SCOPES:
                    if key == atom:
                        scope = subl
                        break
                    else:
                        scope = key.lstrip(".")

                for item in (data[key]):
                    completions.append( {"trigger": data[key][item]["prefix"], "contents": data[key][item]["body"]} )
        except:
            sublime.error_message("Atomizr: Not an Atom snippet file")
            return

        subl = {"scope": scope, "completions": completions}

        selection = sublime.Region(0, self.view.size())
        self.view.replace(edit, selection, json.dumps(subl, sort_keys=False, indent=2, separators=(',', ': ')))

        # set syntax to JSON
        if sublime.version() >= "3103":
            self.view.set_syntax_file('Packages/JavaScript/JSON.sublime-syntax')
        else:
            self.view.set_syntax_file('Packages/JavaScript/JSON.tmLanguage')

# Converts Atom snippets (CSON into JSON)
class AtomToAtomCommand(sublime_plugin.TextCommand):

    def run(self, edit):
        import cson, json

        # read data from view
        selection = self.view.substr(sublime.Region(0, self.view.size()))

        # interprete and validate data
        try:
            data = cson.loads(selection)
        except:
            sublime.error_message("Atomizr: Invalid CSON, aborting conversion")
            return

        # write converted data to view
        selection = sublime.Region(0, self.view.size())
        self.view.replace(edit, selection, json.dumps(data, sort_keys=False, indent=2, separators=(',', ': ')))
<|endoftext|>


<commit_before># coding=utf-8
# Copyright 2020 The Trax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Nesterov momentum optimizer (also known as Nesterov Accelerated Gradient)."""

from trax.math import numpy as np
from trax.optimizers import base

# TODO(jonni): Consider renaming this class to NesterovMomentum.
class Momentum(base.Optimizer):
  r"""A Nesterov momentum optimizer.

  This class implements a Nesterov momentum variant of stochastic gradient
  descent (SGD). The implementation is based on the concepts in Sutskever et
  al. (2013) [http://jmlr.org/proceedings/papers/v28/sutskever13.pdf],
  reformulated in Bengio et al. (2012) [https://arxiv.org/abs/1212.0901],
  to work well with backpropagation (equations 6 and 7):

  $$v_t = \mu_{t-1}v_{t-1} - \epsilon_{t-1}\nabla f(\Theta_{t-1})$$

  $$\Theta_t = \Theta_{t-1} - \mu_{t-1} v_{t-1} + \mu_t v_t + v_t$$

  where $$\mu_{t-1}$$ is the momentum (decay) coefficient at time step $$t-1$$
  and $$\epsilon_{t-1}$$ is the learning rate at time step $$t-1$$.

  Note that the implementation below also includes a weight decay rate
  ($$\alpha$$) on the parameters, independent of the Nesterov momentum.
  """

  def __init__(self, learning_rate, mass=0.9, weight_decay_rate=1e-5):  # pylint: disable=useless-super-delegation
    super(Momentum, self).__init__(
        learning_rate=learning_rate,
        mass=mass,
        weight_decay_rate=weight_decay_rate,
    )

  def init(self, weights):
    return np.zeros_like(weights)

  def update(self, step, grads, weights, velocity, opt_params):
    del step
    v = velocity
    mu = opt_params['mass']
    alpha = opt_params['weight_decay_rate']
    epsilon = opt_params['learning_rate']

    new_v = mu * v + grads
    new_weights = (1 - alpha) * weights - epsilon * (mu * new_v + grads)

    new_weights = new_weights.astype(weights.dtype)
    return (new_weights, new_v)
<commit_msg>Implement non-Nesterov momentum update.<commit_after># coding=utf-8
# Copyright 2020 The Trax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Nesterov momentum optimizer (also known as Nesterov Accelerated Gradient)."""

from trax.math import numpy as np
from trax.optimizers import base

# TODO(jonni): Consider renaming this class to NesterovMomentum.
class Momentum(base.Optimizer):
  r"""A momentum optimizer.

  This class implements two variants of momentum stochastic gradient descent
  (SGD): with and without the Nesterov correction. The implementation of the
  Nesterov update is based on the concepts in Sutskever et al. (2013)
  [http://jmlr.org/proceedings/papers/v28/sutskever13.pdf], reformulated in
  Bengio et al. (2012) [https://arxiv.org/abs/1212.0901], to work well with
  backpropagation (equations 6 and 7):

  $$v_t = \mu_{t-1}v_{t-1} - \epsilon_{t-1}\nabla f(\Theta_{t-1})$$

  $$\Theta_t = \Theta_{t-1} - \mu_{t-1} v_{t-1} + \mu_t v_t + v_t$$

  where $$\mu_{t-1}$$ is the momentum (decay) coefficient at time step $$t-1$$
  and $$\epsilon_{t-1}$$ is the learning rate at time step $$t-1$$.

  Note that the implementation below also includes a weight decay rate
  ($$\alpha$$) on the parameters, independent of the Nesterov momentum.
  """

  def __init__(
      self, learning_rate, mass=0.9, weight_decay_rate=1e-5, nesterov=True
  ):  # pylint: disable=useless-super-delegation
    super(Momentum, self).__init__(
        learning_rate=learning_rate,
        mass=mass,
        weight_decay_rate=weight_decay_rate,
    )
    self._nesterov = nesterov

  def init(self, weights):
    return np.zeros_like(weights)

  def update(self, step, grads, weights, velocity, opt_params):
    del step
    v = velocity
    mu = opt_params['mass']
    alpha = opt_params['weight_decay_rate']
    epsilon = opt_params['learning_rate']

    new_v = mu * v + grads
    if self._nesterov:
      weight_update = mu * new_v + grads
    else:
      weight_update = new_v
    new_weights = (1 - alpha) * weights - epsilon * weight_update

    new_weights = new_weights.astype(weights.dtype)
    return (new_weights, new_v)
<|endoftext|>