<commit_before>#tests the behavior of the column classes
from datetime import datetime, timedelta
from datetime import date
from datetime import tzinfo
from decimal import Decimal as D
from unittest import TestCase
from uuid import uuid4, uuid1
from cassandra import InvalidRequest
from cqlengine import ValidationError
from cqlengine.connection import execute

from cqlengine.tests.base import BaseCassEngTestCase

from cqlengine.columns import Column, TimeUUID
from cqlengine.columns import Bytes
from cqlengine.columns import Ascii
from cqlengine.columns import Text
from cqlengine.columns import Integer
from cqlengine.columns import BigInt
from cqlengine.columns import VarInt
from cqlengine.columns import DateTime
from cqlengine.columns import Date
from cqlengine.columns import UUID
from cqlengine.columns import Boolean
from cqlengine.columns import Float
from cqlengine.columns import Decimal
from cqlengine.columns import Inet

from cqlengine.management import sync_table, drop_table
from cqlengine.models import Model

import sys

class TestDatetime(BaseCassEngTestCase):
    class DatetimeTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        created_at = DateTime()

    @classmethod
    def setUpClass(cls):
        super(TestDatetime, cls).setUpClass()
        sync_table(cls.DatetimeTest)

    @classmethod
    def tearDownClass(cls):
        super(TestDatetime, cls).tearDownClass()
        drop_table(cls.DatetimeTest)

    def test_datetime_io(self):
        now = datetime.now()
        dt = self.DatetimeTest.objects.create(test_id=0, created_at=now)
        dt2 = self.DatetimeTest.objects(test_id=0).first()
        assert dt2.created_at.timetuple()[:6] == now.timetuple()[:6]

    def test_datetime_tzinfo_io(self):
        class TZ(tzinfo):
            def utcoffset(self, date_time):
                return timedelta(hours=-1)
            def dst(self, date_time):
                return None

        now = datetime(1982, 1, 1, tzinfo=TZ())
        dt = self.DatetimeTest.objects.create(test_id=0, created_at=now)
        dt2 = self.DatetimeTest.objects(test_id=0).first()
        assert dt2.created_at.timetuple()[:6] == (now + timedelta(hours=1)).timetuple()[:6]

    def test_datetime_date_support(self):
        today = date.today()
        self.DatetimeTest.objects.create(test_id=0, created_at=today)
        dt2 = self.DatetimeTest.objects(test_id=0).first()
        assert dt2.created_at.isoformat() == datetime(today.year, today.month, today.day).isoformat()

    def test_datetime_none(self):
        dt = self.DatetimeTest.objects.create(test_id=1, created_at=None)
        dt2 = self.DatetimeTest.objects(test_id=1).first()
        assert dt2.created_at is None

        dts = self.DatetimeTest.objects.filter(test_id=1).values_list('created_at')
        assert dts[0][0] is None

class TestBoolDefault(BaseCassEngTestCase):
    class BoolDefaultValueTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        stuff = Boolean(default=True)

    @classmethod
    def setUpClass(cls):
        super(TestBoolDefault, cls).setUpClass()
        sync_table(cls.BoolDefaultValueTest)

    def test_default_is_set(self):
        tmp = self.BoolDefaultValueTest.create(test_id=1)
        self.assertEqual(True, tmp.stuff)
        tmp2 = self.BoolDefaultValueTest.get(test_id=1)
        self.assertEqual(True, tmp2.stuff)

class TestVarInt(BaseCassEngTestCase):
    class VarIntTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        bignum = VarInt(primary_key=True)

    @classmethod
    def setUpClass(cls):
        super(TestVarInt, cls).setUpClass()
        sync_table(cls.VarIntTest)

    @classmethod
    def tearDownClass(cls):
        super(TestVarInt, cls).tearDownClass()
        sync_table(cls.VarIntTest)

    def test_varint_io(self):
        long_int = sys.maxint + 1
        int1 = self.VarIntTest.objects.create(test_id=0, bignum=long_int)
        int2 = self.VarIntTest.objects(test_id=0).first()
        assert int1.bignum == int2.bignum

class TestDate(BaseCassEngTestCase):
    class DateTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        created_at = Date()

    @classmethod
    def setUpClass(cls):
        super(TestDate, cls).setUpClass()
        sync_table(cls.DateTest)

    @classmethod
    def tearDownClass(cls):
        super(TestDate, cls).tearDownClass()
        drop_table(cls.DateTest)

    def test_date_io(self):
        today = date.today()
        self.DateTest.objects.create(test_id=0, created_at=today)
        dt2 = self.DateTest.objects(test_id=0).first()
        assert dt2.created_at.isoformat() == today.isoformat()

    def test_date_io_using_datetime(self):
        now = datetime.utcnow()
        self.DateTest.objects.create(test_id=0, created_at=now)
        dt2 = self.DateTest.objects(test_id=0).first()
        assert not isinstance(dt2.created_at, datetime)
        assert isinstance(dt2.created_at, date)
        assert dt2.created_at.isoformat() == now.date().isoformat()

    def test_date_none(self):
        self.DateTest.objects.create(test_id=1, created_at=None)
        dt2 = self.DateTest.objects(test_id=1).first()
        assert dt2.created_at is None

        dts = self.DateTest.objects(test_id=1).values_list('created_at')
        assert dts[0][0] is None

class TestDecimal(BaseCassEngTestCase):
    class DecimalTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        dec_val = Decimal()

    @classmethod
    def setUpClass(cls):
        super(TestDecimal, cls).setUpClass()
        sync_table(cls.DecimalTest)

    @classmethod
    def tearDownClass(cls):
        super(TestDecimal, cls).tearDownClass()
        drop_table(cls.DecimalTest)

    def test_decimal_io(self):
        dt = self.DecimalTest.objects.create(test_id=0, dec_val=D('0.00'))
        dt2 = self.DecimalTest.objects(test_id=0).first()
        assert dt2.dec_val == dt.dec_val

        dt = self.DecimalTest.objects.create(test_id=0, dec_val=5)
        dt2 = self.DecimalTest.objects(test_id=0).first()
        assert dt2.dec_val == D('5')

class TestUUID(BaseCassEngTestCase):
    class UUIDTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        a_uuid = UUID(default=uuid4())

    @classmethod
    def setUpClass(cls):
        super(TestUUID, cls).setUpClass()
        sync_table(cls.UUIDTest)

    @classmethod
    def tearDownClass(cls):
        super(TestUUID, cls).tearDownClass()
        drop_table(cls.UUIDTest)

    def test_uuid_str_with_dashes(self):
        a_uuid = uuid4()
        t0 = self.UUIDTest.create(test_id=0, a_uuid=str(a_uuid))
        t1 = self.UUIDTest.get(test_id=0)
        assert a_uuid == t1.a_uuid

    def test_uuid_str_no_dashes(self):
        a_uuid = uuid4()
        t0 = self.UUIDTest.create(test_id=1, a_uuid=a_uuid.hex)
        t1 = self.UUIDTest.get(test_id=1)
        assert a_uuid == t1.a_uuid

class TestTimeUUID(BaseCassEngTestCase):
    class TimeUUIDTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        timeuuid = TimeUUID(default=uuid1())

    @classmethod
    def setUpClass(cls):
        super(TestTimeUUID, cls).setUpClass()
        sync_table(cls.TimeUUIDTest)

    @classmethod
    def tearDownClass(cls):
        super(TestTimeUUID, cls).tearDownClass()
        drop_table(cls.TimeUUIDTest)

    def test_timeuuid_io(self):
        """
        ensures that
        :return:
        """
        t0 = self.TimeUUIDTest.create(test_id=0)
        t1 = self.TimeUUIDTest.get(test_id=0)

        assert t1.timeuuid.time == t1.timeuuid.time

class TestInteger(BaseCassEngTestCase):
    class IntegerTest(Model):
        __keyspace__ = 'test'
        test_id = UUID(primary_key=True, default=lambda:uuid4())
        value   = Integer(default=0, required=True)

    def test_default_zero_fields_validate(self):
        """ Tests that integer columns with a default value of 0 validate """
        it = self.IntegerTest()
        it.validate()

class TestBigInt(BaseCassEngTestCase):
    class BigIntTest(Model):
        __keyspace__ = 'test'
        test_id = UUID(primary_key=True, default=lambda:uuid4())
        value   = BigInt(default=0, required=True)

    def test_default_zero_fields_validate(self):
        """ Tests that bigint columns with a default value of 0 validate """
        it = self.BigIntTest()
        it.validate()

class TestText(BaseCassEngTestCase):

    def test_min_length(self):
        #min len defaults to 1
        col = Text()
        col.validate('')

        col.validate('b')

        #test not required defaults to 0
        Text(required=False).validate('')

        #test arbitrary lengths
        Text(min_length=0).validate('')
        Text(min_length=5).validate('blake')
        Text(min_length=5).validate('blaketastic')
        with self.assertRaises(ValidationError):
            Text(min_length=6).validate('blake')

    def test_max_length(self):

        Text(max_length=5).validate('blake')
        with self.assertRaises(ValidationError):
            Text(max_length=5).validate('blaketastic')

    def test_type_checking(self):
        Text().validate('string')
        Text().validate(u'unicode')
        Text().validate(bytearray('bytearray'))

        with self.assertRaises(ValidationError):
            Text(required=True).validate(None)

        with self.assertRaises(ValidationError):
            Text().validate(5)

        with self.assertRaises(ValidationError):
            Text().validate(True)

    def test_non_required_validation(self):
        """ Tests that validation is ok on none and blank values if required is False """
        Text().validate('')
        Text().validate(None)

class TestExtraFieldsRaiseException(BaseCassEngTestCase):
    class TestModel(Model):
        __keyspace__ = 'test'
        id = UUID(primary_key=True, default=uuid4)

    def test_extra_field(self):
        with self.assertRaises(ValidationError):
            self.TestModel.create(bacon=5000)

class TestPythonDoesntDieWhenExtraFieldIsInCassandra(BaseCassEngTestCase):
    class TestModel(Model):
        __keyspace__ = 'test'
        __table_name__ = 'alter_doesnt_break_running_app'
        id = UUID(primary_key=True, default=uuid4)

    def test_extra_field(self):
        drop_table(self.TestModel)
        sync_table(self.TestModel)
        self.TestModel.create()
        execute("ALTER TABLE {} add blah int".format(self.TestModel.column_family_name(include_keyspace=True)))
        self.TestModel.objects().all()

class TestTimeUUIDFromDatetime(TestCase):
    def test_conversion_specific_date(self):
        dt = datetime(1981, 7, 11, microsecond=555000)

        uuid = TimeUUID.from_datetime(dt)

        from uuid import UUID
        assert isinstance(uuid, UUID)

        ts = (uuid.time - 0x01b21dd213814000) / 1e7 # back to a timestamp
        new_dt = datetime.utcfromtimestamp(ts)

        # checks that we created a UUID1 with the proper timestamp
        assert new_dt == dt

class TestInet(BaseCassEngTestCase):

    class InetTestModel(Model):
        id = UUID(primary_key=True, default=uuid4)
        address = Inet()

    def setUp(self):
        drop_table(self.InetTestModel)
        sync_table(self.InetTestModel)

    def test_inet_saves(self):
        tmp = self.InetTestModel.create(address="192.168.1.1")

        m = self.InetTestModel.get(id=tmp.id)

        assert m.address == "192.168.1.1"

    def test_non_address_fails(self):
        with self.assertRaises(InvalidRequest):
            self.InetTestModel.create(address="ham sandwich")

<commit_msg>removal of maxint in python 3<commit_after>#tests the behavior of the column classes
from datetime import datetime, timedelta
from datetime import date
from datetime import tzinfo
from decimal import Decimal as D
from unittest import TestCase
from uuid import uuid4, uuid1
from cassandra import InvalidRequest
import six
from cqlengine import ValidationError
from cqlengine.connection import execute

from cqlengine.tests.base import BaseCassEngTestCase

from cqlengine.columns import Column, TimeUUID
from cqlengine.columns import Bytes
from cqlengine.columns import Ascii
from cqlengine.columns import Text
from cqlengine.columns import Integer
from cqlengine.columns import BigInt
from cqlengine.columns import VarInt
from cqlengine.columns import DateTime
from cqlengine.columns import Date
from cqlengine.columns import UUID
from cqlengine.columns import Boolean
from cqlengine.columns import Float
from cqlengine.columns import Decimal
from cqlengine.columns import Inet

from cqlengine.management import sync_table, drop_table
from cqlengine.models import Model

import sys

class TestDatetime(BaseCassEngTestCase):
    class DatetimeTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        created_at = DateTime()

    @classmethod
    def setUpClass(cls):
        super(TestDatetime, cls).setUpClass()
        sync_table(cls.DatetimeTest)

    @classmethod
    def tearDownClass(cls):
        super(TestDatetime, cls).tearDownClass()
        drop_table(cls.DatetimeTest)

    def test_datetime_io(self):
        now = datetime.now()
        dt = self.DatetimeTest.objects.create(test_id=0, created_at=now)
        dt2 = self.DatetimeTest.objects(test_id=0).first()
        assert dt2.created_at.timetuple()[:6] == now.timetuple()[:6]

    def test_datetime_tzinfo_io(self):
        class TZ(tzinfo):
            def utcoffset(self, date_time):
                return timedelta(hours=-1)
            def dst(self, date_time):
                return None

        now = datetime(1982, 1, 1, tzinfo=TZ())
        dt = self.DatetimeTest.objects.create(test_id=0, created_at=now)
        dt2 = self.DatetimeTest.objects(test_id=0).first()
        assert dt2.created_at.timetuple()[:6] == (now + timedelta(hours=1)).timetuple()[:6]

    def test_datetime_date_support(self):
        today = date.today()
        self.DatetimeTest.objects.create(test_id=0, created_at=today)
        dt2 = self.DatetimeTest.objects(test_id=0).first()
        assert dt2.created_at.isoformat() == datetime(today.year, today.month, today.day).isoformat()

    def test_datetime_none(self):
        dt = self.DatetimeTest.objects.create(test_id=1, created_at=None)
        dt2 = self.DatetimeTest.objects(test_id=1).first()
        assert dt2.created_at is None

        dts = self.DatetimeTest.objects.filter(test_id=1).values_list('created_at')
        assert dts[0][0] is None

class TestBoolDefault(BaseCassEngTestCase):
    class BoolDefaultValueTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        stuff = Boolean(default=True)

    @classmethod
    def setUpClass(cls):
        super(TestBoolDefault, cls).setUpClass()
        sync_table(cls.BoolDefaultValueTest)

    def test_default_is_set(self):
        tmp = self.BoolDefaultValueTest.create(test_id=1)
        self.assertEqual(True, tmp.stuff)
        tmp2 = self.BoolDefaultValueTest.get(test_id=1)
        self.assertEqual(True, tmp2.stuff)

class TestVarInt(BaseCassEngTestCase):
    class VarIntTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        bignum = VarInt(primary_key=True)

    @classmethod
    def setUpClass(cls):
        super(TestVarInt, cls).setUpClass()
        sync_table(cls.VarIntTest)

    @classmethod
    def tearDownClass(cls):
        super(TestVarInt, cls).tearDownClass()
        sync_table(cls.VarIntTest)

    def test_varint_io(self):
        long_int = 92834902384092834092384028340283048239048203480234823048230482304820348239
        int1 = self.VarIntTest.objects.create(test_id=0, bignum=long_int)
        int2 = self.VarIntTest.objects(test_id=0).first()
        assert int1.bignum == int2.bignum

class TestDate(BaseCassEngTestCase):
    class DateTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        created_at = Date()

    @classmethod
    def setUpClass(cls):
        super(TestDate, cls).setUpClass()
        sync_table(cls.DateTest)

    @classmethod
    def tearDownClass(cls):
        super(TestDate, cls).tearDownClass()
        drop_table(cls.DateTest)

    def test_date_io(self):
        today = date.today()
        self.DateTest.objects.create(test_id=0, created_at=today)
        dt2 = self.DateTest.objects(test_id=0).first()
        assert dt2.created_at.isoformat() == today.isoformat()

    def test_date_io_using_datetime(self):
        now = datetime.utcnow()
        self.DateTest.objects.create(test_id=0, created_at=now)
        dt2 = self.DateTest.objects(test_id=0).first()
        assert not isinstance(dt2.created_at, datetime)
        assert isinstance(dt2.created_at, date)
        assert dt2.created_at.isoformat() == now.date().isoformat()

    def test_date_none(self):
        self.DateTest.objects.create(test_id=1, created_at=None)
        dt2 = self.DateTest.objects(test_id=1).first()
        assert dt2.created_at is None

        dts = self.DateTest.objects(test_id=1).values_list('created_at')
        assert dts[0][0] is None

class TestDecimal(BaseCassEngTestCase):
    class DecimalTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        dec_val = Decimal()

    @classmethod
    def setUpClass(cls):
        super(TestDecimal, cls).setUpClass()
        sync_table(cls.DecimalTest)

    @classmethod
    def tearDownClass(cls):
        super(TestDecimal, cls).tearDownClass()
        drop_table(cls.DecimalTest)

    def test_decimal_io(self):
        dt = self.DecimalTest.objects.create(test_id=0, dec_val=D('0.00'))
        dt2 = self.DecimalTest.objects(test_id=0).first()
        assert dt2.dec_val == dt.dec_val

        dt = self.DecimalTest.objects.create(test_id=0, dec_val=5)
        dt2 = self.DecimalTest.objects(test_id=0).first()
        assert dt2.dec_val == D('5')

class TestUUID(BaseCassEngTestCase):
    class UUIDTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        a_uuid = UUID(default=uuid4())

    @classmethod
    def setUpClass(cls):
        super(TestUUID, cls).setUpClass()
        sync_table(cls.UUIDTest)

    @classmethod
    def tearDownClass(cls):
        super(TestUUID, cls).tearDownClass()
        drop_table(cls.UUIDTest)

    def test_uuid_str_with_dashes(self):
        a_uuid = uuid4()
        t0 = self.UUIDTest.create(test_id=0, a_uuid=str(a_uuid))
        t1 = self.UUIDTest.get(test_id=0)
        assert a_uuid == t1.a_uuid

    def test_uuid_str_no_dashes(self):
        a_uuid = uuid4()
        t0 = self.UUIDTest.create(test_id=1, a_uuid=a_uuid.hex)
        t1 = self.UUIDTest.get(test_id=1)
        assert a_uuid == t1.a_uuid

class TestTimeUUID(BaseCassEngTestCase):
    class TimeUUIDTest(Model):
        __keyspace__ = 'test'
        test_id = Integer(primary_key=True)
        timeuuid = TimeUUID(default=uuid1())

    @classmethod
    def setUpClass(cls):
        super(TestTimeUUID, cls).setUpClass()
        sync_table(cls.TimeUUIDTest)

    @classmethod
    def tearDownClass(cls):
        super(TestTimeUUID, cls).tearDownClass()
        drop_table(cls.TimeUUIDTest)

    def test_timeuuid_io(self):
        """
        ensures that
        :return:
        """
        t0 = self.TimeUUIDTest.create(test_id=0)
        t1 = self.TimeUUIDTest.get(test_id=0)

        assert t1.timeuuid.time == t1.timeuuid.time

class TestInteger(BaseCassEngTestCase):
    class IntegerTest(Model):
        __keyspace__ = 'test'
        test_id = UUID(primary_key=True, default=lambda:uuid4())
        value   = Integer(default=0, required=True)

    def test_default_zero_fields_validate(self):
        """ Tests that integer columns with a default value of 0 validate """
        it = self.IntegerTest()
        it.validate()

class TestBigInt(BaseCassEngTestCase):
    class BigIntTest(Model):
        __keyspace__ = 'test'
        test_id = UUID(primary_key=True, default=lambda:uuid4())
        value   = BigInt(default=0, required=True)

    def test_default_zero_fields_validate(self):
        """ Tests that bigint columns with a default value of 0 validate """
        it = self.BigIntTest()
        it.validate()

class TestText(BaseCassEngTestCase):

    def test_min_length(self):
        #min len defaults to 1
        col = Text()
        col.validate('')

        col.validate('b')

        #test not required defaults to 0
        Text(required=False).validate('')

        #test arbitrary lengths
        Text(min_length=0).validate('')
        Text(min_length=5).validate('blake')
        Text(min_length=5).validate('blaketastic')
        with self.assertRaises(ValidationError):
            Text(min_length=6).validate('blake')

    def test_max_length(self):

        Text(max_length=5).validate('blake')
        with self.assertRaises(ValidationError):
            Text(max_length=5).validate('blaketastic')

    def test_type_checking(self):
        Text().validate('string')
        Text().validate(u'unicode')
        Text().validate(bytearray('bytearray'))

        with self.assertRaises(ValidationError):
            Text(required=True).validate(None)

        with self.assertRaises(ValidationError):
            Text().validate(5)

        with self.assertRaises(ValidationError):
            Text().validate(True)

    def test_non_required_validation(self):
        """ Tests that validation is ok on none and blank values if required is False """
        Text().validate('')
        Text().validate(None)

class TestExtraFieldsRaiseException(BaseCassEngTestCase):
    class TestModel(Model):
        __keyspace__ = 'test'
        id = UUID(primary_key=True, default=uuid4)

    def test_extra_field(self):
        with self.assertRaises(ValidationError):
            self.TestModel.create(bacon=5000)

class TestPythonDoesntDieWhenExtraFieldIsInCassandra(BaseCassEngTestCase):
    class TestModel(Model):
        __keyspace__ = 'test'
        __table_name__ = 'alter_doesnt_break_running_app'
        id = UUID(primary_key=True, default=uuid4)

    def test_extra_field(self):
        drop_table(self.TestModel)
        sync_table(self.TestModel)
        self.TestModel.create()
        execute("ALTER TABLE {} add blah int".format(self.TestModel.column_family_name(include_keyspace=True)))
        self.TestModel.objects().all()

class TestTimeUUIDFromDatetime(TestCase):
    def test_conversion_specific_date(self):
        dt = datetime(1981, 7, 11, microsecond=555000)

        uuid = TimeUUID.from_datetime(dt)

        from uuid import UUID
        assert isinstance(uuid, UUID)

        ts = (uuid.time - 0x01b21dd213814000) / 1e7 # back to a timestamp
        new_dt = datetime.utcfromtimestamp(ts)

        # checks that we created a UUID1 with the proper timestamp
        assert new_dt == dt

class TestInet(BaseCassEngTestCase):

    class InetTestModel(Model):
        id = UUID(primary_key=True, default=uuid4)
        address = Inet()

    def setUp(self):
        drop_table(self.InetTestModel)
        sync_table(self.InetTestModel)

    def test_inet_saves(self):
        tmp = self.InetTestModel.create(address="192.168.1.1")

        m = self.InetTestModel.get(id=tmp.id)

        assert m.address == "192.168.1.1"

    def test_non_address_fails(self):
        with self.assertRaises(InvalidRequest):
            self.InetTestModel.create(address="ham sandwich")

<|endoftext|>


<commit_before>'''
Copyright 2014 Tyler Palsulich

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.
'''

# This bot is a result of a tutoral covered on http://shellium.org/wiki.
import socket
import sys
from datetime import datetime

# Check if we need to print help.
if len(sys.argv) == 1 or len(sys.argv) > 4 or sys.argv[1].find("#") != 0:
  print("Usage: python autoack.py #channel [nick [server]]")
  sys.exit()

# Variables for how to connect the bot.
channel = sys.argv[1] # Channel
botnick = "AutoAck" if len(sys.argv) < 3 else sys.argv[2] # Nickname
server = "chat.freenode.net" if len(sys.argv) < 4 else sys.argv[3] # Server

# Substring used to split the received message into the actual message content
splitter = "PRIVMSG " + channel + " :"

# Time used to prevent sending messages while in quiet mode.
can_send_after = datetime.now()

# Map from keywords to how the bot will respond in chat.
default_commands = {
            "ack":  "ack",
            "git":  "#gitpush",
            "aye":  "aye, mate!",
            "+1":   "+1",
            "boom": "kaboom!!!",
            "beum": "kabeum!!!",
            "bewm": "ba-bewm!!!",
            "seen": "seen like an eaten jelly bean"}

# Map where chatroom members can have the bot "learn" commands.
user_commands = {}

# Respond to a PING from the server.
def pong(data):
  ircsock.send("PONG " + data.split()[1] + "\n")

# Send a message to the connected server.
def send(message):
  if datetime.now() > can_send_after:
    ircsock.send("PRIVMSG " + channel + " :" + message + "\n")

def join_channel(channel):
  ircsock.send("JOIN " + channel + "\n")

# Respond to any keywords from the map `commands` in the string `message`.
def handle(message, commands):
  for key in commands:
    if key in message:
      send((commands[key] + " ") * message.count(key, 0))

# Store the given key and value in the user_commands map. But, do not
# allow the users to change default commands.
def learn(key, value):
  if key not in default_commands:
    if key in user_commands:
      send("Relearned " + key)
    else:
      send("Learned " + key)
    user_commands[key] = " ".join(value)
  else:
    send("Go away!")

# Forget the user command with the given key.
def forget(key):
  if key in default_commands:
    send("No.")
  elif key in user_commands:
    user_commands.pop(key)
    send("Dropped like a bad habit.")
  else:
    send("Maybe you're the one forgetting...")

def send_help():
  send("Available commands:")
  send("   " + botnick + ": learn [key] [value]")
  send("   " + botnick + ": forget [key]")

# Connect to the server.
ircsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
ircsock.connect((server, 6667)) # Connect to the server using port 6667.
ircsock.send("USER " + botnick + " " + botnick + " " + botnick + " :.\n") # Authenticate the bot.
ircsock.send("NICK " + botnick + "\n") # Assign the nickname to the bot.

join_channel(channel)

# Loop forever, waiting for messages to arrive.
while 1:
  message = ircsock.recv(2048) # Receive data from the server.
  message = message.strip('\n\r') # Remove any unnecessary linebreaks.

  print(message)

  if "PING :" in message: pong(message)

  # Only respond to chat from the current chatroom (not private or administrative log in messages).
  if splitter not in message: continue

  # Get the content of the message.
  message = message.split(splitter)[1]

  # Convert to lowercase and split the message based on whitespace.
  split = message.lower().split()

  if split[0] == botnick.lower() + ":":   # Command addressed to the bot (e.g. learn or forget).
    if split[1] == "learn" and len(split) > 2:
      learn(split[2], message.split()[3:])
    if split[1] == "forget":
      forget(split[2])
    if split[1] == "help":
      send_help()
    if split[1] == "quiet":
      can_send_after = datetime.now() + datetime.timedelta(seconds=30)
  else:   # Only handle messages that aren't sent directly to the bot.
    handle(message.lower(), default_commands)
    handle(message.lower(), user_commands)
<commit_msg>Print which server and port the bot is trying to connect to.<commit_after>'''
Copyright 2014 Tyler Palsulich

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.
'''

# This bot is a result of a tutoral covered on http://shellium.org/wiki.
import socket
import sys
from datetime import datetime

# Check if we need to print help.
if len(sys.argv) == 1 or len(sys.argv) > 4 or sys.argv[1].find("#") != 0:
  print("Usage: python autoack.py #channel [nick [server]]")
  sys.exit()

# Variables for how to connect the bot.
channel = sys.argv[1] # Channel
botnick = "AutoAck" if len(sys.argv) < 3 else sys.argv[2] # Nickname
server = "chat.freenode.net" if len(sys.argv) < 4 else sys.argv[3] # Server
port = 6667 # Port to connect on.

# Substring used to split the received message into the actual message content
splitter = "PRIVMSG " + channel + " :"

# Time used to prevent sending messages while in quiet mode.
can_send_after = datetime.now()

# Map from keywords to how the bot will respond in chat.
default_commands = {
            "ack":  "ack",
            "git":  "#gitpush",
            "aye":  "aye, mate!",
            "+1":   "+1",
            "boom": "kaboom!!!",
            "beum": "kabeum!!!",
            "bewm": "ba-bewm!!!",
            "seen": "seen like an eaten jelly bean"}

# Map where chatroom members can have the bot "learn" commands.
user_commands = {}

# Respond to a PING from the server.
def pong(data):
  ircsock.send("PONG " + data.split()[1] + "\n")

# Send a message to the connected server.
def send(message):
  if datetime.now() > can_send_after:
    ircsock.send("PRIVMSG " + channel + " :" + message + "\n")

def join_channel(channel):
  ircsock.send("JOIN " + channel + "\n")

# Respond to any keywords from the map `commands` in the string `message`.
def handle(message, commands):
  for key in commands:
    if key in message:
      send((commands[key] + " ") * message.count(key, 0))

# Store the given key and value in the user_commands map. But, do not
# allow the users to change default commands.
def learn(key, value):
  if key not in default_commands:
    if key in user_commands:
      send("Relearned " + key)
    else:
      send("Learned " + key)
    user_commands[key] = " ".join(value)
  else:
    send("Go away!")

# Forget the user command with the given key.
def forget(key):
  if key in default_commands:
    send("No.")
  elif key in user_commands:
    user_commands.pop(key)
    send("Dropped like a bad habit.")
  else:
    send("Maybe you're the one forgetting...")

def send_help():
  send("Available commands:")
  send("   " + botnick + ": learn [key] [value]")
  send("   " + botnick + ": forget [key]")

# Connect to the server.
ircsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
print("Attempting to connect to " + server + " on port " + str(port))
ircsock.connect((server, port)) # Connect to the server using port 6667.
ircsock.send("USER " + botnick + " " + botnick + " " + botnick + " :.\n") # Authenticate the bot.
ircsock.send("NICK " + botnick + "\n") # Assign the nickname to the bot.

join_channel(channel)

# Loop forever, waiting for messages to arrive.
while 1:
  message = ircsock.recv(2048) # Receive data from the server.
  message = message.strip('\n\r') # Remove any unnecessary linebreaks.

  print(message)

  if "PING :" in message: pong(message)

  # Only respond to chat from the current chatroom (not private or administrative log in messages).
  if splitter not in message: continue

  # Get the content of the message.
  message = message.split(splitter)[1]

  # Convert to lowercase and split the message based on whitespace.
  split = message.lower().split()

  if split[0] == botnick.lower() + ":":   # Command addressed to the bot (e.g. learn or forget).
    if split[1] == "learn" and len(split) > 2:
      learn(split[2], message.split()[3:])
    if split[1] == "forget":
      forget(split[2])
    if split[1] == "help":
      send_help()
    if split[1] == "quiet":
      can_send_after = datetime.now() + datetime.timedelta(seconds=30)
  else:   # Only handle messages that aren't sent directly to the bot.
    handle(message.lower(), default_commands)
    handle(message.lower(), user_commands)
<|endoftext|>


<commit_before>"""
Utility functions and classes for receiver function calculation.
"""
import collections
import inspect
import itertools
from pkg_resources import resource_filename

from decorator import decorator
import numpy as np

DEG2KM = 111.2  #: Conversion factor from degrees epicentral distance to km

def _get_stations(inventory):
    channels = inventory.get_contents()['channels']
    stations = {ch[:-1] + '?': ch[-1] for ch in channels}
    return stations

def iter_event_data(events, inventory, get_waveforms, phase='P',
                    request_window=None, pad=10, pbar=None, **kwargs):
    """
    Return iterator yielding three component streams per station and event.

    :param events: list of events or `~obspy.core.event.Catalog` instance
    :param inventory: `~obspy.core.inventory.inventory.Inventory` instance
        with station and channel information
    :param get_waveforms: Function returning the data. It has to take the
        arguments network, station, location, channel, starttime, endtime.
    :param phase: Considered phase, e.g. 'P', 'S', 'PP'
    :type request_window: tuple (start, end)
    :param request_window: requested time window around the onset of the phase
    :param float pad: add specified time in seconds to request window and
       trim afterwards again
    :param pbar: tqdm_ instance for displaying a progressbar

    :return: three component streams with raw data

    Example usage with progressbar::

        from tqdm import tqdm
        from rf.util import iter_event_data
        with tqdm() as t:
            for stream3c in iter_event_data(*args, pbar=t):
                do_something(stream3c)

    .. _tqdm: https://pypi.python.org/pypi/tqdm
    """
    from rf.rfstream import rfstats, RFStream
    method = phase[-1].upper()
    if request_window is None:
        request_window = (-50, 150) if method == 'P' else (-100, 50)
    stations = _get_stations(inventory)
    if pbar is not None:
        pbar.total = len(events) * len(stations)
    for event, seedid in itertools.product(events, stations):
        if pbar is not None:
            pbar.update(1)
        origin_time = (event.preferred_origin() or event.origins[0])['time']
        try:
            args = (seedid[:-1] + stations[seedid], origin_time)
            coords = inventory.get_coordinates(*args)
        except:  # station not available at that time
            continue
        stats = rfstats(station=coords, event=event, phase=phase, **kwargs)
        if not stats:
            continue
        net, sta, loc, cha = seedid.split('.')
        starttime = stats.onset + request_window[0]
        endtime = stats.onset + request_window[1]
        kws = {'network': net, 'station': sta, 'location': loc,
               'channel': cha, 'starttime': starttime - pad,
               'endtime': endtime + pad}
        try:
            stream = get_waveforms(**kws)
        except:  # no data available
            continue
        stream.trim(starttime, endtime)
        stream.merge()
        if len(stream) != 3:
            from warnings import warn
            warn('Need 3 component seismograms. %d components '
                 'detected for event %s, station %s.'
                 % (len(stream), event.resource_id, seedid))
            continue
        if any(isinstance(tr.data, np.ma.masked_array) for tr in stream):
            from warnings import warn
            warn('Gaps or overlaps detected for event %s, station %s.'
                 % (event.resource_id, seedid))
            continue
        for tr in stream:
            tr.stats.update(stats)
        yield RFStream(stream, warn=False)

def iter_event_metadata(events, inventory, pbar=None):
    """
    Return iterator yielding metadata per station and event.

    :param events: list of events or `~obspy.core.event.Catalog` instance
    :param inventory: `~obspy.core.inventory.inventory.Inventory` instance
        with station and channel information
    :param pbar: tqdm_ instance for displaying a progressbar
    """
    stations = _get_stations(inventory)
    if events is None:
        events = [None]
    if pbar is not None:
        pbar.total = len(events) * len(stations)
    for event, seedid in itertools.product(events, stations):
        if pbar is not None:
            pbar.update(1)
        net, sta, loc, cha = seedid.split('.')
        meta = {'network': net, 'station': sta, 'location': loc,
                'channel': cha}
        if event is not None:
            ot = (event.preferred_origin() or event.origins[0])['time']
            meta['event_time'] = ot
        yield meta

class IterMultipleComponents(object):

    """
    Return iterable to iterate over associated components of a stream.

    :param stream: Stream with different, possibly many traces. It is
        split into substreams with the same seed id (only last character
        i.e. component may vary)
    :type key: str or None
    :param key: Additionally, the stream is grouped by the values of
         the given stats entry to differentiate between e.g. different events
         (for example key='starttime', key='onset')
    :type number_components: int, tuple of ints or None
    :param number_components: Only iterate through substreams with
         matching number of components.
    """

    def __init__(self, stream, key=None, number_components=None):
        substreams = collections.defaultdict(stream.__class__)
        for tr in stream:
            k = (tr.id[:-1], str(tr.stats[key]) if key is not None else None)
            substreams[k].append(tr)
        n = number_components
        self.substreams = [s for _, s in sorted(substreams.items())
                           if n is None or len(s) == n or len(s) in n]

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

    def __iter__(self):
        for s in self.substreams:
            yield s

def direct_geodetic(latlon, azi, dist):
    """
    Solve direct geodetic problem with geographiclib.

    :param tuple latlon: coordinates of first point
    :param azi: azimuth of direction
    :param dist: distance in km

    :return: coordinates (lat, lon) of second point on a WGS84 globe
    """
    from geographiclib.geodesic import Geodesic
    coords = Geodesic.WGS84.Direct(latlon[0], latlon[1], azi, dist * 1000)
    return coords['lat2'], coords['lon2']

__CACHE = {}

def minimal_example_rf():
    """
    Return receiver functions calculated from the data returned by read_rf()
    """
    cache_key = 'minimal_example_rf'
    if cache_key in __CACHE:
        return __CACHE[cache_key].copy()
    from rf.rfstream import read_rf, rfstats
    stream = read_rf()
    rfstats(stream)
    stream.filter('bandpass', freqmin=0.5, freqmax=2)
    stream.trim2(10, 110, reftime='starttime')
    stream.rf(winsrc=(-5, 25, 5))
    stream.moveout()
    stream.trim2(-10, 80, reftime='onset')
    stream.ppoints(50)
    __CACHE[cache_key] = stream
    return stream.copy()

def minimal_example_Srf():
    """
    Return S receiver functions calculated from example data
    """
    cache_key = 'minimal_example_Srf'
    if cache_key in __CACHE:
        return __CACHE[cache_key].copy()
    from rf.rfstream import read_rf, rfstats
    fname = resource_filename('rf', 'example/minimal_example_S.tar.gz')
    stream = read_rf(fname)
    rfstats(stream, phase='S')
    stream.filter('bandpass', freqmin=0.2, freqmax=0.5)
    stream.trim2(10, 120, reftime='starttime')
    stream.rf(method='S', winsrc=(-5, 15, 5))
    stream.moveout(phase='Sp')
    stream.ppoints(50, pp_phase='P')
    __CACHE[cache_key] = stream
    return stream.copy()

@decorator
def _add_processing_info(func, *args, **kwargs):
    from rf import __version__
    args_ = inspect.getcallargs(func, *args, **kwargs)
    if args_.pop('self', None) is None:
        args_.pop('stream')
    kw = args_.pop('kwargs', {})
    kw.update(args_)
    boxes = kw.pop('boxes', None)
    if boxes:
        kw['latlon'] = boxes[0]['profile']['latlon']
        kw['azimuth'] = boxes[0]['profile']['azimuth']
        kw['length'] = boxes[0]['profile']['length']
    info = 'rf {version}: {function}(%s)'.format(
        version=__version__, function=func.__name__)
    arguments = ['%s=%s' % (k, repr(v)) if not isinstance(v, str) else
                 "%s='%s'" % (k, v) for k, v in kw.items()]
    info = info % '::'.join(sorted(arguments))
    stream = func(*args, **kwargs)
    try:
        for tr in stream:
            tr._internal_add_processing_info(info)
    except:
        pass
    return stream
<commit_msg>minor typo<commit_after>"""
Utility functions and classes for receiver function calculation.
"""
import collections
import inspect
import itertools
from pkg_resources import resource_filename

from decorator import decorator
import numpy as np

DEG2KM = 111.2  #: Conversion factor from degrees epicentral distance to km

def _get_stations(inventory):
    channels = inventory.get_contents()['channels']
    stations = {ch[:-1] + '?': ch[-1] for ch in channels}
    return stations

def iter_event_data(events, inventory, get_waveforms, phase='P',
                    request_window=None, pad=10, pbar=None, **kwargs):
    """
    Return iterator yielding three component streams per station and event.

    :param events: list of events or `~obspy.core.event.Catalog` instance
    :param inventory: `~obspy.core.inventory.inventory.Inventory` instance
        with station and channel information
    :param get_waveforms: Function returning the data. It has to take the
        arguments network, station, location, channel, starttime, endtime.
    :param phase: Considered phase, e.g. 'P', 'S', 'PP'
    :type request_window: tuple (start, end)
    :param request_window: requested time window around the onset of the phase
    :param float pad: add specified time in seconds to request window and
       trim afterwards again
    :param pbar: tqdm_ instance for displaying a progressbar

    :return: three component streams with raw data

    Example usage with progressbar::

        from tqdm import tqdm
        from rf.util import iter_event_data
        with tqdm() as t:
            for stream3c in iter_event_data(*args, pbar=t):
                do_something(stream3c)

    .. _tqdm: https://pypi.python.org/pypi/tqdm
    """
    from rf.rfstream import rfstats, RFStream
    method = phase[-1].upper()
    if request_window is None:
        request_window = (-50, 150) if method == 'P' else (-100, 50)
    stations = _get_stations(inventory)
    if pbar is not None:
        pbar.total = len(events) * len(stations)
    for event, seedid in itertools.product(events, stations):
        if pbar is not None:
            pbar.update(1)
        origin_time = (event.preferred_origin() or event.origins[0])['time']
        try:
            args = (seedid[:-1] + stations[seedid], origin_time)
            coords = inventory.get_coordinates(*args)
        except:  # station not available at that time
            continue
        stats = rfstats(station=coords, event=event, phase=phase, **kwargs)
        if not stats:
            continue
        net, sta, loc, cha = seedid.split('.')
        starttime = stats.onset + request_window[0]
        endtime = stats.onset + request_window[1]
        kws = {'network': net, 'station': sta, 'location': loc,
               'channel': cha, 'starttime': starttime - pad,
               'endtime': endtime + pad}
        try:
            stream = get_waveforms(**kws)
        except:  # no data available
            continue
        stream.trim(starttime, endtime)
        stream.merge()
        if len(stream) != 3:
            from warnings import warn
            warn('Need 3 component seismograms. %d components '
                 'detected for event %s, station %s.'
                 % (len(stream), event.resource_id, seedid))
            continue
        if any(isinstance(tr.data, np.ma.masked_array) for tr in stream):
            from warnings import warn
            warn('Gaps or overlaps detected for event %s, station %s.'
                 % (event.resource_id, seedid))
            continue
        for tr in stream:
            tr.stats.update(stats)
        yield RFStream(stream, warn=False)

def iter_event_metadata(events, inventory, pbar=None):
    """
    Return iterator yielding metadata per station and event.

    :param events: list of events or `~obspy.core.event.Catalog` instance
    :param inventory: `~obspy.core.inventory.inventory.Inventory` instance
        with station and channel information
    :param pbar: tqdm_ instance for displaying a progressbar
    """
    stations = _get_stations(inventory)
    if events is None:
        events = [None]
    if pbar is not None:
        pbar.total = len(events) * len(stations)
    for event, seedid in itertools.product(events, stations):
        if pbar is not None:
            pbar.update(1)
        net, sta, loc, cha = seedid.split('.')
        meta = {'network': net, 'station': sta, 'location': loc,
                'channel': cha}
        if event is not None:
            ot = (event.preferred_origin() or event.origins[0])['time']
            meta['event_time'] = ot
        yield meta

class IterMultipleComponents(object):

    """
    Return iterable to iterate over associated components of a stream.

    :param stream: Stream with different, possibly many traces. It is
        split into substreams with the same seed id (only last character
        i.e. component may vary)
    :type key: str or None
    :param key: Additionally, the stream is grouped by the values of
         the given stats entry to differentiate between e.g. different events
         (for example key='starttime', key='onset')
    :type number_components: int, tuple of ints or None
    :param number_components: Only iterate through substreams with
         matching number of components.
    """

    def __init__(self, stream, key=None, number_components=None):
        substreams = collections.defaultdict(stream.__class__)
        for tr in stream:
            k = (tr.id[:-1], str(tr.stats[key]) if key is not None else None)
            substreams[k].append(tr)
        n = number_components
        self.substreams = [s for _, s in sorted(substreams.items())
                           if n is None or len(s) == n or len(s) in n]

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

    def __iter__(self):
        for s in self.substreams:
            yield s

def direct_geodetic(latlon, azi, dist):
    """
    Solve direct geodetic problem with geographiclib.

    :param tuple latlon: coordinates of first point
    :param azi: azimuth of direction
    :param dist: distance in km

    :return: coordinates (lat, lon) of second point on a WGS84 globe
    """
    from geographiclib.geodesic import Geodesic
    coords = Geodesic.WGS84.Direct(latlon[0], latlon[1], azi, dist * 1000)
    return coords['lat2'], coords['lon2']

__CACHE = {}

def minimal_example_rf():
    """
    Return receiver functions calculated from the data returned by read_rf().
    """
    cache_key = 'minimal_example_rf'
    if cache_key in __CACHE:
        return __CACHE[cache_key].copy()
    from rf.rfstream import read_rf, rfstats
    stream = read_rf()
    rfstats(stream)
    stream.filter('bandpass', freqmin=0.5, freqmax=2)
    stream.trim2(10, 110, reftime='starttime')
    stream.rf(winsrc=(-5, 25, 5))
    stream.moveout()
    stream.trim2(-10, 80, reftime='onset')
    stream.ppoints(50)
    __CACHE[cache_key] = stream
    return stream.copy()

def minimal_example_Srf():
    """
    Return S receiver functions calculated from example data.
    """
    cache_key = 'minimal_example_Srf'
    if cache_key in __CACHE:
        return __CACHE[cache_key].copy()
    from rf.rfstream import read_rf, rfstats
    fname = resource_filename('rf', 'example/minimal_example_S.tar.gz')
    stream = read_rf(fname)
    rfstats(stream, phase='S')
    stream.filter('bandpass', freqmin=0.2, freqmax=0.5)
    stream.trim2(10, 120, reftime='starttime')
    stream.rf(method='S', winsrc=(-5, 15, 5))
    stream.moveout(phase='Sp')
    stream.ppoints(50, pp_phase='P')
    __CACHE[cache_key] = stream
    return stream.copy()

@decorator
def _add_processing_info(func, *args, **kwargs):
    from rf import __version__
    args_ = inspect.getcallargs(func, *args, **kwargs)
    if args_.pop('self', None) is None:
        args_.pop('stream')
    kw = args_.pop('kwargs', {})
    kw.update(args_)
    boxes = kw.pop('boxes', None)
    if boxes:
        kw['latlon'] = boxes[0]['profile']['latlon']
        kw['azimuth'] = boxes[0]['profile']['azimuth']
        kw['length'] = boxes[0]['profile']['length']
    info = 'rf {version}: {function}(%s)'.format(
        version=__version__, function=func.__name__)
    arguments = ['%s=%s' % (k, repr(v)) if not isinstance(v, str) else
                 "%s='%s'" % (k, v) for k, v in kw.items()]
    info = info % '::'.join(sorted(arguments))
    stream = func(*args, **kwargs)
    try:
        for tr in stream:
            tr._internal_add_processing_info(info)
    except:
        pass
    return stream
<|endoftext|>


<commit_before>"""
kombu.transport.mongodb
=======================

MongoDB transport.

:copyright: (c) 2010 - 2012 by Flavio Percoco Premoli.
:license: BSD, see LICENSE for more details.

"""
from __future__ import absolute_import

import pymongo

from pymongo import errors
from anyjson import loads, dumps
from pymongo.connection import Connection

from kombu.exceptions import StdConnectionError, StdChannelError
from kombu.five import Empty

from . import virtual

DEFAULT_HOST = '127.0.0.1'
DEFAULT_PORT = 27017

__author__ = """\
Flavio [FlaPer87] Percoco Premoli <flaper87@flaper87.org>;\
Scott Lyons <scottalyons@gmail.com>;\
"""

class Channel(virtual.Channel):
    _client = None
    supports_fanout = True
    _fanout_queues = {}

    def __init__(self, *vargs, **kwargs):
        super_ = super(Channel, self)
        super_.__init__(*vargs, **kwargs)

        self._queue_cursors = {}
        self._queue_readcounts = {}

    def _new_queue(self, queue, **kwargs):
        pass

    def _get(self, queue):
        try:
            if queue in self._fanout_queues:
                msg = next(self._queue_cursors[queue])
                self._queue_readcounts[queue] += 1
                return loads(msg['payload'])
            else:
                msg = self.client.command(
                    'findandmodify', 'messages',
                    query={'queue': queue},
                    sort={'_id': pymongo.ASCENDING}, remove=True,
                )
        except errors.OperationFailure as exc:
            if 'No matching object found' in exc.args[0]:
                raise Empty()
            raise
        except StopIteration:
            raise Empty()

        # as of mongo 2.0 empty results won't raise an error
        if msg['value'] is None:
            raise Empty()
        return loads(msg['value']['payload'])

    def _size(self, queue):
        if queue in self._fanout_queues:
            return (self._queue_cursors[queue].count() -
                    self._queue_readcounts[queue])

        return self.client.messages.find({'queue': queue}).count()

    def _put(self, queue, message, **kwargs):
        self.client.messages.insert({'payload': dumps(message),
                                     'queue': queue})

    def _purge(self, queue):
        size = self._size(queue)
        if queue in self._fanout_queues:
            cursor = self._queue_cursors[queue]
            cursor.rewind()
            self._queue_cursors[queue] = cursor.skip(cursor.count())
        else:
            self.client.messages.remove({'queue': queue})
        return size

    def close(self):
        super(Channel, self).close()
        if self._client:
            self._client.connection.end_request()

    def _open(self):
        """
        See mongodb uri documentation:
        http://www.mongodb.org/display/DOCS/Connections
        """
        conninfo = self.connection.client

        dbname = None
        hostname = None

        if not conninfo.hostname:
            conninfo.hostname = DEFAULT_HOST

        for part in conninfo.hostname.split('/'):
            if not hostname:
                hostname = 'mongodb://' + part
                continue

            dbname = part
            if '?' in part:
                # In case someone is passing options
                # to the mongodb connection. Right now
                # it is not permitted by kombu
                dbname, options = part.split('?')
                hostname += '/?' + options

        hostname = "%s/%s" % (
            hostname, dbname in [None, "/"] and "admin" or dbname,
        )
        if not dbname or dbname == "/":
            dbname = "kombu_default"

        # At this point we expect the hostname to be something like
        # (considering replica set form too):
        #
        #   mongodb://[username:password@]host1[:port1][,host2[:port2],
        #   ...[,hostN[:portN]]][/[?options]]
        mongoconn = Connection(host=hostname)
        version = mongoconn.server_info()['version']
        if tuple(map(int, version.split('.')[:2])) < (1, 3):
            raise NotImplementedError(
                'Kombu requires MongoDB version 1.3+ (server is {0})'.format(
                    version))

        database = getattr(mongoconn, dbname)

        # This is done by the connection uri
        # if conninfo.userid:
        #     database.authenticate(conninfo.userid, conninfo.password)
        self.db = database
        col = database.messages
        col.ensure_index([('queue', 1), ('_id', 1)], background=True)

        if 'messages.broadcast' not in database.collection_names():
            capsize = conninfo.transport_options.get(
                'capped_queue_size') or 100000
            database.create_collection('messages.broadcast',
                                       size=capsize, capped=True)

        self.bcast = getattr(database, 'messages.broadcast')
        self.bcast.ensure_index([('queue', 1)])

        self.routing = getattr(database, 'messages.routing')
        self.routing.ensure_index([('queue', 1), ('exchange', 1)])
        return database

    #TODO: Store a more complete exchange metatable in the routing collection
    def get_table(self, exchange):
        """Get table of bindings for ``exchange``."""
        localRoutes = frozenset(self.state.exchanges[exchange]['table'])
        brokerRoutes = self.client.messages.routing.find(
            {'exchange': exchange}
        )

        return localRoutes | frozenset((r['routing_key'],
                                        r['pattern'],
                                        r['queue']) for r in brokerRoutes)

    def _put_fanout(self, exchange, message, **kwargs):
        """Deliver fanout message."""
        self.client.messages.broadcast.insert({'payload': dumps(message),
                                               'queue': exchange})

    def _queue_bind(self, exchange, routing_key, pattern, queue):
        if self.typeof(exchange).type == 'fanout':
            cursor = self.bcast.find(query={'queue': exchange},
                                     sort=[('$natural', 1)], tailable=True)
            # Fast forward the cursor past old events
            self._queue_cursors[queue] = cursor.skip(cursor.count())
            self._queue_readcounts[queue] = cursor.count()
            self._fanout_queues[queue] = exchange

        meta = {'exchange': exchange,
                'queue': queue,
                'routing_key': routing_key,
                'pattern': pattern}
        self.client.messages.routing.update(meta, meta, upsert=True)

    def queue_delete(self, queue, **kwargs):
        self.routing.remove({'queue': queue})
        super(Channel, self).queue_delete(queue, **kwargs)
        if queue in self._fanout_queues:
            self._queue_cursors[queue].close()
            self._queue_cursors.pop(queue, None)
            self._fanout_queues.pop(queue, None)

    @property
    def client(self):
        if self._client is None:
            self._client = self._open()
        return self._client

class Transport(virtual.Transport):
    Channel = Channel

    polling_interval = 1
    default_port = DEFAULT_PORT
    connection_errors = (StdConnectionError, errors.ConnectionFailure)
    channel_errors = (StdChannelError,
                      errors.ConnectionFailure,
                      errors.OperationFailure)
    driver_type = 'mongodb'
    driver_name = 'pymongo'

    def driver_version(self):
        return pymongo.version
<commit_msg>ssl option added to transport mongodb transport<commit_after>"""
kombu.transport.mongodb
=======================

MongoDB transport.

:copyright: (c) 2010 - 2012 by Flavio Percoco Premoli.
:license: BSD, see LICENSE for more details.

"""
from __future__ import absolute_import

import pymongo

from pymongo import errors
from anyjson import loads, dumps
from pymongo.connection import Connection

from kombu.exceptions import StdConnectionError, StdChannelError
from kombu.five import Empty

from . import virtual

DEFAULT_HOST = '127.0.0.1'
DEFAULT_PORT = 27017

__author__ = """\
Flavio [FlaPer87] Percoco Premoli <flaper87@flaper87.org>;\
Scott Lyons <scottalyons@gmail.com>;\
"""

class Channel(virtual.Channel):
    _client = None
    supports_fanout = True
    _fanout_queues = {}

    def __init__(self, *vargs, **kwargs):
        super_ = super(Channel, self)
        super_.__init__(*vargs, **kwargs)

        self._queue_cursors = {}
        self._queue_readcounts = {}

    def _new_queue(self, queue, **kwargs):
        pass

    def _get(self, queue):
        try:
            if queue in self._fanout_queues:
                msg = next(self._queue_cursors[queue])
                self._queue_readcounts[queue] += 1
                return loads(msg['payload'])
            else:
                msg = self.client.command(
                    'findandmodify', 'messages',
                    query={'queue': queue},
                    sort={'_id': pymongo.ASCENDING}, remove=True,
                )
        except errors.OperationFailure as exc:
            if 'No matching object found' in exc.args[0]:
                raise Empty()
            raise
        except StopIteration:
            raise Empty()

        # as of mongo 2.0 empty results won't raise an error
        if msg['value'] is None:
            raise Empty()
        return loads(msg['value']['payload'])

    def _size(self, queue):
        if queue in self._fanout_queues:
            return (self._queue_cursors[queue].count() -
                    self._queue_readcounts[queue])

        return self.client.messages.find({'queue': queue}).count()

    def _put(self, queue, message, **kwargs):
        self.client.messages.insert({'payload': dumps(message),
                                     'queue': queue})

    def _purge(self, queue):
        size = self._size(queue)
        if queue in self._fanout_queues:
            cursor = self._queue_cursors[queue]
            cursor.rewind()
            self._queue_cursors[queue] = cursor.skip(cursor.count())
        else:
            self.client.messages.remove({'queue': queue})
        return size

    def close(self):
        super(Channel, self).close()
        if self._client:
            self._client.connection.end_request()

    def _open(self):
        """
        See mongodb uri documentation:
        http://www.mongodb.org/display/DOCS/Connections
        """
        conninfo = self.connection.client

        dbname = None
        hostname = None

        if not conninfo.hostname:
            conninfo.hostname = DEFAULT_HOST

        for part in conninfo.hostname.split('/'):
            if not hostname:
                hostname = 'mongodb://' + part
                continue

            dbname = part
            if '?' in part:
                # In case someone is passing options
                # to the mongodb connection. Right now
                # it is not permitted by kombu
                dbname, options = part.split('?')
                hostname += '/?' + options

        hostname = "%s/%s" % (
            hostname, dbname in [None, "/"] and "admin" or dbname,
        )
        if not dbname or dbname == "/":
            dbname = "kombu_default"

        # At this point we expect the hostname to be something like
        # (considering replica set form too):
        #
        #   mongodb://[username:password@]host1[:port1][,host2[:port2],
        #   ...[,hostN[:portN]]][/[?options]]
        mongoconn = Connection(host=hostname, ssl=conninfo.ssl)
        version = mongoconn.server_info()['version']
        if tuple(map(int, version.split('.')[:2])) < (1, 3):
            raise NotImplementedError(
                'Kombu requires MongoDB version 1.3+ (server is {0})'.format(
                    version))

        database = getattr(mongoconn, dbname)

        # This is done by the connection uri
        # if conninfo.userid:
        #     database.authenticate(conninfo.userid, conninfo.password)
        self.db = database
        col = database.messages
        col.ensure_index([('queue', 1), ('_id', 1)], background=True)

        if 'messages.broadcast' not in database.collection_names():
            capsize = conninfo.transport_options.get(
                'capped_queue_size') or 100000
            database.create_collection('messages.broadcast',
                                       size=capsize, capped=True)

        self.bcast = getattr(database, 'messages.broadcast')
        self.bcast.ensure_index([('queue', 1)])

        self.routing = getattr(database, 'messages.routing')
        self.routing.ensure_index([('queue', 1), ('exchange', 1)])
        return database

    #TODO: Store a more complete exchange metatable in the routing collection
    def get_table(self, exchange):
        """Get table of bindings for ``exchange``."""
        localRoutes = frozenset(self.state.exchanges[exchange]['table'])
        brokerRoutes = self.client.messages.routing.find(
            {'exchange': exchange}
        )

        return localRoutes | frozenset((r['routing_key'],
                                        r['pattern'],
                                        r['queue']) for r in brokerRoutes)

    def _put_fanout(self, exchange, message, **kwargs):
        """Deliver fanout message."""
        self.client.messages.broadcast.insert({'payload': dumps(message),
                                               'queue': exchange})

    def _queue_bind(self, exchange, routing_key, pattern, queue):
        if self.typeof(exchange).type == 'fanout':
            cursor = self.bcast.find(query={'queue': exchange},
                                     sort=[('$natural', 1)], tailable=True)
            # Fast forward the cursor past old events
            self._queue_cursors[queue] = cursor.skip(cursor.count())
            self._queue_readcounts[queue] = cursor.count()
            self._fanout_queues[queue] = exchange

        meta = {'exchange': exchange,
                'queue': queue,
                'routing_key': routing_key,
                'pattern': pattern}
        self.client.messages.routing.update(meta, meta, upsert=True)

    def queue_delete(self, queue, **kwargs):
        self.routing.remove({'queue': queue})
        super(Channel, self).queue_delete(queue, **kwargs)
        if queue in self._fanout_queues:
            self._queue_cursors[queue].close()
            self._queue_cursors.pop(queue, None)
            self._fanout_queues.pop(queue, None)

    @property
    def client(self):
        if self._client is None:
            self._client = self._open()
        return self._client

class Transport(virtual.Transport):
    Channel = Channel

    polling_interval = 1
    default_port = DEFAULT_PORT
    connection_errors = (StdConnectionError, errors.ConnectionFailure)
    channel_errors = (StdChannelError,
                      errors.ConnectionFailure,
                      errors.OperationFailure)
    driver_type = 'mongodb'
    driver_name = 'pymongo'

    def driver_version(self):
        return pymongo.version
<|endoftext|>


<commit_before>import time

import numpy as np
import pytest

from utilities.norm import (chi2_model_norms, chi2_model_norms_slow, continuum,
                            continuum_slow, get_continuum_points,
                            local_normalization, local_normalization_slow,
                            spec_local_norm, spec_local_norm_slow)

def test_continuums_are_equvalient(host, norm_method):
    x, y = host.xaxis, host.flux

    start = time.time()
    cont_s = continuum_slow(x, y, method=norm_method)
    mid = time.time()
    cont = continuum(x, y, method=norm_method)
    end = time.time()
    t1, t2 = mid - start, end - mid

    assert np.allclose(cont_s, cont)
    assert t2 < t1
    assert t2 < t1 * 0.20  # Reduction by 80%

def test_local_normalization_equivelent(host, norm_method):
    x, y = host.xaxis, host.flux

    assert np.allclose(local_normalization_slow(x, y, splits=50, method=norm_method, plot=False, top=5),
                       local_normalization(x, y, splits=50, method=norm_method, plot=False, top=5))

@pytest.mark.parametrize("method1,method2", [
    ("scalar", "linear"),
    ("quadratic", "scalar"),
    ("quadratic", "exponential")])
def test_local_normalization_methods_work(host, method1, method2):
    x, y = host.xaxis, host.flux
    norm_1 = local_normalization(x, y, splits=50, method=method1, top=5)
    norm_2 = local_normalization(x, y, splits=50, method=method2, top=5)

    assert not np.allclose(norm_1, norm_2)

def test_spec_local_norm_equivalent(host, norm_method):
    norm = spec_local_norm(host, method=norm_method, top=5)
    norm_s = spec_local_norm_slow(host, method=norm_method, top=5)

    assert np.allclose(norm.xaxis, norm_s.xaxis)
    assert np.allclose(norm.flux, norm_s.flux)

def test_spec_local_norm_applies_local_normalization(host, norm_method):
    norm = spec_local_norm_slow(host, method=norm_method, top=5)
    local_norm = local_normalization_slow(host.xaxis, host.flux, method=norm_method, top=5)

    assert np.allclose(norm.flux, local_norm)

def test_manual_normalization():
    x = np.arange(1, 100)
    y = np.arange(1, 100)

    local_norm = local_normalization(x, y, splits=5, plot=False, top=5, method="linear")
    local_norm_s = local_normalization_slow(x, y, splits=5, plot=False, top=5, method="linear")

    cont = continuum(x, y, splits=3, plot=False, top=2, method="linear")
    cont_s = continuum_slow(x, y, splits=3, plot=False, top=2, method="linear")

    assert np.allclose(cont_s, y)
    assert np.allclose(local_norm, np.ones_like(y))

    assert np.allclose(cont_s, cont)
    assert np.allclose(local_norm, local_norm_s)

def test_chi2_model_norms(host, tcm_model, norm_method):
    host.wav_select(2100.5, 2104.5)   # cut to avoid Nans from doppler shifts
    wave = host.xaxis
    obs = host.xaxis
    models = tcm_model(wave)

    start = time.time()
    chi2norm_s = chi2_model_norms_slow(wave, obs, models, method=norm_method, splits=15, top=10)
    mid = time.time()
    chi2norm = chi2_model_norms(wave, obs, models, method=norm_method, splits=15, top=10)
    end = time.time()

    assert chi2norm_s.shape == chi2norm.shape
    assert np.allclose(chi2norm_s, chi2norm)

    t1, t2 = mid - start, end - mid
    assert t2 < t1
    assert t2 < t1 * 0.50   # Reduces time by 50 %

@pytest.mark.parametrize("splits", [13, 27, 50])
def test_shortening_array(splits):
    x = np.arange(2033)
    rem = len(x) % splits
    z = x[:-rem]

    while len(x) % splits != 0:
        x = x[:-1]

    assert len(x) == len(z)
    assert np.allclose(x, z)

@pytest.mark.parametrize("splits,top,size", [
    (10, 5, (10,)),   # zero remainder
    (11, 4, (11,)),
    (51, 5, (51,)),
    (100, 7, (100,))  # zero remainder
])
def test_get_continuum_points(splits, top, size):

    x = np.arange(200)
    x1, x2 = get_continuum_points(x, x, splits=splits, top=top)

    assert np.all(x1 == x2)
    assert x1.shape == size

def test_continuum():
    continuum()
    assert 0<commit_msg>Some continuum tests.<commit_after>import time

import numpy as np
import pytest

from utilities.norm import (chi2_model_norms, chi2_model_norms_slow, continuum,
                            continuum_slow, get_continuum_points,
                            local_normalization, local_normalization_slow,
                            spec_local_norm, spec_local_norm_slow)

def test_continuums_are_equvalient(host, norm_method):
    x, y = host.xaxis, host.flux

    start = time.time()
    cont_s = continuum_slow(x, y, method=norm_method)
    mid = time.time()
    cont = continuum(x, y, method=norm_method)
    end = time.time()
    t1, t2 = mid - start, end - mid

    assert np.allclose(cont_s, cont)
    assert t2 < t1
    assert t2 < t1 * 0.20  # Reduction by 80%

def test_local_normalization_equivelent(host, norm_method):
    x, y = host.xaxis, host.flux

    assert np.allclose(local_normalization_slow(x, y, splits=50, method=norm_method, plot=False, top=5),
                       local_normalization(x, y, splits=50, method=norm_method, plot=False, top=5))

@pytest.mark.parametrize("method1,method2", [
    ("scalar", "linear"),
    ("quadratic", "scalar"),
    ("quadratic", "exponential")])
def test_local_normalization_methods_work(host, method1, method2):
    x, y = host.xaxis, host.flux
    norm_1 = local_normalization(x, y, splits=50, method=method1, top=5)
    norm_2 = local_normalization(x, y, splits=50, method=method2, top=5)

    assert not np.allclose(norm_1, norm_2)

def test_spec_local_norm_equivalent(host, norm_method):
    norm = spec_local_norm(host, method=norm_method, top=5)
    norm_s = spec_local_norm_slow(host, method=norm_method, top=5)

    assert np.allclose(norm.xaxis, norm_s.xaxis)
    assert np.allclose(norm.flux, norm_s.flux)

def test_spec_local_norm_applies_local_normalization(host, norm_method):
    norm = spec_local_norm_slow(host, method=norm_method, top=5)
    local_norm = local_normalization_slow(host.xaxis, host.flux, method=norm_method, top=5)

    assert np.allclose(norm.flux, local_norm)

def test_manual_normalization():
    x = np.arange(1, 100)
    y = np.arange(1, 100)

    local_norm = local_normalization(x, y, splits=5, plot=False, top=5, method="linear")
    local_norm_s = local_normalization_slow(x, y, splits=5, plot=False, top=5, method="linear")

    cont = continuum(x, y, splits=3, plot=False, top=2, method="linear")
    cont_s = continuum_slow(x, y, splits=3, plot=False, top=2, method="linear")

    assert np.allclose(cont_s, y)
    assert np.allclose(local_norm, np.ones_like(y))

    assert np.allclose(cont_s, cont)
    assert np.allclose(local_norm, local_norm_s)

def test_chi2_model_norms(host, tcm_model, norm_method):
    host.wav_select(2110.5, 2114.5)   # cut to avoid Nans from doppler shifts
    wave = host.xaxis
    obs = host.xaxis
    models = tcm_model(wave)

    start = time.time()
    chi2norm_s = chi2_model_norms_slow(wave, obs, models, method=norm_method, splits=15, top=10)
    mid = time.time()
    chi2norm = chi2_model_norms(wave, obs, models, method=norm_method, splits=15, top=10)
    end = time.time()

    assert chi2norm_s.shape == chi2norm.shape
    assert np.allclose(chi2norm_s, chi2norm)

    t1, t2 = mid - start, end - mid
    assert t2 < t1
    assert t2 < t1 * 0.50   # Reduces time by 50 %

@pytest.mark.parametrize("splits", [13, 27, 50])
def test_shortening_array(splits):
    x = np.arange(2033)
    rem = len(x) % splits
    z = x[:-rem]

    while len(x) % splits != 0:
        x = x[:-1]

    assert len(x) == len(z)
    assert np.allclose(x, z)

@pytest.mark.parametrize("splits,top,size", [
    (10, 5, (10,)),   # zero remainder
    (11, 4, (11,)),
    (51, 5, (51,)),
    (100, 7, (100,))  # zero remainder
])
def test_get_continuum_points(splits, top, size):

    x = np.arange(200)
    x1, x2 = get_continuum_points(x, x, splits=splits, top=top)

    assert np.all(x1 == x2)
    assert x1.shape == size

@pytest.mark.parametrize("scale", [1,2,3,10])
def test_continuum_scalar(scale):
    x = np.linspace(2000, 2100, 2000)
    y = scale * np.ones(2000)
    cont = continuum(x, y, method="scalar")
    assert np.allclose(np.mean(cont), scale)
    assert np.allclose(np.mean(y / cont), 1)

@pytest.mark.parametrize("x1, x0", [(0.1, .5), (0.002, .7), (0.08, 1)])
def test_continuum_linear(x1, x0):
    x = np.linspace(2000, 2100, 2000)
    y = x1*x + x0

    cont = continuum(x, y, method="linear")
    assert np.allclose(np.mean(y / cont), 1)
    assert np.allclose(cont, y)

@pytest.mark.parametrize("x1, x0", [(0.1, .5), (0.002, .7), (0.08, 1)])
def test_continuum_exponential(x1, x0):
    x = np.linspace(2000, 2100, 2000)
    y = np.exp(x1*x) + x0

    cont = continuum(x, y, method="exponential")
    assert np.allclose(np.mean(y / cont), 1)
    assert np.allclose(cont, y, 2)

<|endoftext|>


<commit_before>from __future__ import print_function

import os
import matplotlib as mpl
mpl.use("Agg")

from math import sqrt
import numpy as np
import math
from pandas.util.testing import DataFrame, Series
import matplotlib.pyplot as plt
import pandas as pd

__author__ = 'dietz'

from argparse import ArgumentParser

# query metric value
# C09-1	ndcg	0.27478
# C09-1	ndcg5	0.47244
# C09-1	ndcg10	0.32972
# C09-1	ndcg20	0.25703
# C09-1	ERR	0.18652
# C09-1	ERR10	0.16907
# C09-1	ERR20	0.17581
# C09-1	P1	1.00000

def is_valid_file(parser, arg):
    if not os.path.exists(arg):
        parser.error("The file %s does not exist!" % arg)
    else:
        return arg

tooldescription = """
Classic bar chart indicating the mean of values for the given
metric across all queries with error bars indicating the standard
error.
"""
parser = ArgumentParser(description=tooldescription)
parser.add_argument('--out', help='outputfilename', metavar='FILE',  required=True)
parser.add_argument('--metric', help='metric for comparison', required=True)
parser.add_argument('--format', help='trec_eval output or galago_eval output', default='trec_eval')
parser.add_argument('-c', help='instead of average, also count non-existing queries', default=False, action='store_true')
parser.add_argument('--sort', help='sort methods in plot', action='store_true', default=False)
parser.add_argument(dest='runs', nargs='+', type=lambda x: is_valid_file(parser, x))
args = parser.parse_args()

numQueries_key = "num_q"

print("column.py metric="+args.metric+" out="+args.out)

def read_ssv(fname):
    lines = [line.split() for line in open(fname, 'r')]
    if args.format.lower() == 'galago_eval':
        return lines
    elif args.format.lower() == 'trec_eval':
        return [[line[1], line[0]] + line[2:] for line in lines]

def readNumQueries(run):
    tsv = read_ssv(run)
    data = [int(row[2]) for row in tsv if row[0] == "all" and row[1] == numQueries_key]
    return data[0]

def findQueriesWithNanValues(run):
    tsv = read_ssv(run)
    # print ("tsv,", tsv)
    queriesWithNan = {row[0] for row in tsv if row[1] == 'num_rel' and (float(row[2]) == 0.0 or math.isnan(float(row[2])))}
    return queriesWithNan

def fetchValues(run):
    tsv = read_ssv(run)
    data = {row[0]: float(row[2]) for row in tsv if row[1] == args.metric and not math.isnan(float(row[2]))}
    return data

datas = {run: fetchValues(run) for run in args.runs}

# deal with nans
queriesWithNanValues = {'all'}.union(*[findQueriesWithNanValues(run) for run in args.runs])
basedata=datas[args.runs[0]]
queries = set(basedata.keys()).difference(queriesWithNanValues)
numQueries = readNumQueries(args.runs[0]) if args.c else len(queries)

seriesDict = {'mean':dict(), 'stderr':dict()}

for run in datas:
    data = datas[run]

    if sum(not key in data for key in queries) > 0:
        print("data for run "+run+" does not contain all queries "+" ".join(queries))

    mean = np.sum([data.get(key, 0.0) for key in queries]) / numQueries
    stderr = np.std([data.get(key, 0.0) for key in queries] + ([0.0]* (numQueries - len(queries)))) / sqrt(numQueries)
    seriesDict['mean'][run]=mean
    seriesDict['stderr'][run]=stderr

print( "dropping queries because of NaN values: "+ " ".join(queriesWithNanValues))

print ('\t'.join(['run', 'mean/stderr']))
for run in datas:
    #if not run == args.runs[0]:
    print ('\t'.join([run, str(seriesDict['mean'][run]), str(seriesDict['stderr'][run])]))

df1 = DataFrame(seriesDict, index=args.runs)
if args.sort:
	df1.sort_values('mean',ascending=False,inplace=True)
df2 = df1['mean']
df2.index=[os.path.basename(label) for label in df1.index]

fig, ax = plt.subplots()
#plt.figure()
df2.plot(kind='bar', yerr = df1['stderr'], color=['0.0', '0.6', '0.4', '0.8','0.0', '0.6', '0.4', '0.8','0.0', '0.6', '0.4', '0.8'], ax=ax )
ax.grid()
plt.ylabel(args.metric, fontsize=20)
plt.tick_params(axis='both', which='major', labelsize=20)
plt.xticks(rotation=90)
plt.savefig(args.out, bbox_inches='tight')

# plt.show()
<commit_msg>minir-plots documented that run files  is actually a eval files<commit_after>from __future__ import print_function

import os
import matplotlib as mpl
mpl.use("Agg")

from math import sqrt
import numpy as np
import math
from pandas.util.testing import DataFrame, Series
import matplotlib.pyplot as plt
import pandas as pd

__author__ = 'dietz'

from argparse import ArgumentParser

# query metric value
# C09-1	ndcg	0.27478
# C09-1	ndcg5	0.47244
# C09-1	ndcg10	0.32972
# C09-1	ndcg20	0.25703
# C09-1	ERR	0.18652
# C09-1	ERR10	0.16907
# C09-1	ERR20	0.17581
# C09-1	P1	1.00000

def is_valid_file(parser, arg):
    if not os.path.exists(arg):
        parser.error("The file %s does not exist!" % arg)
    else:
        return arg

tooldescription = """
Classic bar chart indicating the mean of values for the given
metric across all queries with error bars indicating the standard
error.  (Run file is a misnomer, it should be called eval file, because it is
the output of trec_eval -c -q)
"""
parser = ArgumentParser(description=tooldescription)
parser.add_argument('--out', help='outputfilename', metavar='FILE',  required=True)
parser.add_argument('--metric', help='metric for comparison', required=True)
parser.add_argument('--format', help='trec_eval output or galago_eval output', default='trec_eval')
parser.add_argument('-c', help='instead of average, also count non-existing queries', default=False, action='store_true')
parser.add_argument('--sort', help='sort methods in plot', action='store_true', default=False)
parser.add_argument(dest='runs', nargs='+', type=lambda x: is_valid_file(parser, x))
args = parser.parse_args()

numQueries_key = "num_q"

print("column.py metric="+args.metric+" out="+args.out)

def read_ssv(fname):
    lines = [line.split() for line in open(fname, 'r')]
    if args.format.lower() == 'galago_eval':
        return lines
    elif args.format.lower() == 'trec_eval':
        return [[line[1], line[0]] + line[2:] for line in lines]

def readNumQueries(run):
    tsv = read_ssv(run)
    data = [int(row[2]) for row in tsv if row[0] == "all" and row[1] == numQueries_key]
    return data[0]

def findQueriesWithNanValues(run):
    tsv = read_ssv(run)
    # print ("tsv,", tsv)
    queriesWithNan = {row[0] for row in tsv if row[1] == 'num_rel' and (float(row[2]) == 0.0 or math.isnan(float(row[2])))}
    return queriesWithNan

def fetchValues(run):
    tsv = read_ssv(run)
    data = {row[0]: float(row[2]) for row in tsv if row[1] == args.metric and not math.isnan(float(row[2]))}
    return data

datas = {run: fetchValues(run) for run in args.runs}

# deal with nans
queriesWithNanValues = {'all'}.union(*[findQueriesWithNanValues(run) for run in args.runs])
basedata=datas[args.runs[0]]
queries = set(basedata.keys()).difference(queriesWithNanValues)
numQueries = readNumQueries(args.runs[0]) if args.c else len(queries)

seriesDict = {'mean':dict(), 'stderr':dict()}

for run in datas:
    data = datas[run]

    if sum(not key in data for key in queries) > 0:
        print("data for run "+run+" does not contain all queries "+" ".join(queries))

    mean = np.sum([data.get(key, 0.0) for key in queries]) / numQueries
    stderr = np.std([data.get(key, 0.0) for key in queries] + ([0.0]* (numQueries - len(queries)))) / sqrt(numQueries)
    seriesDict['mean'][run]=mean
    seriesDict['stderr'][run]=stderr

print( "dropping queries because of NaN values: "+ " ".join(queriesWithNanValues))

print ('\t'.join(['run', 'mean/stderr']))
for run in datas:
    #if not run == args.runs[0]:
    print ('\t'.join([run, str(seriesDict['mean'][run]), str(seriesDict['stderr'][run])]))

df1 = DataFrame(seriesDict, index=args.runs)
if args.sort:
	df1.sort_values('mean',ascending=False,inplace=True)
df2 = df1['mean']
df2.index=[os.path.basename(label) for label in df1.index]

fig, ax = plt.subplots()
#plt.figure()
df2.plot(kind='bar', yerr = df1['stderr'], color=['0.0', '0.6', '0.4', '0.8','0.0', '0.6', '0.4', '0.8','0.0', '0.6', '0.4', '0.8'], ax=ax )
ax.grid()
plt.ylabel(args.metric, fontsize=20)
plt.tick_params(axis='both', which='major', labelsize=20)
plt.xticks(rotation=90)
plt.savefig(args.out, bbox_inches='tight')

# plt.show()
<|endoftext|>