<commit_before># No shebang line, this module is meant to be imported
#
# Copyright 2013 Oliver Palmer
#
# 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.

"""
Enums
=====

Provides enum values for certain aspect of PyFarm.  See below for more
detailed information.

Operating System
----------------

Describes an operating system type.

.. csv-table:: **OperatingSystem**
    :header: Attribute, Description
    :widths: 10, 50

    LINUX, operating system on agent is a Linux variant
    WINDOWS, operating system on agent is a Windows variant
    MAC, operating system on agent is an Apple OS variant

Job Type Load Mode
------------------

Determines how a custom job type will be loaded.

.. csv-table::
    :header: Attribute, Description
    :widths: 10, 50

    DOWNLOAD, download the jobtype file from a url
    OPEN, open the jobtype file from a url
    IMPORT, import the jobtype from the given string (ex. `foo.bar.ClassName`)

Agent State
-----------

The last known state of the remote agent, used for making queue decisions
and locking off resources.

.. csv-table::
    :header: Attribute, Description
    :widths: 10, 50

    OFFLINE, agent cannot be reached
    ONLINE, agent is waiting for work
    DISABLED, agent is online but cannot accept work
    RUNNING, agent is currently processing work
    ALLOC, special internal state used when the agent entry is being built

Work State
----------

The state a job or task is currently in.  These values apply more directly
to tasks as job statuses are built from task status values.

.. csv-table::
    :header: Attribute, Description
    :widths: 10, 50

    PAUSED, this task cannot be assigned right now but can be once unpaused
    BLOCKED, this task cannot be assigned to an agent at this point in time
    QUEUED, waiting on queue to assign this work
    ASSIGN, work has been assigned to an agent but is waiting to start
    RUNNING, work is currently being processed
    DONE, work is finished (previous failures may be present)
    FAILED, work as failed and cannot be continued
    ALLOC, special internal state for a job or task entry is being built

REST API Errors
---------------

Various error which the REST api may throw.  Numerical values will remain
constant however the error message may be rewritten.

.. csv-table::
    :header: Attribute, Integer Value, Description
    :widths: 10, 5, 50

    JSON_DECODE_FAILED, 0, failed to decode any json data from the request
    UNEXPECTED_DATATYPE, 1, the base data type decoded for the json class was not what was expected
    MISSING_FIELDS, 2, one or more of the expected fields were missing in the request
    UNEXPECTED_NULL, 3, a null value was found in a field that requires a non-null value
    DATABASE_ERROR, 4, problem inserting or updating entry in database
    EXTRA_FIELDS_ERROR, 5, an unexpected number of fields or columns were provided

Use Agent Address
-----------------

Describes which address should be used to contact the agent

.. csv-table::
    :header: Attribute, Description
    :widths: 10, 50

    LOCAL, use the address which was provided by the agent
    REMOTE, use the address which we received the request from
    HOSTNAME, disregard both the local IP and the remote IP and use the hostname
    PASSIVE, agent cannot be contacted but will still request work and process jobs
"""

import sys
from warnings import warn

try:
    from collections import namedtuple
except ImportError:  # pragma: no cover
    from pyfarm.core.backports import namedtuple

from pyfarm.core.warning import NotImplementedWarning

def Enum(classname, **kwargs):
    """
    Produce an enum object using :func:`.namedtuple`

    >>> Foo = Enum("Foo", A=1, B=2)
    >>> assert Foo.A == 1 and Foo.B == 2
    >>> FooTemplate = Enum("Foo", A=int, instance=False)
    >>> Foo = FooTemplate(A=1)
    >>> assert Foo.A == 1

    :param str classname:
        the name of the class to produce

    :keyword to_dict:
        a callable function to add to the named tuple for
        converting the internal values into a dictionary

    :keyword bool instance:
        by default calling :func:`.Enum` will produce an instanced
        :func:`.namedtuple` object, setting ``instance`` to False
        will instead produce the named tuple without instancing it
    """
    to_dict = kwargs.pop("to_dict", None)
    instance = kwargs.pop("instance", True)
    template = namedtuple(classname, kwargs.keys())

    if to_dict is not None:
        setattr(template, "to_dict", to_dict)

    return template(**kwargs) if instance else template

class Values(namedtuple("EnumValue", ("int", "str"))):
    """
    Stores values to be used in an enum.  Each time this
    class is instanced it will ensure that the input values
    are of the correct type and unique.
    """
    _integers = set()

    def __init__(self, *args, **kwargs):
        if not isinstance(self.int, int):
            raise TypeError("`int` must be an integer")

        if not isinstance(self.str, basestring):
            raise TypeError("`str` must be a string")

        if kwargs.get("unique", True) and self.int in self._integers:
            raise ValueError("value %s is being reused" % self.int)

        self._integers.add(self.int)
        self._values = set([self.int, self.str])

    def __contains__(self, item):
        return item in self._values

    def __eq__(self, other):
        if other is self or other in self:
            return True
        elif isinstance(other, Values):
            return other.int == self.int and other.str == self.str
        else:
            return False

    def __int__(self):
        return self.int

    def __str__(self):
        return self.str

    def __repr__(self):
        return "%s(%s, %s)" % (
            self.__class__.__name__, self.int, repr(self.str))

class EnumValue(object):
    """
    Special object which wraps an ``enum`` and a ``value`` so they can
    be compared to each other.
    """
    def __init__(self, enum, value):
        self._enum = enum
        self._value = value

        if isinstance(value, int):
            self._i = value
            self._s = self._enum._map[value]
        else:
            self._s = value
            self._i = self._enum._map[value]

        self._ints = []
        self._strings = []
        self._values = set([self._i, self._s])

        for value in enum:
            if isinstance(value, int):
                self._ints.append(value)
                self._strings.insert(0, self._enum._map[value])
            else:
                self._strings.insert(0, value)
                self._ints.append(self._enum._map[value])

    def __contains__(self, item):
        return item in self._values

    def __repr__(self):
        return "%s(%s, %s)" % (self.__class__.__name__, self._i, repr(self._s))

    def __ne__(self, other):
        return not self.__eq__(other)

    def __ge__(self, other):
        return self.__eq__(other) or self.__gt__(other)

    def __le__(self, other):
        return self.__eq__(other) or self.__lt__(other)

    def __eq__(self, other):
        if isinstance(other, EnumValue):
            return other._i == self._i
        elif other == self._i or other == self._s:
            return True
        else:
            return False

    def __gt__(self, other):
        if isinstance(other, EnumValue):
            return other._i > self._i
        elif other not in self._enum._map:
            raise ValueError(
                "%s is not part of this enum's values" % repr(other))
        elif isinstance(other, int):
            return self._ints.index(other) > self._ints.index(self._i)
        elif isinstance(other, basestring):
            return self._strings.index(other) > self._strings.index(self._s)
        else:
            return False

    def __lt__(self, other):
        if isinstance(other, EnumValue):
            return other._i < self._i
        elif other not in self._enum._map:
            raise ValueError(
                "%s is not part of this enum's values" % repr(other))
        elif isinstance(other, int):
            return self._ints.index(other) < self._ints.index(self._i)
        elif isinstance(other, basestring):
            return self._strings.index(other) < self._strings.index(self._s)
        else:
            return False

def cast_enum(enum, enum_type):
    """
    Pulls the requested ``enum_type`` from ``enum`` and produce a new
    named tuple which contains only the requested data

    >>> from pyfarm.core.enums import Enum, EnumValue
    >>> FooBase = Enum("Foo", A=EnumValue(int=1, str="1")
    >>> Foo = cast_enum(FooBase, str)
    >>> assert Foo.A == "1"
    >>> Foo = cast_enum(FooBase, int)
    >>> assert Foo.A == 1
    >>> assert Foo._map == {"A": 1, 1: "A"}

    .. warning::
        This function does not perform any kind of caching.  For the most
        efficient usage it should only be called once per process or
        module for a given enum and enum_type combination.
    """
    enum_data = {}
    reverse_map = {}

    # construct the reverse mapping and push
    # the request type into enum_data
    for key, value in enum._asdict().iteritems():
        reverse_map[value.int] = value.str
        reverse_map[value.str] = value.int

        if enum_type is int:
            enum_data[key] = value.int
        elif enum_type is str:
            enum_data[key] = value.str
        else:
            raise ValueError("valid values for `enum_type` are int or str")

    class MappedEnum(namedtuple(enum.__class__.__name__, enum_data.keys())):
        _map = reverse_map
        _enum = enum

        def __contains__(self, item):
            if item in self._map:
                return True
            else:
                for key, value in self._enum._asdict().iteritems():
                    if item in value:
                        return True
            return False

    return MappedEnum(**enum_data)

# 1xx - work states
# NOTE: these values are directly tested test_enums.test_direct_work_values
_WorkState = Enum(
    "WorkState",
    PAUSED=Values(100, "paused"),
    QUEUED=Values(101, "queued"),
    BLOCKED=Values(102, "blocked"),
    ALLOC=Values(103, "alloc"),
    ASSIGN=Values(104, "assign"),
    RUNNING=Values(105, "running"),
    DONE=Values(106, "done"),
    FAILED=Values(107, "failed"),
    JOBTYPE_FAILED_IMPORT=Values(108, "jobtype_failed_import"),
    JOBTYPE_INVALID_CLASS=Values(109, "jobtype_invalid_class"),
    NO_SUCH_COMMAND=Values(110, "no_such_command"))

# 2xx - agent states
# NOTE: these values are directly tested test_enums.test_direct_agent_values
_AgentState = Enum(
    "AgentState",
    DISABLED=Values(200, "disabled"),
    OFFLINE=Values(201, "offline"),
    ONLINE=Values(202, "online"),
    RUNNING=Values(203, "running"))

# 3xx - non-queue related modes or states
# NOTE: these values are directly tested test_enums.test_direct_os_values
_OperatingSystem = Enum(
    "OperatingSystem",
    LINUX=Values(300, "linux"),
    WINDOWS=Values(301, "windows"),
    MAC=Values(302, "mac"),
    OTHER=Values(303, "other"))

# NOTE: these values are directly tested test_enums.test_direct_agent_addr
_UseAgentAddress = Enum(
    "UseAgentAddress",
    LOCAL=Values(310, "local"),
    REMOTE=Values(311, "remote"),
    HOSTNAME=Values(312, "hostname"),
    PASSIVE=Values(313, "passive"))

# NOTE: these values are directly tested test_enums.test_direct_agent_addr
_JobTypeLoadMode = Enum(
    "JobTypeLoadMode",
    DOWNLOAD=Values(320, "download"),
    OPEN=Values(321, "open"),
    IMPORT=Values(322, "import"))

# cast the enums defined above
WorkState = cast_enum(_WorkState, str)
AgentState = cast_enum(_AgentState, str)
OperatingSystem = cast_enum(_OperatingSystem, str)
UseAgentAddress = cast_enum(_UseAgentAddress, str)
JobTypeLoadMode = cast_enum(_JobTypeLoadMode, str)
DBWorkState = cast_enum(_WorkState, int)
DBAgentState = cast_enum(_AgentState, int)
DBOperatingSystem = cast_enum(_OperatingSystem, int)
DBUseAgentAddress = cast_enum(_UseAgentAddress, int)
DBJobTypeLoadMode = cast_enum(_JobTypeLoadMode, int)

RUNNING_WORK_STATES = set([
    WorkState.ALLOC,
    WorkState.ASSIGN,
    WorkState.RUNNING])

DB_RUNNING_WORK_STATES = set([
    DBWorkState.ALLOC,
    DBWorkState.ASSIGN,
    DBWorkState.RUNNING])

FAILED_WORK_STATES = set([
    WorkState.FAILED,
    WorkState.JOBTYPE_FAILED_IMPORT,
    WorkState.JOBTYPE_INVALID_CLASS,
    WorkState.NO_SUCH_COMMAND])

DB_FAILED_WORK_STATES = set([
    DBWorkState.FAILED,
    DBWorkState.JOBTYPE_FAILED_IMPORT,
    DBWorkState.JOBTYPE_INVALID_CLASS,
    DBWorkState.NO_SUCH_COMMAND])

APIErrorValue = Enum(
    "APIErrorValue",
    value=int, description=str,
    instance=False, to_dict=lambda self: {
        "value": self.value, "description": self.description})

APIError = Enum(
    "APIError",
    JSON_DECODE_FAILED=APIErrorValue(
        1, "failed to decode any json data from the request"),
    UNEXPECTED_DATATYPE=APIErrorValue(
        2, "the base data type decoded for the json class was not what was "
           "expected"),
    MISSING_FIELDS=APIErrorValue(
        3, "one or more of the expected fields were missing in the request"),
    UNEXPECTED_NULL=APIErrorValue(
        4, "a null value was found in a field that requires a non-null value"),
    DATABASE_ERROR=APIErrorValue(
        5, "problem inserting or updating entry in database"),
    EXTRA_FIELDS_ERROR=APIErrorValue(
        6, "an unexpected number of fields or columns were provided"))

def get_operating_system(platform=sys.platform):
    """returns the operating system for the given platform"""
    if platform.startswith("linux"):
        return OperatingSystem.LINUX
    elif platform.startswith("win"):
        return OperatingSystem.WINDOWS
    elif platform.startswith("darwin"):
        return OperatingSystem.MAC
    else:
        warn("unknown operating system: %s" % platform, NotImplementedWarning)
        return OperatingSystem.OTHER

OS = get_operating_system()
<commit_msg>adding conversion for int() and str()<commit_after># No shebang line, this module is meant to be imported
#
# Copyright 2013 Oliver Palmer
#
# 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.

"""
Enums
=====

Provides enum values for certain aspect of PyFarm.  See below for more
detailed information.

Operating System
----------------

Describes an operating system type.

.. csv-table:: **OperatingSystem**
    :header: Attribute, Description
    :widths: 10, 50

    LINUX, operating system on agent is a Linux variant
    WINDOWS, operating system on agent is a Windows variant
    MAC, operating system on agent is an Apple OS variant

Job Type Load Mode
------------------

Determines how a custom job type will be loaded.

.. csv-table::
    :header: Attribute, Description
    :widths: 10, 50

    DOWNLOAD, download the jobtype file from a url
    OPEN, open the jobtype file from a url
    IMPORT, import the jobtype from the given string (ex. `foo.bar.ClassName`)

Agent State
-----------

The last known state of the remote agent, used for making queue decisions
and locking off resources.

.. csv-table::
    :header: Attribute, Description
    :widths: 10, 50

    OFFLINE, agent cannot be reached
    ONLINE, agent is waiting for work
    DISABLED, agent is online but cannot accept work
    RUNNING, agent is currently processing work
    ALLOC, special internal state used when the agent entry is being built

Work State
----------

The state a job or task is currently in.  These values apply more directly
to tasks as job statuses are built from task status values.

.. csv-table::
    :header: Attribute, Description
    :widths: 10, 50

    PAUSED, this task cannot be assigned right now but can be once unpaused
    BLOCKED, this task cannot be assigned to an agent at this point in time
    QUEUED, waiting on queue to assign this work
    ASSIGN, work has been assigned to an agent but is waiting to start
    RUNNING, work is currently being processed
    DONE, work is finished (previous failures may be present)
    FAILED, work as failed and cannot be continued
    ALLOC, special internal state for a job or task entry is being built

REST API Errors
---------------

Various error which the REST api may throw.  Numerical values will remain
constant however the error message may be rewritten.

.. csv-table::
    :header: Attribute, Integer Value, Description
    :widths: 10, 5, 50

    JSON_DECODE_FAILED, 0, failed to decode any json data from the request
    UNEXPECTED_DATATYPE, 1, the base data type decoded for the json class was not what was expected
    MISSING_FIELDS, 2, one or more of the expected fields were missing in the request
    UNEXPECTED_NULL, 3, a null value was found in a field that requires a non-null value
    DATABASE_ERROR, 4, problem inserting or updating entry in database
    EXTRA_FIELDS_ERROR, 5, an unexpected number of fields or columns were provided

Use Agent Address
-----------------

Describes which address should be used to contact the agent

.. csv-table::
    :header: Attribute, Description
    :widths: 10, 50

    LOCAL, use the address which was provided by the agent
    REMOTE, use the address which we received the request from
    HOSTNAME, disregard both the local IP and the remote IP and use the hostname
    PASSIVE, agent cannot be contacted but will still request work and process jobs
"""

import sys
from warnings import warn

try:
    from collections import namedtuple
except ImportError:  # pragma: no cover
    from pyfarm.core.backports import namedtuple

from pyfarm.core.warning import NotImplementedWarning

def Enum(classname, **kwargs):
    """
    Produce an enum object using :func:`.namedtuple`

    >>> Foo = Enum("Foo", A=1, B=2)
    >>> assert Foo.A == 1 and Foo.B == 2
    >>> FooTemplate = Enum("Foo", A=int, instance=False)
    >>> Foo = FooTemplate(A=1)
    >>> assert Foo.A == 1

    :param str classname:
        the name of the class to produce

    :keyword to_dict:
        a callable function to add to the named tuple for
        converting the internal values into a dictionary

    :keyword bool instance:
        by default calling :func:`.Enum` will produce an instanced
        :func:`.namedtuple` object, setting ``instance`` to False
        will instead produce the named tuple without instancing it
    """
    to_dict = kwargs.pop("to_dict", None)
    instance = kwargs.pop("instance", True)
    template = namedtuple(classname, kwargs.keys())

    if to_dict is not None:
        setattr(template, "to_dict", to_dict)

    return template(**kwargs) if instance else template

class Values(namedtuple("EnumValue", ("int", "str"))):
    """
    Stores values to be used in an enum.  Each time this
    class is instanced it will ensure that the input values
    are of the correct type and unique.
    """
    _integers = set()

    def __init__(self, *args, **kwargs):
        if not isinstance(self.int, int):
            raise TypeError("`int` must be an integer")

        if not isinstance(self.str, basestring):
            raise TypeError("`str` must be a string")

        if kwargs.get("unique", True) and self.int in self._integers:
            raise ValueError("value %s is being reused" % self.int)

        self._integers.add(self.int)
        self._values = set([self.int, self.str])

    def __contains__(self, item):
        return item in self._values

    def __eq__(self, other):
        if other is self or other in self:
            return True
        elif isinstance(other, Values):
            return other.int == self.int and other.str == self.str
        else:
            return False

    def __int__(self):
        return self.int

    def __str__(self):
        return self.str

    def __repr__(self):
        return "%s(%s, %s)" % (
            self.__class__.__name__, self.int, repr(self.str))

class EnumValue(object):
    """
    Special object which wraps an ``enum`` and a ``value`` so they can
    be compared to each other.
    """
    def __init__(self, enum, value):
        self._enum = enum
        self._value = value

        if isinstance(value, int):
            self._i = value
            self._s = self._enum._map[value]
        else:
            self._s = value
            self._i = self._enum._map[value]

        self._ints = []
        self._strings = []
        self._values = set([self._i, self._s])

        for value in enum:
            if isinstance(value, int):
                self._ints.append(value)
                self._strings.insert(0, self._enum._map[value])
            else:
                self._strings.insert(0, value)
                self._ints.append(self._enum._map[value])

    def __contains__(self, item):
        return item in self._values

    def __repr__(self):
        return "%s(%s, %s)" % (self.__class__.__name__, self._i, repr(self._s))

    def __str__(self):
        return self._s

    def __int__(self):
        return self._i

    def __ne__(self, other):
        return not self.__eq__(other)

    def __ge__(self, other):
        return self.__eq__(other) or self.__gt__(other)

    def __le__(self, other):
        return self.__eq__(other) or self.__lt__(other)

    def __eq__(self, other):
        if isinstance(other, EnumValue):
            return other._i == self._i
        elif other == self._i or other == self._s:
            return True
        else:
            return False

    def __gt__(self, other):
        if isinstance(other, EnumValue):
            return other._i > self._i
        elif other not in self._enum._map:
            raise ValueError(
                "%s is not part of this enum's values" % repr(other))
        elif isinstance(other, int):
            return self._ints.index(other) > self._ints.index(self._i)
        elif isinstance(other, basestring):
            return self._strings.index(other) > self._strings.index(self._s)
        else:
            return False

    def __lt__(self, other):
        if isinstance(other, EnumValue):
            return other._i < self._i
        elif other not in self._enum._map:
            raise ValueError(
                "%s is not part of this enum's values" % repr(other))
        elif isinstance(other, int):
            return self._ints.index(other) < self._ints.index(self._i)
        elif isinstance(other, basestring):
            return self._strings.index(other) < self._strings.index(self._s)
        else:
            return False

def cast_enum(enum, enum_type):
    """
    Pulls the requested ``enum_type`` from ``enum`` and produce a new
    named tuple which contains only the requested data

    >>> from pyfarm.core.enums import Enum, EnumValue
    >>> FooBase = Enum("Foo", A=EnumValue(int=1, str="1")
    >>> Foo = cast_enum(FooBase, str)
    >>> assert Foo.A == "1"
    >>> Foo = cast_enum(FooBase, int)
    >>> assert Foo.A == 1
    >>> assert Foo._map == {"A": 1, 1: "A"}

    .. warning::
        This function does not perform any kind of caching.  For the most
        efficient usage it should only be called once per process or
        module for a given enum and enum_type combination.
    """
    enum_data = {}
    reverse_map = {}

    # construct the reverse mapping and push
    # the request type into enum_data
    for key, value in enum._asdict().iteritems():
        reverse_map[value.int] = value.str
        reverse_map[value.str] = value.int

        if enum_type is int:
            enum_data[key] = value.int
        elif enum_type is str:
            enum_data[key] = value.str
        else:
            raise ValueError("valid values for `enum_type` are int or str")

    class MappedEnum(namedtuple(enum.__class__.__name__, enum_data.keys())):
        _map = reverse_map
        _enum = enum

        def __contains__(self, item):
            if item in self._map:
                return True
            else:
                for key, value in self._enum._asdict().iteritems():
                    if item in value:
                        return True
            return False

    return MappedEnum(**enum_data)

# 1xx - work states
# NOTE: these values are directly tested test_enums.test_direct_work_values
_WorkState = Enum(
    "WorkState",
    PAUSED=Values(100, "paused"),
    QUEUED=Values(101, "queued"),
    BLOCKED=Values(102, "blocked"),
    ALLOC=Values(103, "alloc"),
    ASSIGN=Values(104, "assign"),
    RUNNING=Values(105, "running"),
    DONE=Values(106, "done"),
    FAILED=Values(107, "failed"),
    JOBTYPE_FAILED_IMPORT=Values(108, "jobtype_failed_import"),
    JOBTYPE_INVALID_CLASS=Values(109, "jobtype_invalid_class"),
    NO_SUCH_COMMAND=Values(110, "no_such_command"))

# 2xx - agent states
# NOTE: these values are directly tested test_enums.test_direct_agent_values
_AgentState = Enum(
    "AgentState",
    DISABLED=Values(200, "disabled"),
    OFFLINE=Values(201, "offline"),
    ONLINE=Values(202, "online"),
    RUNNING=Values(203, "running"))

# 3xx - non-queue related modes or states
# NOTE: these values are directly tested test_enums.test_direct_os_values
_OperatingSystem = Enum(
    "OperatingSystem",
    LINUX=Values(300, "linux"),
    WINDOWS=Values(301, "windows"),
    MAC=Values(302, "mac"),
    OTHER=Values(303, "other"))

# NOTE: these values are directly tested test_enums.test_direct_agent_addr
_UseAgentAddress = Enum(
    "UseAgentAddress",
    LOCAL=Values(310, "local"),
    REMOTE=Values(311, "remote"),
    HOSTNAME=Values(312, "hostname"),
    PASSIVE=Values(313, "passive"))

# NOTE: these values are directly tested test_enums.test_direct_agent_addr
_JobTypeLoadMode = Enum(
    "JobTypeLoadMode",
    DOWNLOAD=Values(320, "download"),
    OPEN=Values(321, "open"),
    IMPORT=Values(322, "import"))

# cast the enums defined above
WorkState = cast_enum(_WorkState, str)
AgentState = cast_enum(_AgentState, str)
OperatingSystem = cast_enum(_OperatingSystem, str)
UseAgentAddress = cast_enum(_UseAgentAddress, str)
JobTypeLoadMode = cast_enum(_JobTypeLoadMode, str)
DBWorkState = cast_enum(_WorkState, int)
DBAgentState = cast_enum(_AgentState, int)
DBOperatingSystem = cast_enum(_OperatingSystem, int)
DBUseAgentAddress = cast_enum(_UseAgentAddress, int)
DBJobTypeLoadMode = cast_enum(_JobTypeLoadMode, int)

RUNNING_WORK_STATES = set([
    WorkState.ALLOC,
    WorkState.ASSIGN,
    WorkState.RUNNING])

DB_RUNNING_WORK_STATES = set([
    DBWorkState.ALLOC,
    DBWorkState.ASSIGN,
    DBWorkState.RUNNING])

FAILED_WORK_STATES = set([
    WorkState.FAILED,
    WorkState.JOBTYPE_FAILED_IMPORT,
    WorkState.JOBTYPE_INVALID_CLASS,
    WorkState.NO_SUCH_COMMAND])

DB_FAILED_WORK_STATES = set([
    DBWorkState.FAILED,
    DBWorkState.JOBTYPE_FAILED_IMPORT,
    DBWorkState.JOBTYPE_INVALID_CLASS,
    DBWorkState.NO_SUCH_COMMAND])

APIErrorValue = Enum(
    "APIErrorValue",
    value=int, description=str,
    instance=False, to_dict=lambda self: {
        "value": self.value, "description": self.description})

APIError = Enum(
    "APIError",
    JSON_DECODE_FAILED=APIErrorValue(
        1, "failed to decode any json data from the request"),
    UNEXPECTED_DATATYPE=APIErrorValue(
        2, "the base data type decoded for the json class was not what was "
           "expected"),
    MISSING_FIELDS=APIErrorValue(
        3, "one or more of the expected fields were missing in the request"),
    UNEXPECTED_NULL=APIErrorValue(
        4, "a null value was found in a field that requires a non-null value"),
    DATABASE_ERROR=APIErrorValue(
        5, "problem inserting or updating entry in database"),
    EXTRA_FIELDS_ERROR=APIErrorValue(
        6, "an unexpected number of fields or columns were provided"))

def get_operating_system(platform=sys.platform):
    """returns the operating system for the given platform"""
    if platform.startswith("linux"):
        return OperatingSystem.LINUX
    elif platform.startswith("win"):
        return OperatingSystem.WINDOWS
    elif platform.startswith("darwin"):
        return OperatingSystem.MAC
    else:
        warn("unknown operating system: %s" % platform, NotImplementedWarning)
        return OperatingSystem.OTHER

OS = get_operating_system()
<|endoftext|>


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

from __future__ import print_function

from datetime import datetime

from numpy import nan
from numpy.random import randn
import numpy as np

from pandas import DataFrame, Series, Index, Timestamp, DatetimeIndex
import pandas as pd
import pandas.tseries.offsets as offsets

from pandas.util.testing import (assert_almost_equal,
                                 assert_series_equal,
                                 assert_frame_equal,
                                 assertRaisesRegexp)

import pandas.util.testing as tm

from pandas.tests.frame.common import TestData

class TestDataFrameTimeSeriesMethods(tm.TestCase, TestData):

    _multiprocess_can_split_ = True

    def test_diff(self):
        the_diff = self.tsframe.diff(1)

        assert_series_equal(the_diff['A'],
                            self.tsframe['A'] - self.tsframe['A'].shift(1))

        # int dtype
        a = 10000000000000000
        b = a + 1
        s = Series([a, b])

        rs = DataFrame({'s': s}).diff()
        self.assertEqual(rs.s[1], 1)

        # mixed numeric
        tf = self.tsframe.astype('float32')
        the_diff = tf.diff(1)
        assert_series_equal(the_diff['A'],
                            tf['A'] - tf['A'].shift(1))

        # issue 10907
        df = pd.DataFrame({'y': pd.Series([2]), 'z': pd.Series([3])})
        df.insert(0, 'x', 1)
        result = df.diff(axis=1)
        expected = pd.DataFrame({'x': np.nan, 'y': pd.Series(
            1), 'z': pd.Series(1)}).astype('float64')
        assert_frame_equal(result, expected)

    def test_diff_timedelta(self):
        # GH 4533
        df = DataFrame(dict(time=[Timestamp('20130101 9:01'),
                                  Timestamp('20130101 9:02')],
                            value=[1.0, 2.0]))

        res = df.diff()
        exp = DataFrame([[pd.NaT, np.nan],
                         [pd.Timedelta('00:01:00'), 1]],
                        columns=['time', 'value'])
        assert_frame_equal(res, exp)

    def test_diff_mixed_dtype(self):
        df = DataFrame(np.random.randn(5, 3))
        df['A'] = np.array([1, 2, 3, 4, 5], dtype=object)

        result = df.diff()
        self.assertEqual(result[0].dtype, np.float64)

    def test_diff_neg_n(self):
        rs = self.tsframe.diff(-1)
        xp = self.tsframe - self.tsframe.shift(-1)
        assert_frame_equal(rs, xp)

    def test_diff_float_n(self):
        rs = self.tsframe.diff(1.)
        xp = self.tsframe.diff(1)
        assert_frame_equal(rs, xp)

    def test_diff_axis(self):
        # GH 9727
        df = DataFrame([[1., 2.], [3., 4.]])
        assert_frame_equal(df.diff(axis=1), DataFrame(
            [[np.nan, 1.], [np.nan, 1.]]))
        assert_frame_equal(df.diff(axis=0), DataFrame(
            [[np.nan, np.nan], [2., 2.]]))

    def test_pct_change(self):
        rs = self.tsframe.pct_change(fill_method=None)
        assert_frame_equal(rs, self.tsframe / self.tsframe.shift(1) - 1)

        rs = self.tsframe.pct_change(2)
        filled = self.tsframe.fillna(method='pad')
        assert_frame_equal(rs, filled / filled.shift(2) - 1)

        rs = self.tsframe.pct_change(fill_method='bfill', limit=1)
        filled = self.tsframe.fillna(method='bfill', limit=1)
        assert_frame_equal(rs, filled / filled.shift(1) - 1)

        rs = self.tsframe.pct_change(freq='5D')
        filled = self.tsframe.fillna(method='pad')
        assert_frame_equal(rs, filled / filled.shift(freq='5D') - 1)

    def test_pct_change_shift_over_nas(self):
        s = Series([1., 1.5, np.nan, 2.5, 3.])

        df = DataFrame({'a': s, 'b': s})

        chg = df.pct_change()
        expected = Series([np.nan, 0.5, np.nan, 2.5 / 1.5 - 1, .2])
        edf = DataFrame({'a': expected, 'b': expected})
        assert_frame_equal(chg, edf)

    def test_shift(self):
        # naive shift
        shiftedFrame = self.tsframe.shift(5)
        self.assert_index_equal(shiftedFrame.index, self.tsframe.index)

        shiftedSeries = self.tsframe['A'].shift(5)
        assert_series_equal(shiftedFrame['A'], shiftedSeries)

        shiftedFrame = self.tsframe.shift(-5)
        self.assert_index_equal(shiftedFrame.index, self.tsframe.index)

        shiftedSeries = self.tsframe['A'].shift(-5)
        assert_series_equal(shiftedFrame['A'], shiftedSeries)

        # shift by 0
        unshifted = self.tsframe.shift(0)
        assert_frame_equal(unshifted, self.tsframe)

        # shift by DateOffset
        shiftedFrame = self.tsframe.shift(5, freq=offsets.BDay())
        self.assertEqual(len(shiftedFrame), len(self.tsframe))

        shiftedFrame2 = self.tsframe.shift(5, freq='B')
        assert_frame_equal(shiftedFrame, shiftedFrame2)

        d = self.tsframe.index[0]
        shifted_d = d + offsets.BDay(5)
        assert_series_equal(self.tsframe.xs(d),
                            shiftedFrame.xs(shifted_d), check_names=False)

        # shift int frame
        int_shifted = self.intframe.shift(1)  # noqa

        # Shifting with PeriodIndex
        ps = tm.makePeriodFrame()
        shifted = ps.shift(1)
        unshifted = shifted.shift(-1)
        self.assert_index_equal(shifted.index, ps.index)
        self.assert_index_equal(unshifted.index, ps.index)
        tm.assert_numpy_array_equal(unshifted.ix[:, 0].valid().values,
                                    ps.ix[:-1, 0].values)

        shifted2 = ps.shift(1, 'B')
        shifted3 = ps.shift(1, offsets.BDay())
        assert_frame_equal(shifted2, shifted3)
        assert_frame_equal(ps, shifted2.shift(-1, 'B'))

        assertRaisesRegexp(ValueError, 'does not match PeriodIndex freq',
                           ps.shift, freq='D')

        # shift other axis
        # GH 6371
        df = DataFrame(np.random.rand(10, 5))
        expected = pd.concat([DataFrame(np.nan, index=df.index,
                                        columns=[0]),
                              df.iloc[:, 0:-1]],
                             ignore_index=True, axis=1)
        result = df.shift(1, axis=1)
        assert_frame_equal(result, expected)

        # shift named axis
        df = DataFrame(np.random.rand(10, 5))
        expected = pd.concat([DataFrame(np.nan, index=df.index,
                                        columns=[0]),
                              df.iloc[:, 0:-1]],
                             ignore_index=True, axis=1)
        result = df.shift(1, axis='columns')
        assert_frame_equal(result, expected)

    def test_shift_bool(self):
        df = DataFrame({'high': [True, False],
                        'low': [False, False]})
        rs = df.shift(1)
        xp = DataFrame(np.array([[np.nan, np.nan],
                                 [True, False]], dtype=object),
                       columns=['high', 'low'])
        assert_frame_equal(rs, xp)

    def test_shift_categorical(self):
        # GH 9416
        s1 = pd.Series(['a', 'b', 'c'], dtype='category')
        s2 = pd.Series(['A', 'B', 'C'], dtype='category')
        df = DataFrame({'one': s1, 'two': s2})
        rs = df.shift(1)
        xp = DataFrame({'one': s1.shift(1), 'two': s2.shift(1)})
        assert_frame_equal(rs, xp)

    def test_shift_empty(self):
        # Regression test for #8019
        df = DataFrame({'foo': []})
        rs = df.shift(-1)

        assert_frame_equal(df, rs)

    def test_tshift(self):
        # PeriodIndex
        ps = tm.makePeriodFrame()
        shifted = ps.tshift(1)
        unshifted = shifted.tshift(-1)

        assert_frame_equal(unshifted, ps)

        shifted2 = ps.tshift(freq='B')
        assert_frame_equal(shifted, shifted2)

        shifted3 = ps.tshift(freq=offsets.BDay())
        assert_frame_equal(shifted, shifted3)

        assertRaisesRegexp(ValueError, 'does not match', ps.tshift, freq='M')

        # DatetimeIndex
        shifted = self.tsframe.tshift(1)
        unshifted = shifted.tshift(-1)

        assert_frame_equal(self.tsframe, unshifted)

        shifted2 = self.tsframe.tshift(freq=self.tsframe.index.freq)
        assert_frame_equal(shifted, shifted2)

        inferred_ts = DataFrame(self.tsframe.values,
                                Index(np.asarray(self.tsframe.index)),
                                columns=self.tsframe.columns)
        shifted = inferred_ts.tshift(1)
        unshifted = shifted.tshift(-1)
        assert_frame_equal(shifted, self.tsframe.tshift(1))
        assert_frame_equal(unshifted, inferred_ts)

        no_freq = self.tsframe.ix[[0, 5, 7], :]
        self.assertRaises(ValueError, no_freq.tshift)

    def test_truncate(self):
        ts = self.tsframe[::3]

        start, end = self.tsframe.index[3], self.tsframe.index[6]

        start_missing = self.tsframe.index[2]
        end_missing = self.tsframe.index[7]

        # neither specified
        truncated = ts.truncate()
        assert_frame_equal(truncated, ts)

        # both specified
        expected = ts[1:3]

        truncated = ts.truncate(start, end)
        assert_frame_equal(truncated, expected)

        truncated = ts.truncate(start_missing, end_missing)
        assert_frame_equal(truncated, expected)

        # start specified
        expected = ts[1:]

        truncated = ts.truncate(before=start)
        assert_frame_equal(truncated, expected)

        truncated = ts.truncate(before=start_missing)
        assert_frame_equal(truncated, expected)

        # end specified
        expected = ts[:3]

        truncated = ts.truncate(after=end)
        assert_frame_equal(truncated, expected)

        truncated = ts.truncate(after=end_missing)
        assert_frame_equal(truncated, expected)

        self.assertRaises(ValueError, ts.truncate,
                          before=ts.index[-1] - 1,
                          after=ts.index[0] + 1)

    def test_truncate_copy(self):
        index = self.tsframe.index
        truncated = self.tsframe.truncate(index[5], index[10])
        truncated.values[:] = 5.
        self.assertFalse((self.tsframe.values[5:11] == 5).any())

    def test_asfreq(self):
        offset_monthly = self.tsframe.asfreq(offsets.BMonthEnd())
        rule_monthly = self.tsframe.asfreq('BM')

        assert_almost_equal(offset_monthly['A'], rule_monthly['A'])

        filled = rule_monthly.asfreq('B', method='pad')  # noqa
        # TODO: actually check that this worked.

        # don't forget!
        filled_dep = rule_monthly.asfreq('B', method='pad')  # noqa

        # test does not blow up on length-0 DataFrame
        zero_length = self.tsframe.reindex([])
        result = zero_length.asfreq('BM')
        self.assertIsNot(result, zero_length)

    def test_asfreq_datetimeindex(self):
        df = DataFrame({'A': [1, 2, 3]},
                       index=[datetime(2011, 11, 1), datetime(2011, 11, 2),
                              datetime(2011, 11, 3)])
        df = df.asfreq('B')
        tm.assertIsInstance(df.index, DatetimeIndex)

        ts = df['A'].asfreq('B')
        tm.assertIsInstance(ts.index, DatetimeIndex)

    def test_first_last_valid(self):
        N = len(self.frame.index)
        mat = randn(N)
        mat[:5] = nan
        mat[-5:] = nan

        frame = DataFrame({'foo': mat}, index=self.frame.index)
        index = frame.first_valid_index()

        self.assertEqual(index, frame.index[5])

        index = frame.last_valid_index()
        self.assertEqual(index, frame.index[-6])

        # GH12800
        empty = DataFrame()
        self.assertIsNone(empty.last_valid_index())
        self.assertIsNone(empty.first_valid_index())

    def test_operation_on_NaT(self):
        # Both NaT and Timestamp are in DataFrame.
        df = pd.DataFrame({'foo': [pd.NaT, pd.NaT,
                                   pd.Timestamp('2012-05-01')]})

        res = df.min()
        exp = pd.Series([pd.Timestamp('2012-05-01')], index=["foo"])
        tm.assert_series_equal(res, exp)

        res = df.max()
        exp = pd.Series([pd.Timestamp('2012-05-01')], index=["foo"])
        tm.assert_series_equal(res, exp)

        # GH12941, only NaTs are in DataFrame.
        df = pd.DataFrame({'foo': [pd.NaT, pd.NaT]})

        res = df.min()
        exp = pd.Series([pd.NaT], index=["foo"])
        tm.assert_series_equal(res, exp)

        res = df.max()
        exp = pd.Series([pd.NaT], index=["foo"])
        tm.assert_series_equal(res, exp)

if __name__ == '__main__':
    import nose
    nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
                   exit=False)
<commit_msg>TST: Test datetime array assignment with different units (#7492) (#14884)<commit_after># -*- coding: utf-8 -*-

from __future__ import print_function

from datetime import datetime

from numpy import nan
from numpy.random import randn
import numpy as np

from pandas import DataFrame, Series, Index, Timestamp, DatetimeIndex
import pandas as pd
import pandas.tseries.offsets as offsets

from pandas.util.testing import (assert_almost_equal,
                                 assert_series_equal,
                                 assert_frame_equal,
                                 assertRaisesRegexp)

import pandas.util.testing as tm

from pandas.tests.frame.common import TestData

class TestDataFrameTimeSeriesMethods(tm.TestCase, TestData):

    _multiprocess_can_split_ = True

    def test_diff(self):
        the_diff = self.tsframe.diff(1)

        assert_series_equal(the_diff['A'],
                            self.tsframe['A'] - self.tsframe['A'].shift(1))

        # int dtype
        a = 10000000000000000
        b = a + 1
        s = Series([a, b])

        rs = DataFrame({'s': s}).diff()
        self.assertEqual(rs.s[1], 1)

        # mixed numeric
        tf = self.tsframe.astype('float32')
        the_diff = tf.diff(1)
        assert_series_equal(the_diff['A'],
                            tf['A'] - tf['A'].shift(1))

        # issue 10907
        df = pd.DataFrame({'y': pd.Series([2]), 'z': pd.Series([3])})
        df.insert(0, 'x', 1)
        result = df.diff(axis=1)
        expected = pd.DataFrame({'x': np.nan, 'y': pd.Series(
            1), 'z': pd.Series(1)}).astype('float64')
        assert_frame_equal(result, expected)

    def test_diff_timedelta(self):
        # GH 4533
        df = DataFrame(dict(time=[Timestamp('20130101 9:01'),
                                  Timestamp('20130101 9:02')],
                            value=[1.0, 2.0]))

        res = df.diff()
        exp = DataFrame([[pd.NaT, np.nan],
                         [pd.Timedelta('00:01:00'), 1]],
                        columns=['time', 'value'])
        assert_frame_equal(res, exp)

    def test_diff_mixed_dtype(self):
        df = DataFrame(np.random.randn(5, 3))
        df['A'] = np.array([1, 2, 3, 4, 5], dtype=object)

        result = df.diff()
        self.assertEqual(result[0].dtype, np.float64)

    def test_diff_neg_n(self):
        rs = self.tsframe.diff(-1)
        xp = self.tsframe - self.tsframe.shift(-1)
        assert_frame_equal(rs, xp)

    def test_diff_float_n(self):
        rs = self.tsframe.diff(1.)
        xp = self.tsframe.diff(1)
        assert_frame_equal(rs, xp)

    def test_diff_axis(self):
        # GH 9727
        df = DataFrame([[1., 2.], [3., 4.]])
        assert_frame_equal(df.diff(axis=1), DataFrame(
            [[np.nan, 1.], [np.nan, 1.]]))
        assert_frame_equal(df.diff(axis=0), DataFrame(
            [[np.nan, np.nan], [2., 2.]]))

    def test_pct_change(self):
        rs = self.tsframe.pct_change(fill_method=None)
        assert_frame_equal(rs, self.tsframe / self.tsframe.shift(1) - 1)

        rs = self.tsframe.pct_change(2)
        filled = self.tsframe.fillna(method='pad')
        assert_frame_equal(rs, filled / filled.shift(2) - 1)

        rs = self.tsframe.pct_change(fill_method='bfill', limit=1)
        filled = self.tsframe.fillna(method='bfill', limit=1)
        assert_frame_equal(rs, filled / filled.shift(1) - 1)

        rs = self.tsframe.pct_change(freq='5D')
        filled = self.tsframe.fillna(method='pad')
        assert_frame_equal(rs, filled / filled.shift(freq='5D') - 1)

    def test_pct_change_shift_over_nas(self):
        s = Series([1., 1.5, np.nan, 2.5, 3.])

        df = DataFrame({'a': s, 'b': s})

        chg = df.pct_change()
        expected = Series([np.nan, 0.5, np.nan, 2.5 / 1.5 - 1, .2])
        edf = DataFrame({'a': expected, 'b': expected})
        assert_frame_equal(chg, edf)

    def test_shift(self):
        # naive shift
        shiftedFrame = self.tsframe.shift(5)
        self.assert_index_equal(shiftedFrame.index, self.tsframe.index)

        shiftedSeries = self.tsframe['A'].shift(5)
        assert_series_equal(shiftedFrame['A'], shiftedSeries)

        shiftedFrame = self.tsframe.shift(-5)
        self.assert_index_equal(shiftedFrame.index, self.tsframe.index)

        shiftedSeries = self.tsframe['A'].shift(-5)
        assert_series_equal(shiftedFrame['A'], shiftedSeries)

        # shift by 0
        unshifted = self.tsframe.shift(0)
        assert_frame_equal(unshifted, self.tsframe)

        # shift by DateOffset
        shiftedFrame = self.tsframe.shift(5, freq=offsets.BDay())
        self.assertEqual(len(shiftedFrame), len(self.tsframe))

        shiftedFrame2 = self.tsframe.shift(5, freq='B')
        assert_frame_equal(shiftedFrame, shiftedFrame2)

        d = self.tsframe.index[0]
        shifted_d = d + offsets.BDay(5)
        assert_series_equal(self.tsframe.xs(d),
                            shiftedFrame.xs(shifted_d), check_names=False)

        # shift int frame
        int_shifted = self.intframe.shift(1)  # noqa

        # Shifting with PeriodIndex
        ps = tm.makePeriodFrame()
        shifted = ps.shift(1)
        unshifted = shifted.shift(-1)
        self.assert_index_equal(shifted.index, ps.index)
        self.assert_index_equal(unshifted.index, ps.index)
        tm.assert_numpy_array_equal(unshifted.ix[:, 0].valid().values,
                                    ps.ix[:-1, 0].values)

        shifted2 = ps.shift(1, 'B')
        shifted3 = ps.shift(1, offsets.BDay())
        assert_frame_equal(shifted2, shifted3)
        assert_frame_equal(ps, shifted2.shift(-1, 'B'))

        assertRaisesRegexp(ValueError, 'does not match PeriodIndex freq',
                           ps.shift, freq='D')

        # shift other axis
        # GH 6371
        df = DataFrame(np.random.rand(10, 5))
        expected = pd.concat([DataFrame(np.nan, index=df.index,
                                        columns=[0]),
                              df.iloc[:, 0:-1]],
                             ignore_index=True, axis=1)
        result = df.shift(1, axis=1)
        assert_frame_equal(result, expected)

        # shift named axis
        df = DataFrame(np.random.rand(10, 5))
        expected = pd.concat([DataFrame(np.nan, index=df.index,
                                        columns=[0]),
                              df.iloc[:, 0:-1]],
                             ignore_index=True, axis=1)
        result = df.shift(1, axis='columns')
        assert_frame_equal(result, expected)

    def test_shift_bool(self):
        df = DataFrame({'high': [True, False],
                        'low': [False, False]})
        rs = df.shift(1)
        xp = DataFrame(np.array([[np.nan, np.nan],
                                 [True, False]], dtype=object),
                       columns=['high', 'low'])
        assert_frame_equal(rs, xp)

    def test_shift_categorical(self):
        # GH 9416
        s1 = pd.Series(['a', 'b', 'c'], dtype='category')
        s2 = pd.Series(['A', 'B', 'C'], dtype='category')
        df = DataFrame({'one': s1, 'two': s2})
        rs = df.shift(1)
        xp = DataFrame({'one': s1.shift(1), 'two': s2.shift(1)})
        assert_frame_equal(rs, xp)

    def test_shift_empty(self):
        # Regression test for #8019
        df = DataFrame({'foo': []})
        rs = df.shift(-1)

        assert_frame_equal(df, rs)

    def test_tshift(self):
        # PeriodIndex
        ps = tm.makePeriodFrame()
        shifted = ps.tshift(1)
        unshifted = shifted.tshift(-1)

        assert_frame_equal(unshifted, ps)

        shifted2 = ps.tshift(freq='B')
        assert_frame_equal(shifted, shifted2)

        shifted3 = ps.tshift(freq=offsets.BDay())
        assert_frame_equal(shifted, shifted3)

        assertRaisesRegexp(ValueError, 'does not match', ps.tshift, freq='M')

        # DatetimeIndex
        shifted = self.tsframe.tshift(1)
        unshifted = shifted.tshift(-1)

        assert_frame_equal(self.tsframe, unshifted)

        shifted2 = self.tsframe.tshift(freq=self.tsframe.index.freq)
        assert_frame_equal(shifted, shifted2)

        inferred_ts = DataFrame(self.tsframe.values,
                                Index(np.asarray(self.tsframe.index)),
                                columns=self.tsframe.columns)
        shifted = inferred_ts.tshift(1)
        unshifted = shifted.tshift(-1)
        assert_frame_equal(shifted, self.tsframe.tshift(1))
        assert_frame_equal(unshifted, inferred_ts)

        no_freq = self.tsframe.ix[[0, 5, 7], :]
        self.assertRaises(ValueError, no_freq.tshift)

    def test_truncate(self):
        ts = self.tsframe[::3]

        start, end = self.tsframe.index[3], self.tsframe.index[6]

        start_missing = self.tsframe.index[2]
        end_missing = self.tsframe.index[7]

        # neither specified
        truncated = ts.truncate()
        assert_frame_equal(truncated, ts)

        # both specified
        expected = ts[1:3]

        truncated = ts.truncate(start, end)
        assert_frame_equal(truncated, expected)

        truncated = ts.truncate(start_missing, end_missing)
        assert_frame_equal(truncated, expected)

        # start specified
        expected = ts[1:]

        truncated = ts.truncate(before=start)
        assert_frame_equal(truncated, expected)

        truncated = ts.truncate(before=start_missing)
        assert_frame_equal(truncated, expected)

        # end specified
        expected = ts[:3]

        truncated = ts.truncate(after=end)
        assert_frame_equal(truncated, expected)

        truncated = ts.truncate(after=end_missing)
        assert_frame_equal(truncated, expected)

        self.assertRaises(ValueError, ts.truncate,
                          before=ts.index[-1] - 1,
                          after=ts.index[0] + 1)

    def test_truncate_copy(self):
        index = self.tsframe.index
        truncated = self.tsframe.truncate(index[5], index[10])
        truncated.values[:] = 5.
        self.assertFalse((self.tsframe.values[5:11] == 5).any())

    def test_asfreq(self):
        offset_monthly = self.tsframe.asfreq(offsets.BMonthEnd())
        rule_monthly = self.tsframe.asfreq('BM')

        assert_almost_equal(offset_monthly['A'], rule_monthly['A'])

        filled = rule_monthly.asfreq('B', method='pad')  # noqa
        # TODO: actually check that this worked.

        # don't forget!
        filled_dep = rule_monthly.asfreq('B', method='pad')  # noqa

        # test does not blow up on length-0 DataFrame
        zero_length = self.tsframe.reindex([])
        result = zero_length.asfreq('BM')
        self.assertIsNot(result, zero_length)

    def test_asfreq_datetimeindex(self):
        df = DataFrame({'A': [1, 2, 3]},
                       index=[datetime(2011, 11, 1), datetime(2011, 11, 2),
                              datetime(2011, 11, 3)])
        df = df.asfreq('B')
        tm.assertIsInstance(df.index, DatetimeIndex)

        ts = df['A'].asfreq('B')
        tm.assertIsInstance(ts.index, DatetimeIndex)

    def test_first_last_valid(self):
        N = len(self.frame.index)
        mat = randn(N)
        mat[:5] = nan
        mat[-5:] = nan

        frame = DataFrame({'foo': mat}, index=self.frame.index)
        index = frame.first_valid_index()

        self.assertEqual(index, frame.index[5])

        index = frame.last_valid_index()
        self.assertEqual(index, frame.index[-6])

        # GH12800
        empty = DataFrame()
        self.assertIsNone(empty.last_valid_index())
        self.assertIsNone(empty.first_valid_index())

    def test_operation_on_NaT(self):
        # Both NaT and Timestamp are in DataFrame.
        df = pd.DataFrame({'foo': [pd.NaT, pd.NaT,
                                   pd.Timestamp('2012-05-01')]})

        res = df.min()
        exp = pd.Series([pd.Timestamp('2012-05-01')], index=["foo"])
        tm.assert_series_equal(res, exp)

        res = df.max()
        exp = pd.Series([pd.Timestamp('2012-05-01')], index=["foo"])
        tm.assert_series_equal(res, exp)

        # GH12941, only NaTs are in DataFrame.
        df = pd.DataFrame({'foo': [pd.NaT, pd.NaT]})

        res = df.min()
        exp = pd.Series([pd.NaT], index=["foo"])
        tm.assert_series_equal(res, exp)

        res = df.max()
        exp = pd.Series([pd.NaT], index=["foo"])
        tm.assert_series_equal(res, exp)

    def test_datetime_assignment_with_NaT_and_diff_time_units(self):
        # GH 7492
        data_ns = np.array([1, 'nat'], dtype='datetime64[ns]')
        result = pd.Series(data_ns).to_frame()
        result['new'] = data_ns
        expected = pd.DataFrame({0: [1, None],
                                'new': [1, None]}, dtype='datetime64[ns]')
        tm.assert_frame_equal(result, expected)
        # OutOfBoundsDatetime error shouldn't occur
        data_s = np.array([1, 'nat'], dtype='datetime64[s]')
        result['new'] = data_s
        expected = pd.DataFrame({0: [1, None],
                                'new': [1e9, None]}, dtype='datetime64[ns]')
        tm.assert_frame_equal(result, expected)

if __name__ == '__main__':
    import nose
    nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
                   exit=False)
<|endoftext|>


<commit_before>"""
Provide the Reporter class.
"""

import sys

class Reporter(object):
    """
    Formats the results of pyflakes checks to users.
    """

    def __init__(self, warningStream, errorStream):
        """
        Construct a L{Reporter}.

        @param warningStream: A file-like object where warnings will be
            written to.  The stream's C{write} method must accept unicode.
            C{sys.stdout} is a good value.
        @param errorStream: A file-like object where error output will be
            written to.  The stream's C{write} method must accept unicode.
            C{sys.stderr} is a good value.
        """
        self._stdout = warningStream
        self._stderr = errorStream

    def unexpectedError(self, filename, msg):
        """
        An unexpected error occurred trying to process C{filename}.

        @param filename: The path to a file that we could not process.
        @ptype filename: C{unicode}
        @param msg: A message explaining the problem.
        @ptype msg: C{unicode}
        """
        self._stderr.write("%s: %s\n" % (filename, msg))

    def syntaxError(self, filename, msg, lineno, offset, text):
        """
        There was a syntax errror in C{filename}.

        @param filename: The path to the file with the syntax error.
        @ptype filename: C{unicode}
        @param msg: An explanation of the syntax error.
        @ptype msg: C{unicode}
        @param lineno: The line number where the syntax error occurred.
        @ptype lineno: C{int}
        @param offset: The column on which the syntax error occurred.
        @ptype offset: C{int}
        @param text: The source code containing the syntax error.
        @ptype text: C{unicode}
        """
        line = text.splitlines()[-1]
        if offset is not None:
            offset = offset - (len(text) - len(line))
        self._stderr.write('%s:%d: %s\n' % (filename, lineno, msg))
        self._stderr.write(line)
        self._stderr.write('\n')
        if offset is not None:
            self._stderr.write(" " * (offset + 1) + "^\n")

    def flake(self, message):
        """
        pyflakes found something wrong with the code.

        @param: A L{pyflakes.messages.Message}.
        """
        self._stdout.write(str(message))
        self._stdout.write('\n')

def _makeDefaultReporter():
    """
    Make a reporter that can be used when no reporter is specified.
    """
    return Reporter(sys.stdout, sys.stderr)
<commit_msg>Update reporter.py<commit_after>"""
Provide the Reporter class.
"""

import sys

class Reporter(object):
    """
    Formats the results of pyflakes checks to users.
    """

    def __init__(self, warningStream, errorStream):
        """
        Construct a L{Reporter}.

        @param warningStream: A file-like object where warnings will be
            written to.  The stream's C{write} method must accept unicode.
            C{sys.stdout} is a good value.
        @param errorStream: A file-like object where error output will be
            written to.  The stream's C{write} method must accept unicode.
            C{sys.stderr} is a good value.
        """
        self._stdout = warningStream
        self._stderr = errorStream

    def unexpectedError(self, filename, msg):
        """
        An unexpected error occurred trying to process C{filename}.

        @param filename: The path to a file that we could not process.
        @ptype filename: C{unicode}
        @param msg: A message explaining the problem.
        @ptype msg: C{unicode}
        """
        self._stderr.write("%s: %s\n" % (filename, msg))

    def syntaxError(self, filename, msg, lineno, offset, text):
        """
        There was a syntax error in C{filename}.

        @param filename: The path to the file with the syntax error.
        @ptype filename: C{unicode}
        @param msg: An explanation of the syntax error.
        @ptype msg: C{unicode}
        @param lineno: The line number where the syntax error occurred.
        @ptype lineno: C{int}
        @param offset: The column on which the syntax error occurred.
        @ptype offset: C{int}
        @param text: The source code containing the syntax error.
        @ptype text: C{unicode}
        """
        line = text.splitlines()[-1]
        if offset is not None:
            offset = offset - (len(text) - len(line))
        self._stderr.write('%s:%d: %s\n' % (filename, lineno, msg))
        self._stderr.write(line)
        self._stderr.write('\n')
        if offset is not None:
            self._stderr.write(" " * (offset + 1) + "^\n")

    def flake(self, message):
        """
        pyflakes found something wrong with the code.

        @param: A L{pyflakes.messages.Message}.
        """
        self._stdout.write(str(message))
        self._stdout.write('\n')

def _makeDefaultReporter():
    """
    Make a reporter that can be used when no reporter is specified.
    """
    return Reporter(sys.stdout, sys.stderr)
<|endoftext|>


<commit_before>from pandasqt.compat import QtCore, QtGui, Qt, Signal, Slot

class OverlayProgressWidget(QtGui.QFrame):
    def __init__(self, parent, workers=[], debug=True, margin=0):
        super(OverlayProgressWidget, self).__init__(parent)
        self._debug = debug
        self._workers = workers
        self._detailProgressBars = []
        self._addedBars = 0
        self._minHeight = 50
        self._width = parent.width() * 0.38
        self._margin = margin
        self._totalProgress = 0

        self.initUi()

        for worker in workers:
            self._addProgressBar(worker)

    def initUi(self):
        self.sizePolicy = QtGui.QSizePolicy(QtGui.QSizePolicy.MinimumExpanding, QtGui.QSizePolicy.Expanding)

        self._pbHeight = 30

        self.setMinimumWidth(self._width)
        #self.setMaximumWidth(self._width)
        self.setMinimumHeight(self._minHeight)

        self.glayout = QtGui.QGridLayout(self)

        self.totalProgressBar = QtGui.QProgressBar(self)
        self.totalProgressBar.setMinimumHeight(self._pbHeight)
        self.totalProgressBar.setMaximumHeight(self._pbHeight)

        self.toggleButton = QtGui.QPushButton('Details', self)
        self.toggleButton.setCheckable(True)
        self.toggleButton.toggled.connect(self.showDetails)
        self.glayout.addWidget(self.totalProgressBar, 0, 0, 1, 1)
        self.glayout.addWidget(self.toggleButton, 0, 1, 1, 1)

        #styleSheet = """.QProgressBar {
            #border: none;
            #border-radius: 3px;
            #text-align: center;
            #background-color: rgba(37, 37, 37, 50%);
            #color: white;
            #margin: 1px;
            #border-bottom-left-radius:5px;
            #border-top-left-radius:5px;
        #}

        #.QProgressBar::chunk {
            #background-color: #05B8CC;
            #border-radius: 3px;
        #}

        #.OverlayProgressWidget {
            #background-color: white;
        #}

        #"""
        ## set stylesheet for all progressbars in this widget
        #self.setStyleSheet(styleSheet)

        parent = self.parent()
        xAnchor = parent.width() - self._width - self._margin
        yAnchor = self._margin
        self.setGeometry(xAnchor, yAnchor, self._width, self._minHeight)

    @Slot(bool)
    def showDetails(self, toggled):
        for (progressBar, label) in self._detailProgressBars:
            progressBar.setVisible(toggled)
            label.setVisible(toggled)
        self.resizeFrame()

    def _addProgressBar(self, worker):
        progressBar = QtGui.QProgressBar(self)
        progressBar.setMinimumHeight(self._pbHeight - 5)
        progressBar.setMaximumHeight(self._pbHeight - 5)
        label = QtGui.QLabel(worker.name, self)
        if not self.toggleButton.isChecked():
            progressBar.hide()
            label.hide()
        row = self._addedBars + 1
        self.glayout.addWidget(progressBar, row, 0, 1, 1)
        self.glayout.addWidget(label, row, 1, 1, 1)
        self._addedBars += 1
        self._detailProgressBars.append((progressBar, label))
        worker.progressChanged.connect(progressBar.setValue)
        worker.progressChanged.connect(self.calculateTotalProgress)

        worker.progressChanged.connect(self.debugProgressChanged)

    def debugProgressChanged(self, value):

    def addWorker(self, worker):
        self._workers.append(worker)
        self._addProgressBar(worker)
        self.resizeFrame()

    def resizeFrame(self):
        size = self.glayout.sizeHint()
        self.resize(size.width(), size.height())

    @Slot()
    def calculateTotalProgress(self):
        bars = len(self._detailProgressBars)
        if bars:
            progress = 0
            for (progressBar, label) in self._detailProgressBars:
                value = progressBar.value()
                progress += value
            progress = progress / bars
        else:
            progress = 100

        self.totalProgressBar.setValue(progress)
<commit_msg>missed one<commit_after>from pandasqt.compat import QtCore, QtGui, Qt, Signal, Slot

class OverlayProgressWidget(QtGui.QFrame):
    def __init__(self, parent, workers=[], debug=True, margin=0):
        super(OverlayProgressWidget, self).__init__(parent)
        self._debug = debug
        self._workers = workers
        self._detailProgressBars = []
        self._addedBars = 0
        self._minHeight = 50
        self._width = parent.width() * 0.38
        self._margin = margin
        self._totalProgress = 0

        self.initUi()

        for worker in workers:
            self._addProgressBar(worker)

    def initUi(self):
        self.sizePolicy = QtGui.QSizePolicy(QtGui.QSizePolicy.MinimumExpanding, QtGui.QSizePolicy.Expanding)

        self._pbHeight = 30

        self.setMinimumWidth(self._width)
        #self.setMaximumWidth(self._width)
        self.setMinimumHeight(self._minHeight)

        self.glayout = QtGui.QGridLayout(self)

        self.totalProgressBar = QtGui.QProgressBar(self)
        self.totalProgressBar.setMinimumHeight(self._pbHeight)
        self.totalProgressBar.setMaximumHeight(self._pbHeight)

        self.toggleButton = QtGui.QPushButton('Details', self)
        self.toggleButton.setCheckable(True)
        self.toggleButton.toggled.connect(self.showDetails)
        self.glayout.addWidget(self.totalProgressBar, 0, 0, 1, 1)
        self.glayout.addWidget(self.toggleButton, 0, 1, 1, 1)

        #styleSheet = """.QProgressBar {
            #border: none;
            #border-radius: 3px;
            #text-align: center;
            #background-color: rgba(37, 37, 37, 50%);
            #color: white;
            #margin: 1px;
            #border-bottom-left-radius:5px;
            #border-top-left-radius:5px;
        #}

        #.QProgressBar::chunk {
            #background-color: #05B8CC;
            #border-radius: 3px;
        #}

        #.OverlayProgressWidget {
            #background-color: white;
        #}

        #"""
        ## set stylesheet for all progressbars in this widget
        #self.setStyleSheet(styleSheet)

        parent = self.parent()
        xAnchor = parent.width() - self._width - self._margin
        yAnchor = self._margin
        self.setGeometry(xAnchor, yAnchor, self._width, self._minHeight)

    @Slot(bool)
    def showDetails(self, toggled):
        for (progressBar, label) in self._detailProgressBars:
            progressBar.setVisible(toggled)
            label.setVisible(toggled)
        self.resizeFrame()

    def _addProgressBar(self, worker):
        progressBar = QtGui.QProgressBar(self)
        progressBar.setMinimumHeight(self._pbHeight - 5)
        progressBar.setMaximumHeight(self._pbHeight - 5)
        label = QtGui.QLabel(worker.name, self)
        if not self.toggleButton.isChecked():
            progressBar.hide()
            label.hide()
        row = self._addedBars + 1
        self.glayout.addWidget(progressBar, row, 0, 1, 1)
        self.glayout.addWidget(label, row, 1, 1, 1)
        self._addedBars += 1
        self._detailProgressBars.append((progressBar, label))
        worker.progressChanged.connect(progressBar.setValue)
        worker.progressChanged.connect(self.calculateTotalProgress)

        worker.progressChanged.connect(self.debugProgressChanged)

    def debugProgressChanged(self, value):
        print(("debugProgressChanged", value))

    def addWorker(self, worker):
        self._workers.append(worker)
        self._addProgressBar(worker)
        self.resizeFrame()

    def resizeFrame(self):
        size = self.glayout.sizeHint()
        self.resize(size.width(), size.height())

    @Slot()
    def calculateTotalProgress(self):
        bars = len(self._detailProgressBars)
        if bars:
            progress = 0
            for (progressBar, label) in self._detailProgressBars:
                value = progressBar.value()
                progress += value
            progress = progress / bars
        else:
            progress = 100

        self.totalProgressBar.setValue(progress)
<|endoftext|>


<commit_before># Copyright (c) 2015, Vienna University of Technology, Department of Geodesy
# and Geoinformation. All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#   * Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the Vienna University of Technology, Department of
#     Geodesy and Geoinformation nor the names of its contributors may be
#     used to endorse or promote products derived from this software without
#     specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL VIENNA UNIVERSITY OF TECHNOLOGY,
# DEPARTMENT OF GEODESY AND GEOINFORMATION BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
import abc
import numpy as np

class StaticBase(object):
    pass

class TsBase(object):

    """
    The TsBase class serves as a template for i/o objects used in
    GriddedTsBase.
    """
    __metaclass__ = abc.ABCMeta

    def __init__(self, filename, mode='r', **kwargs):
        """
        Initialization of i/o object.

        Parameters
        ----------
        filename : str
            File name.
        mode : str, optional
            Opening mode. Default: r
        """
        self.filename = filename
        self.mode = mode
        self.kwargs = kwargs

    @abc.abstractmethod
    def read_ts(self, gpi, **kwargs):
        """
        Read time series data for given grid point.

        Parameters
        ----------
        gpi : int
            Grid point index.

        Returns
        -------
        data : numpy.ndarray
            Data set.
        """
        return

    @abc.abstractmethod
    def write_ts(self, gpi, data, **kwargs):
        """
        Write data.

        Parameters
        ----------
        gpi : int
            Grid point index.
        data : numpy.ndarray
            Data records.
        """
        return

    @abc.abstractmethod
    def flush(self):
        """
        Flush data.
        """
        return

    @abc.abstractmethod
    def close(self):
        """
        Close file.
        """
        return

class ImageBase(object):
    pass

class GriddedStaticBase(object):

    """
    The GriddedStaticBase class uses another IO class together with a grid
    object to read/write a dataset under the given path.

    Parameters
    ----------
    path : string
        Path to dataset.
    grid : pytesmo.grid.grids.BasicGrid of CellGrid instance
        Grid on which the time series data is stored.
    ioclass : class
        IO class.
    mode : str, optional
        File mode and can be read 'r', write 'w' or append 'a'. Default: 'r'
    fn_format : str, optional
        The string format of the cell files. Default: '{:04d}'
    ioclass_kws : dict
        Additional keyword arguments for the ioclass.
    """
    __metaclass__ = abc.ABCMeta

    def __init__(self, path, grid, ioclass, mode='r', fn_format='{:04d}',
                 ioclass_kws=None):

        self.path = path
        self.grid = grid
        self.ioclass = ioclass
        self.mode = mode
        self.fn_format = fn_format
        self.previous_cell = None
        self.fid = None

        if ioclass_kws is None:
            self.ioclass_kws = {}
        else:
            self.ioclass_kws = ioclass_kws

    def __enter__(self):
        """
        Context manager initialization.

        Returns
        -------
        self : GriddedStaticBase object
            self
        """
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """
        Exit the runtime context related to this object. The file will be
        closed. The parameters describe the exception that caused the
        context to be exited.

        exc_type :

        exc_value :

        traceback :

        """
        self.close()

    def _open(self, gpi):
        """
        Open file.

        Parameters
        ----------
        gpi : int
            Grid point index.
        """
        cell = self.grid.gpi2cell(gpi)
        filename = os.path.join(self.path, self.fn_format.format(cell))

        if self.previous_cell != cell:
            if self.fid is not None:
                self.close()

            self.fid = self.ioclass(filename, mode=self.mode,
                                    **self.ioclass_kws)
            self.previous_cell = cell

    def read(self, *args, **kwargs):
        """
        Takes either 1 or 2 arguments and calls the correct function
        which is either reading the gpi directly or finding
        the nearest gpi from given lat,lon coordinates and then reading it
        """
        if len(args) == 1:
            data = self.read_gp(args[0], **kwargs)
        if len(args) == 2:
            data = self._read_lonlat(args[0], args[1], **kwargs)

        return data

    def _read_lonlat(self, lon, lat, **kwargs):
        """
        Reading data for given longitude and latitude coordinate.

        Parameters
        ----------
        lon : float
            Longitude coordinate.
        lat : float
            Latitude coordinate.

        Returns
        -------
        data : dict of values
            data record.
        """
        gp, _ = self.grid.find_nearest_gpi(lon, lat)

        return self.read_gp(gp, **kwargs)

    def read_gp(self, gpi):
        """
        Read data for given grid point.

        Parameters
        ----------
        gpi : int
            Grid point index.

        Returns
        -------
        data : numpy.ndarray
            Time series data.
        """
        self._open(gpi)

        return self.fid.read(gpi)

    def iter_gp(self):
        """
        Yield all values for all grid points.

        Yields
        ------
        data : pandas.DataFrame
            pandas.DateFrame with DateTimeIndex
        """

        gpi_info = list(self.grid.grid_points())
        gps = np.array(gpi_info, dtype=np.int)[:, 0]

        for gp in gps:
            yield self.read_gp(gp), gp

    def write(self, data):
        """
        Write data.

        Parameters
        ----------
        data : numpy.ndarray
            Data records. A field 'gpi', indicating the grid point index
            has to be included.
        """
        self.write_gp(data['gpi'], data)

    def write_gp(self, gpi, data):
        """
        Write data for given grid point.

        Parameters
        ----------
        gpi : int
            Grid point index.
        data : numpy.ndarray
            Data
        """
        self._open(gpi)
        self.fid.write(data)

    def flush(self):
        """
        Flush data.
        """
        if self.fid is not None:
            self.fid.flush()

    def close(self):
        """
        Close file.
        """
        if self.fid is not None:
            self.fid.close()

class GriddedTsBase(object):

    """
    The GriddedTsBase class uses another IO class together with a grid object
    to read/write a time series dataset under the given path.

    Parameters
    ----------
    path : string
        Path to dataset.
    grid : pytesmo.grid.grids.BasicGrid of CellGrid instance
        Grid on which the time series data is stored.
    ioclass : class
        IO class.
    mode : str, optional
        File mode and can be read 'r', write 'w' or append 'a'. Default: 'r'
    fn_format : str, optional
        The string format of the cell files. Default: '{:04d}'
    ioclass_kws : dict
        Additional keyword arguments for the ioclass.
    """
    __metaclass__ = abc.ABCMeta

    def __init__(self, path, grid, ioclass, mode='r', fn_format='{:04d}',
                 ioclass_kws=None):

        self.path = path
        self.grid = grid
        self.ioclass = ioclass
        self.mode = mode
        self.fn_format = fn_format
        self.previous_cell = None
        self.fid = None

        if ioclass_kws is None:
            self.ioclass_kws = {}
        else:
            self.ioclass_kws = ioclass_kws

    def __enter__(self):
        """
        Context manager initialization.

        Returns
        -------
        self : GriddedBaseTs object
            self
        """
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """
        Exit the runtime context related to this object. The file will be
        closed. The parameters describe the exception that caused the
        context to be exited.

        exc_type :

        exc_value :

        traceback :

        """
        self.close()

    def _open(self, gpi):
        """
        Open file.

        Parameters
        ----------
        gpi : int
            Grid point index.
        """
        cell = self.grid.gpi2cell(gpi)
        filename = os.path.join(self.path, self.fn_format.format(cell))

        if self.previous_cell != cell:
            self.close()
            self.fid = self.ioclass(filename, mode=self.mode,
                                    **self.ioclass_kws)
            self.previous_cell = cell

    def _read_lonlat(self, lon, lat, **kwargs):
        """
        Reading time series for given longitude and latitude coordinate.

        Parameters
        ----------
        lon : float
            Longitude coordinate.
        lat : float
            Latitude coordinate.

        Returns
        -------
        data : pandas.DataFrame
            pandas.DateFrame with DateTimeIndex.
        """
        gp, _ = self.grid.find_nearest_gpi(lon, lat)

        return self.read_gp(gp, **kwargs)

    def _write_lonlat(self, lon, lat, data, **kwargs):
        """
        Write time series for given longitude and latitude coordinate.

        Parameters
        ----------
        lon : float
            Longitude coordinate.
        lat : float
            Latitude coordinate.
        data : numpy.ndarray
            Data records.
        """
        gp, _ = self.grid.find_nearest_gpi(lon, lat)

        return self.write_gp(gp, **kwargs)

    def get_nearest_gp_info(self, lon, lat):
        """
        get info for nearest grid point

        Parameters
        ----------
        lon : float
            Longitude coordinate.
        lat : float
            Latitude coordinate.

        Returns
        -------
        gpi : int
            Grid point index of nearest grid point.
        gp_lon : float
            Lontitude coordinate of nearest grid point.
        gp_lat : float
            Latitude coordinate of nearest grid point.
        gp_dist : float
            Geodetic distance to nearest grid point.
        """
        gpi, gp_dist = self.grid.find_nearest_gpi(lon, lat)
        gp_lon, gp_lat = self.grid.gpi2lonlat(gpi)

        return gpi, gp_lon, gp_lat, gp_dist

    def write_gp(self, gpi, data, **kwargs):
        """
        Write data for given grid point.

        Parameters
        ----------
        gpi : int
            Grid point index.
        data : numpy.ndarray
            Data records.
        """
        self._open(gpi)

        if self.mode in ['r']:
            raise IOError("File is not open in write/append mode")

        lon, lat = self.grid.gpi2lonlat(gpi)

        self.fid.write_ts(gpi, data, lon=lon, lat=lat, **kwargs)

    def write_ts(self, *args, **kwargs):
        """
        Takes either 2 or 3 arguments (the last one always needs to be the
        data to be written) and calls the correct function which is either
        writing the gpi directly or finding the nearest gpi from given
        lon, lat coordinates and then reading it.
        """
        if len(args) == 2:
            self.write_gp(args[0], args[1], **kwargs)
        if len(args) == 3:
            self._write_lonlat(args[0], args[1], args[2], **kwargs)
        if len(args) < 2 or len(args) > 3:
            raise ValueError("Wrong number of arguments")

    def read_gp(self, gpi, **kwargs):
        """
        Reads time series for a given grid point index.

        Parameters
        ----------
        gpi : int
            grid point index

        Returns
        -------
        data : pandas.DataFrame
            pandas.DateFrame with DateTimeIndex
        """
        self._open(gpi)

        if self.mode in ['w', 'a']:
            raise IOError("File is not open in read mode")

        return self.fid.read_ts(gpi, **kwargs)

    def read_ts(self, *args, **kwargs):
        """
        Takes either 1 or 2 arguments and calls the correct function
        which is either reading the gpi directly or finding
        the nearest gpi from given lat,lon coordinates and then reading it
        """
        if len(args) == 1:
            data = self.read_gp(args[0], **kwargs)
        if len(args) == 2:
            data = self._read_lonlat(args[0], args[1], **kwargs)
        if len(args) < 1 or len(args) > 2:
            raise ValueError("Wrong number of arguments")

        return data

    def iter_ts(self):
        """
        Yield time series for all grid points.

        Yields
        ------
        data : pandas.DataFrame
            pandas.DateFrame with DateTimeIndex
        gpi : int
            Grid point index
        """
        gpi_info = list(self.grid.grid_points())
        gps = np.array(gpi_info, dtype=np.int)[:, 0]

        for gp in gps:
            yield self.read_gp(gp), gp

    def flush(self):
        """
        Flush data.
        """
        if self.fid is not None:
            self.fid.flush()

    def close(self):
        """
        Close file.
        """
        if self.fid is not None:
            self.fid.close()
            self.fid = None
<commit_msg>add implementation of StaticBase<commit_after># Copyright (c) 2015, Vienna University of Technology, Department of Geodesy
# and Geoinformation. All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#   * Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the Vienna University of Technology, Department of
#     Geodesy and Geoinformation nor the names of its contributors may be
#     used to endorse or promote products derived from this software without
#     specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL VIENNA UNIVERSITY OF TECHNOLOGY,
# DEPARTMENT OF GEODESY AND GEOINFORMATION BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
import abc
import numpy as np

class StaticBase(object):

    """
    The StaticBase class serves as a template for i/o objects used in
    GriddedStaticBase.
    """
    __metaclass__ = abc.ABCMeta

    def __init__(self, filename, mode='r', **kwargs):
        """
        Initialization of i/o object.

        Parameters
        ----------
        filename : str
            File name.
        mode : str, optional
            Opening mode. Default: r
        """
        self.filename = filename
        self.mode = mode
        self.kwargs = kwargs

    @abc.abstractmethod
    def read(self, gpi):
        """
        Read data for given grid point.

        Parameters
        ----------
        gpi : int
            Grid point index.

        Returns
        -------
        data : numpy.ndarray
            Data set.
        """
        return

    @abc.abstractmethod
    def write(self, data):
        """
        Write data.

        Parameters
        ----------
        data : numpy.ndarray
            Data records.
        """
        return

    @abc.abstractmethod
    def flush(self):
        """
        Flush data.
        """
        return

    @abc.abstractmethod
    def close(self):
        """
        Close file.
        """
        return

class TsBase(object):

    """
    The TsBase class serves as a template for i/o objects used in
    GriddedTsBase.
    """
    __metaclass__ = abc.ABCMeta

    def __init__(self, filename, mode='r', **kwargs):
        """
        Initialization of i/o object.

        Parameters
        ----------
        filename : str
            File name.
        mode : str, optional
            Opening mode. Default: r
        """
        self.filename = filename
        self.mode = mode
        self.kwargs = kwargs

    @abc.abstractmethod
    def read_ts(self, gpi, **kwargs):
        """
        Read time series data for given grid point.

        Parameters
        ----------
        gpi : int
            Grid point index.

        Returns
        -------
        data : numpy.ndarray
            Data set.
        """
        return

    @abc.abstractmethod
    def write_ts(self, gpi, data, **kwargs):
        """
        Write data.

        Parameters
        ----------
        gpi : int
            Grid point index.
        data : numpy.ndarray
            Data records.
        """
        return

    @abc.abstractmethod
    def flush(self):
        """
        Flush data.
        """
        return

    @abc.abstractmethod
    def close(self):
        """
        Close file.
        """
        return

class ImageBase(object):
    pass

class GriddedStaticBase(object):

    """
    The GriddedStaticBase class uses another IO class together with a grid
    object to read/write a dataset under the given path.

    Parameters
    ----------
    path : string
        Path to dataset.
    grid : pytesmo.grid.grids.BasicGrid of CellGrid instance
        Grid on which the time series data is stored.
    ioclass : class
        IO class.
    mode : str, optional
        File mode and can be read 'r', write 'w' or append 'a'. Default: 'r'
    fn_format : str, optional
        The string format of the cell files. Default: '{:04d}'
    ioclass_kws : dict
        Additional keyword arguments for the ioclass.
    """
    __metaclass__ = abc.ABCMeta

    def __init__(self, path, grid, ioclass, mode='r', fn_format='{:04d}',
                 ioclass_kws=None):

        self.path = path
        self.grid = grid
        self.ioclass = ioclass
        self.mode = mode
        self.fn_format = fn_format
        self.previous_cell = None
        self.fid = None

        if ioclass_kws is None:
            self.ioclass_kws = {}
        else:
            self.ioclass_kws = ioclass_kws

    def __enter__(self):
        """
        Context manager initialization.

        Returns
        -------
        self : GriddedStaticBase object
            self
        """
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """
        Exit the runtime context related to this object. The file will be
        closed. The parameters describe the exception that caused the
        context to be exited.

        exc_type :

        exc_value :

        traceback :

        """
        self.close()

    def _open(self, gpi):
        """
        Open file.

        Parameters
        ----------
        gpi : int
            Grid point index.
        """
        cell = self.grid.gpi2cell(gpi)
        filename = os.path.join(self.path, self.fn_format.format(cell))

        if self.previous_cell != cell:
            if self.fid is not None:
                self.close()

            self.fid = self.ioclass(filename, mode=self.mode,
                                    **self.ioclass_kws)
            self.previous_cell = cell

    def read(self, *args, **kwargs):
        """
        Takes either 1 or 2 arguments and calls the correct function
        which is either reading the gpi directly or finding
        the nearest gpi from given lat,lon coordinates and then reading it
        """
        if len(args) == 1:
            data = self.read_gp(args[0], **kwargs)
        if len(args) == 2:
            data = self._read_lonlat(args[0], args[1], **kwargs)

        return data

    def _read_lonlat(self, lon, lat, **kwargs):
        """
        Reading data for given longitude and latitude coordinate.

        Parameters
        ----------
        lon : float
            Longitude coordinate.
        lat : float
            Latitude coordinate.

        Returns
        -------
        data : dict of values
            data record.
        """
        gp, _ = self.grid.find_nearest_gpi(lon, lat)

        return self.read_gp(gp, **kwargs)

    def read_gp(self, gpi):
        """
        Read data for given grid point.

        Parameters
        ----------
        gpi : int
            Grid point index.

        Returns
        -------
        data : numpy.ndarray
            Data set.
        """
        self._open(gpi)

        return self.fid.read(gpi)

    def iter_gp(self):
        """
        Yield all values for all grid points.

        Yields
        ------
        data : pandas.DataFrame
            pandas.DateFrame with DateTimeIndex
        """
        gpi_info = list(self.grid.grid_points())
        gps = np.array(gpi_info, dtype=np.int)[:, 0]

        for gp in gps:
            yield self.read_gp(gp), gp

    def write(self, data):
        """
        Write data.

        Parameters
        ----------
        data : numpy.ndarray
            Data records. A field 'gpi', indicating the grid point index
            has to be included.
        """
        self.write_gp(data['gpi'], data)

    def write_gp(self, gpi, data):
        """
        Write data for given grid point.

        Parameters
        ----------
        gpi : int
            Grid point index.
        data : numpy.ndarray
            Data
        """
        self._open(gpi)
        self.fid.write(data)

    def flush(self):
        """
        Flush data.
        """
        if self.fid is not None:
            self.fid.flush()

    def close(self):
        """
        Close file.
        """
        if self.fid is not None:
            self.fid.close()

class GriddedTsBase(object):

    """
    The GriddedTsBase class uses another IO class together with a grid object
    to read/write a time series dataset under the given path.

    Parameters
    ----------
    path : string
        Path to dataset.
    grid : pytesmo.grid.grids.BasicGrid of CellGrid instance
        Grid on which the time series data is stored.
    ioclass : class
        IO class.
    mode : str, optional
        File mode and can be read 'r', write 'w' or append 'a'. Default: 'r'
    fn_format : str, optional
        The string format of the cell files. Default: '{:04d}'
    ioclass_kws : dict
        Additional keyword arguments for the ioclass.
    """
    __metaclass__ = abc.ABCMeta

    def __init__(self, path, grid, ioclass, mode='r', fn_format='{:04d}',
                 ioclass_kws=None):

        self.path = path
        self.grid = grid
        self.ioclass = ioclass
        self.mode = mode
        self.fn_format = fn_format
        self.previous_cell = None
        self.fid = None

        if ioclass_kws is None:
            self.ioclass_kws = {}
        else:
            self.ioclass_kws = ioclass_kws

    def __enter__(self):
        """
        Context manager initialization.

        Returns
        -------
        self : GriddedBaseTs object
            self
        """
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """
        Exit the runtime context related to this object. The file will be
        closed. The parameters describe the exception that caused the
        context to be exited.

        exc_type :

        exc_value :

        traceback :

        """
        self.close()

    def _open(self, gpi):
        """
        Open file.

        Parameters
        ----------
        gpi : int
            Grid point index.
        """
        cell = self.grid.gpi2cell(gpi)
        filename = os.path.join(self.path, self.fn_format.format(cell))

        if self.previous_cell != cell:
            self.close()
            self.fid = self.ioclass(filename, mode=self.mode,
                                    **self.ioclass_kws)
            self.previous_cell = cell

    def _read_lonlat(self, lon, lat, **kwargs):
        """
        Reading time series for given longitude and latitude coordinate.

        Parameters
        ----------
        lon : float
            Longitude coordinate.
        lat : float
            Latitude coordinate.

        Returns
        -------
        data : pandas.DataFrame
            pandas.DateFrame with DateTimeIndex.
        """
        gp, _ = self.grid.find_nearest_gpi(lon, lat)

        return self.read_gp(gp, **kwargs)

    def _write_lonlat(self, lon, lat, data, **kwargs):
        """
        Write time series for given longitude and latitude coordinate.

        Parameters
        ----------
        lon : float
            Longitude coordinate.
        lat : float
            Latitude coordinate.
        data : numpy.ndarray
            Data records.
        """
        gp, _ = self.grid.find_nearest_gpi(lon, lat)

        return self.write_gp(gp, **kwargs)

    def get_nearest_gp_info(self, lon, lat):
        """
        get info for nearest grid point

        Parameters
        ----------
        lon : float
            Longitude coordinate.
        lat : float
            Latitude coordinate.

        Returns
        -------
        gpi : int
            Grid point index of nearest grid point.
        gp_lon : float
            Lontitude coordinate of nearest grid point.
        gp_lat : float
            Latitude coordinate of nearest grid point.
        gp_dist : float
            Geodetic distance to nearest grid point.
        """
        gpi, gp_dist = self.grid.find_nearest_gpi(lon, lat)
        gp_lon, gp_lat = self.grid.gpi2lonlat(gpi)

        return gpi, gp_lon, gp_lat, gp_dist

    def write_gp(self, gpi, data, **kwargs):
        """
        Write data for given grid point.

        Parameters
        ----------
        gpi : int
            Grid point index.
        data : numpy.ndarray
            Data records.
        """
        self._open(gpi)

        if self.mode in ['r']:
            raise IOError("File is not open in write/append mode")

        lon, lat = self.grid.gpi2lonlat(gpi)

        self.fid.write_ts(gpi, data, lon=lon, lat=lat, **kwargs)

    def write_ts(self, *args, **kwargs):
        """
        Takes either 2 or 3 arguments (the last one always needs to be the
        data to be written) and calls the correct function which is either
        writing the gpi directly or finding the nearest gpi from given
        lon, lat coordinates and then reading it.
        """
        if len(args) == 2:
            self.write_gp(args[0], args[1], **kwargs)
        if len(args) == 3:
            self._write_lonlat(args[0], args[1], args[2], **kwargs)
        if len(args) < 2 or len(args) > 3:
            raise ValueError("Wrong number of arguments")

    def read_gp(self, gpi, **kwargs):
        """
        Reads time series for a given grid point index.

        Parameters
        ----------
        gpi : int
            grid point index

        Returns
        -------
        data : pandas.DataFrame
            pandas.DateFrame with DateTimeIndex
        """
        self._open(gpi)

        if self.mode in ['w', 'a']:
            raise IOError("File is not open in read mode")

        return self.fid.read_ts(gpi, **kwargs)

    def read_ts(self, *args, **kwargs):
        """
        Takes either 1 or 2 arguments and calls the correct function
        which is either reading the gpi directly or finding
        the nearest gpi from given lat,lon coordinates and then reading it
        """
        if len(args) == 1:
            data = self.read_gp(args[0], **kwargs)
        if len(args) == 2:
            data = self._read_lonlat(args[0], args[1], **kwargs)
        if len(args) < 1 or len(args) > 2:
            raise ValueError("Wrong number of arguments")

        return data

    def iter_ts(self):
        """
        Yield time series for all grid points.

        Yields
        ------
        data : pandas.DataFrame
            pandas.DateFrame with DateTimeIndex
        gpi : int
            Grid point index
        """
        gpi_info = list(self.grid.grid_points())
        gps = np.array(gpi_info, dtype=np.int)[:, 0]

        for gp in gps:
            yield self.read_gp(gp), gp

    def flush(self):
        """
        Flush data.
        """
        if self.fid is not None:
            self.fid.flush()

    def close(self):
        """
        Close file.
        """
        if self.fid is not None:
            self.fid.close()
            self.fid = None
<|endoftext|>