<commit_before>#Create EC2 instance, EBS, and RDS alarms for all resources in a region

import boto3
import collections

#SNS Topic Definition for EC2, EBS, and RDS
ec2_sns = 'arn:aws:sns:us-east-1:659177528321:NotifyMe'
ebs_sns = 'arn:aws:sns:us-east-1:659177528321:NotifyMe'
rds_sns = 'arn:aws:sns:us-east-1:659177528321:NotifyMe'

#AWS Account and Region Definition for Reboot Actions
akid = '659177528321'
region = 'us-east-1'

#Create AWS clients
ec = boto3.client('ec2')
rd = boto3.client('rds')
cw = boto3.client('cloudwatch')

def lambda_handler(event, context):

    #Enumerate EC2 instances
    reservations = ec.describe_instances().get('Reservations', [])
    instances = sum(
        [
            [i for i in r['Instances']]
            for r in reservations
        ], [])

    for instance in instances:
        for tag in instance['Tags']:
            if tag['Key'] == 'Name':
                name_tag = tag['Value']
                print "Found instance %s with name %s" % (instance['InstanceId'], name_tag)

        #Create Metric "CPU Utilization Greater than 95% for 15+ Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s High CPU Utilization Warning" % (name_tag, instance['InstanceId']),
            AlarmDescription='CPU Utilization Greater than 95% for 15+ Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ec2_sns,
            ],
            MetricName='CPUUtilization',
            Namespace='AWS/EC2',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'InstanceId',
                    'Value': instance['InstanceId']
                },
            ],
            Period=300,
            EvaluationPeriods=3,
            Threshold=95.0,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

        #Create Metric "CPU Utilization Greater than 95% for 60+ Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s High CPU Utilization Critical" % (name_tag, instance['InstanceId']),
            AlarmDescription='CPU Utilization Greater than 95% for 60+ Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ec2_sns,
            ],
            MetricName='CPUUtilization',
            Namespace='AWS/EC2',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'InstanceId',
                    'Value': instance['InstanceId']
                },
            ],
            Period=300,
            EvaluationPeriods=12,
            Threshold=95.0,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

        #Create CPU Credit Alarms only on T2 instances
        if "t2" in instance['InstanceType']:
            #Create Metric "CPU Credit Balance <= 25 for 30 Minutes"
            response = cw.put_metric_alarm(
                AlarmName="%s %s Credit Balance Warning" % (name_tag, instance['InstanceId']),
                AlarmDescription='CPU Credit Balance <= 25 for 30 Minutes',
                ActionsEnabled=True,
                AlarmActions=[
                    ec2_sns,
                ],
                MetricName='CPUCreditBalance',
                Namespace='AWS/EC2',
                Statistic='Average',
                Dimensions=[
                    {
                        'Name': 'InstanceId',
                        'Value': instance['InstanceId']
                    },
                ],
                Period=300,
                EvaluationPeriods=6,
                Threshold=25.0,
                ComparisonOperator='LessThanOrEqualToThreshold'
            )

            #Create Metric "CPU Credit Balance <= 5 for 10 Minutes"
            response = cw.put_metric_alarm(
                AlarmName="%s %s Credit Balance Critical" % (name_tag, instance['InstanceId']),
                AlarmDescription='CPU Credit Balance <= 5 for 10 Minutes',
                ActionsEnabled=True,
                AlarmActions=[
                    ec2_sns,
                ],
                MetricName='CPUCreditBalance',
                Namespace='AWS/EC2',
                Statistic='Average',
                Dimensions=[
                    {
                        'Name': 'InstanceId',
                        'Value': instance['InstanceId']
                    },
                ],
                Period=300,
                EvaluationPeriods=2,
                Threshold=5.0,
                ComparisonOperator='LessThanOrEqualToThreshold'
            )

        #Create Metric "Status Check Failed (System) for 5 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s System Check Failed" % (name_tag, instance['InstanceId']),
            AlarmDescription='Status Check Failed (System) for 5 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ec2_sns,
                "arn:aws:automate:%s:ec2:recover" % region,
            ],
            MetricName='StatusCheckFailed_System',
            Namespace='AWS/EC2',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'InstanceId',
                    'Value': instance['InstanceId']
                },
            ],
            Period=60,
            EvaluationPeriods=5,
            Threshold=1.0,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

        #Create Metric "Status Check Failed (Instance) for 20 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s Instance Check Failed" % (name_tag, instance['InstanceId']),
            AlarmDescription='Status Check Failed (Instance) for 20 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ec2_sns,
                "arn:aws:swf:%s:%s:action/actions/AWS_EC2.InstanceId.Reboot/1.0" % (region, akid)
            ],
            MetricName='StatusCheckFailed_Instance',
            Namespace='AWS/EC2',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'InstanceId',
                    'Value': instance['InstanceId']
                },
            ],
            Period=60,
            EvaluationPeriods=20,
            Threshold=1.0,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

    #Enumerate EBS devices attached to EC2 instances
    for instance in instances:
        for dev in instance['BlockDeviceMappings']:
            if dev.get('Ebs', None) is None:
                continue
            vol_id = dev['Ebs']['VolumeId']
            print "Found EBS volume %s on instance %s" % (
                vol_id, instance['InstanceId'])

        #Create Metric "Volume Idle Time <= 30 sec (of 5 minutes) for 30 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s High Volume Activity Warning" % (vol_id, instance['InstanceId']),
            AlarmDescription='Volume Idle Time <= 30 sec (of 5 minutes) for 30 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ebs_sns,
            ],
            MetricName='VolumeIdleTime',
            Namespace='AWS/EBS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'VolumeId',
                    'Value': vol_id
                },
            ],
            Period=300,
            EvaluationPeriods=6,
            Threshold=30.0,
            ComparisonOperator='LessThanOrEqualToThreshold'
        )

        #Create Metric "Volume Idle Time <= 30 sec (of 5 minutes) for 60 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s High Volume Activity Critical" % (vol_id, instance['InstanceId']),
            AlarmDescription='Volume Idle Time <= 30 sec (of 5 minutes) for 60 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ebs_sns,
            ],
            MetricName='VolumeIdleTime',
            Namespace='AWS/EBS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'VolumeId',
                    'Value': vol_id
                },
            ],
            Period=300,
            EvaluationPeriods=12,
            Threshold=30.0,
            ComparisonOperator='LessThanOrEqualToThreshold'
        )

    #Enumerate RDS database instances
    database_instances = rd.describe_db_instances().get('DBInstances', [])
    for database in database_instances:
        print "Found RDS database %s" % (
            database['DBInstanceIdentifier'])

        #Create CPU Credit Alarms only on T2 DB Instance Classes
        if "t2" in database['DBInstanceClass']:

            #Create Metric "CPU Credit Balance <= 25 for 30 Minutes"
            response = cw.put_metric_alarm(
                AlarmName="%s RDS Credit Balance Warning" % (database['DBInstanceIdentifier']),
                AlarmDescription='CPU Credit Balance <= 25 for 30 Minutes',
                ActionsEnabled=True,
                AlarmActions=[
                    rds_sns,
                ],
                MetricName='CPUCreditBalance',
                Namespace='AWS/RDS',
                Statistic='Average',
                Dimensions=[
                    {
                        'Name': 'DBInstanceIdentifier',
                        'Value': database['DBInstanceIdentifier']
                    },
                ],
                Period=300,
                EvaluationPeriods=6,
                Threshold=25.0,
                ComparisonOperator='LessThanOrEqualToThreshold'
            )

            #Create Metric "CPU Credit Balance <= 10 for 10 Minutes"
            response = cw.put_metric_alarm(
                AlarmName="%s RDS Credit Balance Critical" % (database['DBInstanceIdentifier']),
                AlarmDescription='CPU Credit Balance <= 10 for 10 Minutes',
                ActionsEnabled=True,
                AlarmActions=[
                    rds_sns,
                ],
                MetricName='CPUCreditBalance',
                Namespace='AWS/RDS',
                Statistic='Average',
                Dimensions=[
                    {
                        'Name': 'DBInstanceIdentifier',
                        'Value': database['DBInstanceIdentifier']
                    },
                ],
                Period=300,
                EvaluationPeriods=2,
                Threshold=10.0,
                ComparisonOperator='LessThanOrEqualToThreshold'
            )

        #Create Metric "Read Latency >= 3ms for 15 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s RDS Read Latency Warning" % (database['DBInstanceIdentifier']),
            AlarmDescription='Read Latency >= 3ms for 15 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                rds_sns,
            ],
            MetricName='ReadLatency',
            Namespace='AWS/RDS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'DBInstanceIdentifier',
                    'Value': database['DBInstanceIdentifier']
                },
            ],
            Period=300,
            EvaluationPeriods=3,
            Threshold=.003,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

        #Create Metric "Read Latency >= 5ms for 15 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s RDS Read Latency Critical" % (database['DBInstanceIdentifier']),
            AlarmDescription='Read Latency >= 5ms for 15 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                rds_sns,
            ],
            MetricName='ReadLatency',
            Namespace='AWS/RDS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'DBInstanceIdentifier',
                    'Value': database['DBInstanceIdentifier']
                },
            ],
            Period=300,
            EvaluationPeriods=3,
            Threshold=.005,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )
        #Create Metric "Write Latency >= 3ms for 15 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s RDS Write Latency Warning" % (database['DBInstanceIdentifier']),
            AlarmDescription='Write Latency >= 3ms for 15 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                rds_sns,
            ],
            MetricName='WriteLatency',
            Namespace='AWS/RDS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'DBInstanceIdentifier',
                    'Value': database['DBInstanceIdentifier']
                },
            ],
            Period=300,
            EvaluationPeriods=3,
            Threshold=.003,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

        #Create Metric "Write Latency >= 5ms for 15 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s RDS Write Latency Critical" % (database['DBInstanceIdentifier']),
            AlarmDescription='Write Latency >= 5ms for 15 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                rds_sns,
            ],
            MetricName='WriteLatency',
            Namespace='AWS/RDS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'DBInstanceIdentifier',
                    'Value': database['DBInstanceIdentifier']
                },
            ],
            Period=300,
            EvaluationPeriods=3,
            Threshold=.005,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )
<commit_msg>Remove Info<commit_after>#Create EC2 instance, EBS, and RDS alarms for all resources in a region

import boto3
import collections

#SNS Topic Definition for EC2, EBS, and RDS
ec2_sns = '<SNS_TOPIC_ARN>'
ebs_sns = '<SNS_TOPIC_ARN>'
rds_sns = '<SNS_TOPIC_ARN>'

#AWS Account and Region Definition for Reboot Actions
akid = '<ACCOUNT_ID>'
region = '<REGION_NAME>'

#Create AWS clients
ec = boto3.client('ec2')
rd = boto3.client('rds')
cw = boto3.client('cloudwatch')

def lambda_handler(event, context):

    #Enumerate EC2 instances
    reservations = ec.describe_instances().get('Reservations', [])
    instances = sum(
        [
            [i for i in r['Instances']]
            for r in reservations
        ], [])

    for instance in instances:
        for tag in instance['Tags']:
            if tag['Key'] == 'Name':
                name_tag = tag['Value']
                print "Found instance %s with name %s" % (instance['InstanceId'], name_tag)

        #Create Metric "CPU Utilization Greater than 95% for 15+ Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s High CPU Utilization Warning" % (name_tag, instance['InstanceId']),
            AlarmDescription='CPU Utilization Greater than 95% for 15+ Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ec2_sns,
            ],
            MetricName='CPUUtilization',
            Namespace='AWS/EC2',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'InstanceId',
                    'Value': instance['InstanceId']
                },
            ],
            Period=300,
            EvaluationPeriods=3,
            Threshold=95.0,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

        #Create Metric "CPU Utilization Greater than 95% for 60+ Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s High CPU Utilization Critical" % (name_tag, instance['InstanceId']),
            AlarmDescription='CPU Utilization Greater than 95% for 60+ Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ec2_sns,
            ],
            MetricName='CPUUtilization',
            Namespace='AWS/EC2',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'InstanceId',
                    'Value': instance['InstanceId']
                },
            ],
            Period=300,
            EvaluationPeriods=12,
            Threshold=95.0,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

        #Create CPU Credit Alarms only on T2 instances
        if "t2" in instance['InstanceType']:
            #Create Metric "CPU Credit Balance <= 25 for 30 Minutes"
            response = cw.put_metric_alarm(
                AlarmName="%s %s Credit Balance Warning" % (name_tag, instance['InstanceId']),
                AlarmDescription='CPU Credit Balance <= 25 for 30 Minutes',
                ActionsEnabled=True,
                AlarmActions=[
                    ec2_sns,
                ],
                MetricName='CPUCreditBalance',
                Namespace='AWS/EC2',
                Statistic='Average',
                Dimensions=[
                    {
                        'Name': 'InstanceId',
                        'Value': instance['InstanceId']
                    },
                ],
                Period=300,
                EvaluationPeriods=6,
                Threshold=25.0,
                ComparisonOperator='LessThanOrEqualToThreshold'
            )

            #Create Metric "CPU Credit Balance <= 5 for 10 Minutes"
            response = cw.put_metric_alarm(
                AlarmName="%s %s Credit Balance Critical" % (name_tag, instance['InstanceId']),
                AlarmDescription='CPU Credit Balance <= 5 for 10 Minutes',
                ActionsEnabled=True,
                AlarmActions=[
                    ec2_sns,
                ],
                MetricName='CPUCreditBalance',
                Namespace='AWS/EC2',
                Statistic='Average',
                Dimensions=[
                    {
                        'Name': 'InstanceId',
                        'Value': instance['InstanceId']
                    },
                ],
                Period=300,
                EvaluationPeriods=2,
                Threshold=5.0,
                ComparisonOperator='LessThanOrEqualToThreshold'
            )

        #Create Metric "Status Check Failed (System) for 5 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s System Check Failed" % (name_tag, instance['InstanceId']),
            AlarmDescription='Status Check Failed (System) for 5 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ec2_sns,
                "arn:aws:automate:%s:ec2:recover" % region,
            ],
            MetricName='StatusCheckFailed_System',
            Namespace='AWS/EC2',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'InstanceId',
                    'Value': instance['InstanceId']
                },
            ],
            Period=60,
            EvaluationPeriods=5,
            Threshold=1.0,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

        #Create Metric "Status Check Failed (Instance) for 20 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s Instance Check Failed" % (name_tag, instance['InstanceId']),
            AlarmDescription='Status Check Failed (Instance) for 20 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ec2_sns,
                "arn:aws:swf:%s:%s:action/actions/AWS_EC2.InstanceId.Reboot/1.0" % (region, akid)
            ],
            MetricName='StatusCheckFailed_Instance',
            Namespace='AWS/EC2',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'InstanceId',
                    'Value': instance['InstanceId']
                },
            ],
            Period=60,
            EvaluationPeriods=20,
            Threshold=1.0,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

    #Enumerate EBS devices attached to EC2 instances
    for instance in instances:
        for dev in instance['BlockDeviceMappings']:
            if dev.get('Ebs', None) is None:
                continue
            vol_id = dev['Ebs']['VolumeId']
            print "Found EBS volume %s on instance %s" % (
                vol_id, instance['InstanceId'])

        #Create Metric "Volume Idle Time <= 30 sec (of 5 minutes) for 30 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s High Volume Activity Warning" % (vol_id, instance['InstanceId']),
            AlarmDescription='Volume Idle Time <= 30 sec (of 5 minutes) for 30 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ebs_sns,
            ],
            MetricName='VolumeIdleTime',
            Namespace='AWS/EBS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'VolumeId',
                    'Value': vol_id
                },
            ],
            Period=300,
            EvaluationPeriods=6,
            Threshold=30.0,
            ComparisonOperator='LessThanOrEqualToThreshold'
        )

        #Create Metric "Volume Idle Time <= 30 sec (of 5 minutes) for 60 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s %s High Volume Activity Critical" % (vol_id, instance['InstanceId']),
            AlarmDescription='Volume Idle Time <= 30 sec (of 5 minutes) for 60 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                ebs_sns,
            ],
            MetricName='VolumeIdleTime',
            Namespace='AWS/EBS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'VolumeId',
                    'Value': vol_id
                },
            ],
            Period=300,
            EvaluationPeriods=12,
            Threshold=30.0,
            ComparisonOperator='LessThanOrEqualToThreshold'
        )

    #Enumerate RDS database instances
    database_instances = rd.describe_db_instances().get('DBInstances', [])
    for database in database_instances:
        print "Found RDS database %s" % (
            database['DBInstanceIdentifier'])

        #Create CPU Credit Alarms only on T2 DB Instance Classes
        if "t2" in database['DBInstanceClass']:

            #Create Metric "CPU Credit Balance <= 25 for 30 Minutes"
            response = cw.put_metric_alarm(
                AlarmName="%s RDS Credit Balance Warning" % (database['DBInstanceIdentifier']),
                AlarmDescription='CPU Credit Balance <= 25 for 30 Minutes',
                ActionsEnabled=True,
                AlarmActions=[
                    rds_sns,
                ],
                MetricName='CPUCreditBalance',
                Namespace='AWS/RDS',
                Statistic='Average',
                Dimensions=[
                    {
                        'Name': 'DBInstanceIdentifier',
                        'Value': database['DBInstanceIdentifier']
                    },
                ],
                Period=300,
                EvaluationPeriods=6,
                Threshold=25.0,
                ComparisonOperator='LessThanOrEqualToThreshold'
            )

            #Create Metric "CPU Credit Balance <= 10 for 10 Minutes"
            response = cw.put_metric_alarm(
                AlarmName="%s RDS Credit Balance Critical" % (database['DBInstanceIdentifier']),
                AlarmDescription='CPU Credit Balance <= 10 for 10 Minutes',
                ActionsEnabled=True,
                AlarmActions=[
                    rds_sns,
                ],
                MetricName='CPUCreditBalance',
                Namespace='AWS/RDS',
                Statistic='Average',
                Dimensions=[
                    {
                        'Name': 'DBInstanceIdentifier',
                        'Value': database['DBInstanceIdentifier']
                    },
                ],
                Period=300,
                EvaluationPeriods=2,
                Threshold=10.0,
                ComparisonOperator='LessThanOrEqualToThreshold'
            )

        #Create Metric "Read Latency >= 3ms for 15 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s RDS Read Latency Warning" % (database['DBInstanceIdentifier']),
            AlarmDescription='Read Latency >= 3ms for 15 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                rds_sns,
            ],
            MetricName='ReadLatency',
            Namespace='AWS/RDS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'DBInstanceIdentifier',
                    'Value': database['DBInstanceIdentifier']
                },
            ],
            Period=300,
            EvaluationPeriods=3,
            Threshold=.003,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

        #Create Metric "Read Latency >= 5ms for 15 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s RDS Read Latency Critical" % (database['DBInstanceIdentifier']),
            AlarmDescription='Read Latency >= 5ms for 15 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                rds_sns,
            ],
            MetricName='ReadLatency',
            Namespace='AWS/RDS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'DBInstanceIdentifier',
                    'Value': database['DBInstanceIdentifier']
                },
            ],
            Period=300,
            EvaluationPeriods=3,
            Threshold=.005,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )
        #Create Metric "Write Latency >= 3ms for 15 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s RDS Write Latency Warning" % (database['DBInstanceIdentifier']),
            AlarmDescription='Write Latency >= 3ms for 15 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                rds_sns,
            ],
            MetricName='WriteLatency',
            Namespace='AWS/RDS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'DBInstanceIdentifier',
                    'Value': database['DBInstanceIdentifier']
                },
            ],
            Period=300,
            EvaluationPeriods=3,
            Threshold=.003,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )

        #Create Metric "Write Latency >= 5ms for 15 Minutes"
        response = cw.put_metric_alarm(
            AlarmName="%s RDS Write Latency Critical" % (database['DBInstanceIdentifier']),
            AlarmDescription='Write Latency >= 5ms for 15 Minutes',
            ActionsEnabled=True,
            AlarmActions=[
                rds_sns,
            ],
            MetricName='WriteLatency',
            Namespace='AWS/RDS',
            Statistic='Average',
            Dimensions=[
                {
                    'Name': 'DBInstanceIdentifier',
                    'Value': database['DBInstanceIdentifier']
                },
            ],
            Period=300,
            EvaluationPeriods=3,
            Threshold=.005,
            ComparisonOperator='GreaterThanOrEqualToThreshold'
        )
<|endoftext|>


<commit_before>#
#   Cython -- Things that don't belong
#            anywhere else in particular
#

import os, sys, re, codecs

def replace_suffix(path, newsuf):
    base, _ = os.path.splitext(path)
    return base + newsuf

def open_new_file(path):
    #  Open and truncate existing file to
    #  preserve metadata on the Mac.
    return open(path, "w+")

def castrate_file(path, st):
    #  Remove junk contents from an output file after a
    #  failed compilation, but preserve metadata on Mac.
    #  Also sets access and modification times back to
    #  those specified by st (a stat struct).
    try:
        f = open(path, "r+")
    except EnvironmentError:
        pass
    else:
        f.seek(0, 0)
        f.truncate()
        f.write(
            "#error Do not use this file, it is the result of a failed Cython compilation.\n")
        f.close()
        if st:
            os.utime(path, (st.st_atime, st.st_mtime-1))

def modification_time(path):
    st = os.stat(path)
    return st.st_mtime

def file_newer_than(path, time):
    ftime = modification_time(path)
    return ftime > time

# support for source file encoding detection

def encode_filename(filename):
    if isinstance(filename, unicode):
        return filename
    try:
        filename_encoding = sys.getfilesystemencoding()
        if filename_encoding is None:
            filename_encoding = sys.getdefaultencoding()
        filename = filename.decode(filename_encoding)
    except UnicodeDecodeError:
        pass
    return filename

_match_file_encoding = re.compile(u"coding[:=]\s*([-\w.]+)").search

def detect_file_encoding(source_filename):
    # PEPs 263 and 3120
    f = codecs.open(source_filename, "rU", encoding="UTF-8")
    try:
        chars = []
        for i in range(2):
            c = f.read(1)
            while c and c != u'\n':
                chars.append(c)
                c = f.read(1)
            encoding = _match_file_encoding(u''.join(chars))
            if encoding:
                return encoding.group(1)
    finally:
        f.close()
    return "UTF-8"

def open_source_file(source_filename, mode="rU"):
    encoding = detect_file_encoding(source_filename)
    return codecs.open(source_filename, mode=mode, encoding=encoding)

def long_literal(value):
    if isinstance(value, basestring):
        if len(value) < 2:
            value = int(value)
        elif value[0] == 0:
            value = int(value, 8)
        elif value[1] in 'xX':
            value = int(value[2:], 16)
        else:
            value = int(value)
    return not -2**31 <= value < 2**31

# a simple class that simplifies the usage of utility code

class UtilityCode(object):
    def __init__(self, proto=None, impl=None, init=None, cleanup=None):
        self.proto = proto
        self.impl = impl
        self.init = init
        self.cleanup = cleanup

    def write_init_code(self, writer, pos):
        if not self.init:
            return
        if callable(self.init):
            self.init(writer, pos)
        else:
            writer.put(self.init)

    def write_cleanup_code(self, writer, pos):
        if not self.cleanup:
            return
        if callable(self.cleanup):
            self.cleanup(writer, pos)
        else:
            writer.put(self.cleanup)
<commit_msg>Create new c file to play better with hard links (Trac #220)<commit_after>#
#   Cython -- Things that don't belong
#            anywhere else in particular
#

import os, sys, re, codecs

def replace_suffix(path, newsuf):
    base, _ = os.path.splitext(path)
    return base + newsuf

def open_new_file(path):
    if os.path.exists(path):
        # Make sure to create a new file here so we can
        # safely hard link the output files.
        os.unlink(path)
    return open(path, "w")

def castrate_file(path, st):
    #  Remove junk contents from an output file after a
    #  failed compilation.
    #  Also sets access and modification times back to
    #  those specified by st (a stat struct).
    try:
        f = open_new_file(path)
    except EnvironmentError:
        pass
    else:
        f.write(
            "#error Do not use this file, it is the result of a failed Cython compilation.\n")
        f.close()
        if st:
            os.utime(path, (st.st_atime, st.st_mtime-1))

def modification_time(path):
    st = os.stat(path)
    return st.st_mtime

def file_newer_than(path, time):
    ftime = modification_time(path)
    return ftime > time

# support for source file encoding detection

def encode_filename(filename):
    if isinstance(filename, unicode):
        return filename
    try:
        filename_encoding = sys.getfilesystemencoding()
        if filename_encoding is None:
            filename_encoding = sys.getdefaultencoding()
        filename = filename.decode(filename_encoding)
    except UnicodeDecodeError:
        pass
    return filename

_match_file_encoding = re.compile(u"coding[:=]\s*([-\w.]+)").search

def detect_file_encoding(source_filename):
    # PEPs 263 and 3120
    f = codecs.open(source_filename, "rU", encoding="UTF-8")
    try:
        chars = []
        for i in range(2):
            c = f.read(1)
            while c and c != u'\n':
                chars.append(c)
                c = f.read(1)
            encoding = _match_file_encoding(u''.join(chars))
            if encoding:
                return encoding.group(1)
    finally:
        f.close()
    return "UTF-8"

def open_source_file(source_filename, mode="rU"):
    encoding = detect_file_encoding(source_filename)
    return codecs.open(source_filename, mode=mode, encoding=encoding)

def long_literal(value):
    if isinstance(value, basestring):
        if len(value) < 2:
            value = int(value)
        elif value[0] == 0:
            value = int(value, 8)
        elif value[1] in 'xX':
            value = int(value[2:], 16)
        else:
            value = int(value)
    return not -2**31 <= value < 2**31

# a simple class that simplifies the usage of utility code

class UtilityCode(object):
    def __init__(self, proto=None, impl=None, init=None, cleanup=None):
        self.proto = proto
        self.impl = impl
        self.init = init
        self.cleanup = cleanup

    def write_init_code(self, writer, pos):
        if not self.init:
            return
        if callable(self.init):
            self.init(writer, pos)
        else:
            writer.put(self.init)

    def write_cleanup_code(self, writer, pos):
        if not self.cleanup:
            return
        if callable(self.cleanup):
            self.cleanup(writer, pos)
        else:
            writer.put(self.cleanup)
<|endoftext|>


<commit_before><commit_msg>Update for cases where year & month are in the index<commit_after><|endoftext|>


<commit_before>##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2013 Stanford University and the Authors
#
# Authors: Gregory Bowman
# Contributors: Robert McGibbon
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
##############################################################################

import bz2
import glob
import gzip
import os
import os.path
import tables
import tempfile
import unittest

from mdtraj.testing import eq
from mdtraj import io

import numpy as np

fd, temp = tempfile.mkstemp(suffix='.h5')
os.close(fd)

def teardown_module(module):
    """remove the temporary file created by tests in this file
    this gets automatically called by nose"""
    os.unlink(temp)

def test_overwrite_1():
    fid, fn = tempfile.mkstemp()
    try:
        a = np.arange(10)
        b = a + 1
        io.saveh(fn, a=a)
        io.saveh(fn, b=b)
        eq(io.loadh(fn, 'a'), a)
        eq(io.loadh(fn, 'b'), b)
    except:
        raise
    finally:
        if os.path.exists(fn):
            os.close(fid)
            os.unlink(fn)

def test_overwrite_2():
    fid, fn = tempfile.mkstemp()
    try:
        a = np.arange(10)
        b = a + 1
        io.saveh(fn, a=a)
        io.saveh(fn, a=b)
        eq(io.loadh(fn, 'a'), b)
    except:
        raise
    finally:
        if os.path.exists(fn):
            os.close(fid)
            os.unlink(fn)

class test_io(unittest.TestCase):
    def setUp(self):
        """setup() is called before very test and just creates a temporary work space for reading/writing files."""
        fid, self.filename1 = tempfile.mkstemp()
        fid, self.filename2 = tempfile.mkstemp()
        self.data = np.arange(10000, dtype=np.float32)

        #Write Data to an HDF5 file as a compressed CArray.
        hdfFile = tables.File(self.filename1, 'a')
        #The filter is the same used to save MSMB2 data
        hdfFile.createCArray("/", "arr_0", tables.Float32Atom(), self.data.shape, filters=io.COMPRESSION)
        hdfFile.root.arr_0[:] = self.data[:]
        hdfFile.flush()
        hdfFile.close()

    def test_load_1(self):
        "Load by specifying array name"
        TestData = io.loadh(self.filename1, 'arr_0')
        eq(TestData, self.data)

    def test_load_2(self):
        "load using deferred=False"
        TestData = io.loadh(self.filename1, deferred=False)['arr_0']
        eq(TestData, self.data)

    def test_load_2(self):
        "load using deferred=True"
        deferred = io.loadh(self.filename1, deferred=True)
        eq(deferred['arr_0'], self.data)
        deferred.close()

    def test_save(self):
        """Save HDF5 to disk and load it back up"""
        io.saveh(self.filename2, self.data)
        TestData = io.loadh(self.filename2, 'arr_0')
        eq(TestData, self.data)

    def teardown(self):
        os.remove(self.filename1)
        os.remove(self.filename2)

class test_io_int(test_io):
    "Run the same test as the class above, but using int64 data"
    def setUp(self):
        """setup() is called before very test and just creates a temporary work space for reading/writing files."""
        fid, self.filename1 = tempfile.mkstemp()
        fid, self.filename2 = tempfile.mkstemp()
        self.data = np.arange(10000, dtype=np.int64)

        #Write Data to an HDF5 file as a compressed CArray.
        hdfFile = tables.File(self.filename1, 'a')
        #The filter is the same used to save MSMB2 data
        hdfFile.createCArray("/", "arr_0", tables.Int64Atom(), self.data.shape, filters=io.COMPRESSION)
        hdfFile.root.arr_0[:] = self.data[:]
        hdfFile.flush()
        hdfFile.close()

def test_groups():
    """Test to ensure that files are loaded correctly even if they contain nested
    groups and stuff"""
    x = np.random.randn(10)
    y = np.random.randn(11)
    f = tables.openFile(temp, 'w')
    f.createGroup(where='/', name='mygroup')
    if tables.__version__ >= '3.0.0':
        f.createArray(where='/mygroup', name='myarray', obj=x)
        f.createArray(where='/', name='mya2', obj=y)
    else:
        f.createArray(where='/mygroup', name='myarray', object=x)
        f.createArray(where='/', name='mya2', object=y)
    f.close()

    yield lambda: eq(io.loadh(temp)['mygroup/myarray'], x)
    yield lambda: eq(io.loadh(temp)['mya2'], y)
    yield lambda: eq(io.loadh(temp, deferred=False)['mygroup/myarray'], x)
    yield lambda: eq(io.loadh(temp, deferred=False)['mya2'], y)
    yield lambda: eq(io.loadh(temp, 'mygroup/myarray'), x)
    yield lambda: eq(io.loadh(temp, 'mya2'), y)
<commit_msg>Modernize test_io.py<commit_after>##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2013 Stanford University and the Authors
#
# Authors: Gregory Bowman
# Contributors: Robert McGibbon
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
##############################################################################

import os
import tempfile
import unittest

import numpy as np
import tables

from mdtraj import io
from mdtraj.testing import eq

fd, temp = tempfile.mkstemp(suffix='.h5')
os.close(fd)

def teardown_module(module):
    """remove the temporary file created by tests in this file
    this gets automatically called by nose"""
    os.unlink(temp)

def test_overwrite_1():
    fid, fn = tempfile.mkstemp()
    try:
        a = np.arange(10)
        b = a + 1
        io.saveh(fn, a=a)
        io.saveh(fn, b=b)
        eq(io.loadh(fn, 'a'), a)
        eq(io.loadh(fn, 'b'), b)
    except:
        raise
    finally:
        if os.path.exists(fn):
            os.close(fid)
            os.unlink(fn)

def test_overwrite_2():
    fid, fn = tempfile.mkstemp()
    try:
        a = np.arange(10)
        b = a + 1
        io.saveh(fn, a=a)
        io.saveh(fn, a=b)
        eq(io.loadh(fn, 'a'), b)
    except:
        raise
    finally:
        if os.path.exists(fn):
            os.close(fid)
            os.unlink(fn)

class test_io(unittest.TestCase):
    def setUp(self):
        # setup() is called before very test and just creates a
        # temporary work space for reading/writing files.
        fid, self.filename1 = tempfile.mkstemp()
        fid, self.filename2 = tempfile.mkstemp()
        self.data = np.arange(10000, dtype=np.float32)

        # Write Data to an HDF5 file as a compressed CArray.
        hdf_file = tables.open_file(self.filename1, 'a')
        hdf_file.create_carray("/", "arr_0", tables.Float32Atom(),
                               self.data.shape, filters=io.COMPRESSION)
        hdf_file.root.arr_0[:] = self.data[:]
        hdf_file.flush()
        hdf_file.close()

    def test_load_1(self):
        # Load by specifying array name
        TestData = io.loadh(self.filename1, 'arr_0')
        eq(TestData, self.data)

    def test_load_2(self):
        # load using deferred=False
        TestData = io.loadh(self.filename1, deferred=False)['arr_0']
        eq(TestData, self.data)

    def test_load_3(self):
        # load using deferred=True
        deferred = io.loadh(self.filename1, deferred=True)
        eq(deferred['arr_0'], self.data)
        deferred.close()

    def test_save(self):
        # Save HDF5 to disk and load it back up
        io.saveh(self.filename2, self.data)
        TestData = io.loadh(self.filename2, 'arr_0')
        eq(TestData, self.data)

    def teardown(self):
        os.remove(self.filename1)
        os.remove(self.filename2)

class test_io_int(test_io):
    "Run the same test as the class above, but using int64 data"

    def setUp(self):
        # setup() is called before very test and just creates
        # a temporary work space for reading/writing files.
        fid, self.filename1 = tempfile.mkstemp()
        fid, self.filename2 = tempfile.mkstemp()
        self.data = np.arange(10000, dtype=np.int64)

        # Write Data to an HDF5 file as a compressed CArray.
        hdf_file = tables.open_file(self.filename1, 'a')
        hdf_file.create_carray("/", "arr_0", tables.Int64Atom(),
                               self.data.shape, filters=io.COMPRESSION)
        hdf_file.root.arr_0[:] = self.data[:]
        hdf_file.flush()
        hdf_file.close()

def test_groups():
    # Test to ensure that files are loaded correctly even if they contain
    # nested groups and stuff
    x = np.random.randn(10)
    y = np.random.randn(11)
    f = tables.open_file(temp, 'w')
    f.create_group(where='/', name='mygroup')
    f.create_array(where='/mygroup', name='myarray', obj=x)
    f.create_array(where='/', name='mya2', obj=y)
    f.close()

    yield lambda: eq(io.loadh(temp)['mygroup/myarray'], x)
    yield lambda: eq(io.loadh(temp)['mya2'], y)
    yield lambda: eq(io.loadh(temp, deferred=False)['mygroup/myarray'], x)
    yield lambda: eq(io.loadh(temp, deferred=False)['mya2'], y)
    yield lambda: eq(io.loadh(temp, 'mygroup/myarray'), x)
    yield lambda: eq(io.loadh(temp, 'mya2'), y)
<|endoftext|>


<commit_before>"""Convenient HTTP UserAgent class.

This is a subclass of urllib2.OpenerDirector.

Copyright 2003 John J. Lee <jjl@pobox.com>

This code is free software; you can redistribute it and/or modify it under
the terms of the BSD License (see the file COPYING included with the
distribution).

"""

import sys
import urllib2, httplib
import ClientCookie
if sys.version_info[:2] >= (2, 4):
    from urllib2 import OpenerDirector, BaseHandler, \
         HTTPHandler, HTTPSHandler, HTTPErrorProcessor
    class SaneHTTPCookieProcessor(ClientCookie.HTTPCookieProcessor):
        # workaround for RFC 2109 bug (at least if you don't pass your own
        # CookieJar in...)
        def __init__(self, cookiejar=None):
            if cookiejar is None:
                cookiejar = CookieJar(DefaultCookiePolicy(rfc2965=True))
            self.cookiejar = cookiejar
    HTTPCookieProcessor = SaneHTTPCookieProcessor
else:
    from ClientCookie import OpenerDirector, BaseHandler, \
         HTTPHandler, HTTPSHandler, HTTPErrorProcessor, HTTPCookieProcessor

class HTTPRefererProcessor(BaseHandler):
    def http_request(self, request):
        # See RFC 2616 14.36.  The only times we know the source of the
        # request URI has a URI associated with it are redirect, and
        # Browser.click() / Browser.submit() / Browser.follow_link().
        # Otherwise, it's the user's job to add any Referer header before
        # .open()ing.
        if hasattr(request, "redirect_dict"):
            request = self.parent._add_referer_header(
                request, origin_request=False)
        return request

    https_request = http_request

class UserAgent(OpenerDirector):
    """Convenient user-agent class.

    Do not use .add_handler() to add a handler for something already dealt with
    by this code.

    Public attributes:

    addheaders: list of (name, value) pairs specifying headers to send with
     every request, unless they are overridden in the Request instance.

     >>> ua = UserAgent()
     >>> ua.addheaders = [
     ...  ("User-agent", "Mozilla/5.0 (compatible)"),
     ...  ("From", "responsible.person@example.com")]

    """

    handler_classes = {
        # scheme handlers
        "http": HTTPHandler,
        "ftp": urllib2.FTPHandler,  # CacheFTPHandler is buggy in 2.3
        "file": urllib2.FileHandler,
        "gopher": urllib2.GopherHandler,
        # XXX etc.

        # other handlers
        "_unknown": urllib2.UnknownHandler,
        # HTTP{S,}Handler depend on HTTPErrorProcessor too
        "_http_error": HTTPErrorProcessor,
        "_http_request_upgrade": ClientCookie.HTTPRequestUpgradeProcessor,
        "_http_default_error": urllib2.HTTPDefaultErrorHandler,

        # feature handlers
        "_authen": urllib2.HTTPBasicAuthHandler,
        # XXX rest of authentication stuff
        "_redirect": ClientCookie.HTTPRedirectHandler,
        "_cookies": HTTPCookieProcessor,
        "_refresh": ClientCookie.HTTPRefreshProcessor,
        "_referer": HTTPRefererProcessor,  # from this module, note
        "_equiv": ClientCookie.HTTPEquivProcessor,
        "_seek": ClientCookie.SeekableProcessor,
        "_proxy": urllib2.ProxyHandler,
        # XXX there's more to proxies, too

        # debug handlers
        "_debug_redirect": ClientCookie.HTTPRedirectDebugProcessor,
        "_debug_response_body": ClientCookie.HTTPResponseDebugProcessor,
        }

    default_schemes = ["http", "ftp", "file", "gopher"]
    default_others = ["_unknown", "_http_error", "_http_request_upgrade",
                      "_http_default_error"]
    default_features = ["_authen", "_redirect", "_cookies", "_seek", "_proxy"]
    if hasattr(httplib, 'HTTPS'):
        handler_classes["https"] = HTTPSHandler
        default_schemes.append("https")
    if hasattr(ClientCookie, "HTTPRobotRulesProcessor"):
        handler_classes["_robots"] = ClientCookie.HTTPRobotRulesProcessor
        default_features.append("_robots")

    def __init__(self):
        OpenerDirector.__init__(self)

        self._ua_handlers = {}
        for scheme in (self.default_schemes+
                       self.default_others+
                       self.default_features):
            klass = self.handler_classes[scheme]
            self._ua_handlers[scheme] = klass()
        for handler in self._ua_handlers.itervalues():
            self.add_handler(handler)

        # special case, requires extra support from mechanize.Browser
        self._handle_referer = True

    def close(self):
        OpenerDirector.close(self)
        self._ua_handlers = None

    # XXX
##     def set_timeout(self, timeout):
##         self._timeout = timeout
##     def set_http_connection_cache(self, conn_cache):
##         self._http_conn_cache = conn_cache
##     def set_ftp_connection_cache(self, conn_cache):
##         # XXX ATM, FTP has cache as part of handler; should it be separate?
##         self._ftp_conn_cache = conn_cache

    def set_handled_schemes(self, schemes):
        """Set sequence of protocol scheme strings.

        If this fails (with ValueError) because you've passed an unknown
        scheme, the set of handled schemes WILL be updated, but schemes in the
        list that come after the unknown scheme won't be handled.

        """
        want = {}
        for scheme in schemes:
            if scheme.startswith("_"):
                raise ValueError("invalid scheme '%s'" % scheme)
            want[scheme] = None

        # get rid of scheme handlers we don't want
        for scheme, oldhandler in self._ua_handlers.items():
            if scheme.startswith("_"): continue  # not a scheme handler
            if scheme not in want:
                self._replace_handler(scheme, None)
            else:
                del want[scheme]  # already got it
        # add the scheme handlers that are missing
        for scheme in want.keys():
            if scheme not in self.handler_classes:
                raise ValueError("unknown scheme '%s'")
            self._set_handler(scheme, True)

    def _add_referer_header(self, request):
        raise NotImplementedError(
            "this class can't do HTTP Referer: use mechanize.Browser instead")

    def set_cookiejar(self, cookiejar):
        """Set a ClientCookie.CookieJar, or None."""
        self._set_handler("_cookies", obj=cookiejar)
    def set_credentials(self, credentials):
        """Set a urllib2.HTTPPasswordMgr, or None."""
        # XXX use Greg Stein's httpx instead?
        self._set_handler("_authen", obj=credentials)

    # these methods all take a boolean parameter
    def set_handle_robots(self, handle):
        """Set whether to observe rules from robots.txt."""
        self._set_handler("_robots", handle)
    def set_handle_redirect(self, handle):
        """Set whether to handle HTTP Refresh headers."""
        self._set_handler("_redirect", handle)
    def set_handle_refresh(self, handle):
        """Set whether to handle HTTP Refresh headers."""
        self._set_handler("_refresh", handle)
    def set_handle_equiv(self, handle):
        """Set whether to treat HTML http-equiv headers like HTTP headers.

        Response objects will be .seek()able if this is set.

        """
        self._set_handler("_equiv", handle)
    def set_handle_referer(self, handle):
        """Set whether to add Referer header to each request.

        This base class does not implement this feature (so don't turn this on
        if you're using this base class directly), but the subclass
        mechanize.Browser does.

        """
        self._set_handler("_referer", handle)
        self._handle_referer = True
    def set_seekable_responses(self, handle):
        """Make response objects .seek()able."""
        self._set_handler("_seek", handle)
    def set_debug_redirects(self, handle):
        """Print information about HTTP redirects.

        This includes refreshes, which show up as faked 302 redirections at the
        moment.

        """
        self._set_handler("_debug_redirect", handle)
    def set_debug_responses(self, handle):
        """Print HTTP response bodies."""
        self._set_handler("_debug_response_body", handle)
    def set_debug_http(self, handle):
        """Print HTTP headers."""
        level = int(bool(handle))
        for scheme in "http", "https":
            h = self._ua_handlers.get(scheme)
            if h is not None:
                h.set_http_debuglevel(level)

    def _set_handler(self, name, handle=None, obj=None):
        if handle is None:
            handle = obj is not None
        if handle:
            handler_class = self.handler_classes[name]
            if obj is not None:
                newhandler = handler_class(obj)
            else:
                newhandler = handler_class()
        else:
            newhandler = None
        self._replace_handler(name, newhandler)

    # XXXX I'd *really* rather get rid of this and just rebuild every time.
    #  This is fragile to base class changes, and hard to understand.
    #  Have to make sure there's no state directly stored in handlers, though,
    #  and have appropriate methods for adding state back to the cookie etc.
    #  handlers known to this class (only the ones in urllib2 / ClientCookie --
    #  no need to care about other peoples' as long as it's documented that
    #  calling the set_* methods will in general clobber handler state).
    def _replace_handler(self, name, newhandler=None):
        # first, if handler was previously added, remove it
        if name is not None:
            try:
                handler = self._ua_handlers[name]
            except:
                pass
            else:
                for table in (
                    [self.handle_open,
                     self.process_request, self.process_response]+
                    self.handle_error.values()):
                    for handlers in table.values():
                        remove(handlers, handler)
                    remove(self.handlers, handler)
        # then add the replacement, if any
        if newhandler is not None:
            self.add_handler(newhandler)
            self._ua_handlers[name] = newhandler

def remove(sequence, obj):
    # for use when can't use .remove() because of obj.__cmp__ :-(
    # (ClientCookie only requires Python 2.0, which doesn't have __lt__)
    i = 0
    while i < len(sequence):
        if sequence[i] is obj:
            del sequence[i]
        else:
            i += 1

# XXX
# This is urllib2.Request with a new .set_method() method,
# for HTTP HEAD / PUT -- move into ClientCookie if/when need it.
# Maybe it should have a constructor arg, too.
## class Request:

##     def __init__(self, url, data=None, headers={}):
##         # unwrap('<URL:type://host/path>') --> 'type://host/path'
##         self.__original = unwrap(url)
##         self.type = None
##         # self.__r_type is what's left after doing the splittype
##         self.host = None
##         self.port = None
##         self.data = data
##         self.headers = {}
##         for key, value in headers.items():
##             self.add_header(key, value)
##         if data is None:
##             self._method = "GET"
##         else:
##             self._method = "POST"

##     def __getattr__(self, attr):
##         # XXX this is a fallback mechanism to guard against these
##         # methods getting called in a non-standard order.  this may be
##         # too complicated and/or unnecessary.
##         # XXX should the __r_XXX attributes be public?
##         if attr[:12] == '_Request__r_':
##             name = attr[12:]
##             if hasattr(Request, 'get_' + name):
##                 getattr(self, 'get_' + name)()
##                 return getattr(self, attr)
##         raise AttributeError, attr

##     def get_method(self):
##         return self._method

##     def set_method(self, method):
##         if method == "POST":
##             if data is None:
##                 data = ""
##         else:
##             self.data = None
##         self._method == method

##     def add_data(self, data):
##         self.data = data

##     def has_data(self):
##         return self.data is not None

##     def get_data(self):
##         return self.data

##     def get_full_url(self):
##         return self.__original

##     def get_type(self):
##         if self.type is None:
##             self.type, self.__r_type = splittype(self.__original)
##             if self.type is None:
##                 raise ValueError, "unknown url type: %s" % self.__original
##         return self.type

##     def get_host(self):
##         if self.host is None:
##             self.host, self.__r_host = splithost(self.__r_type)
##             if self.host:
##                 self.host = unquote(self.host)
##         return self.host

##     def get_selector(self):
##         return self.__r_host

##     def set_proxy(self, host, type):
##         self.host, self.type = host, type
##         self.__r_host = self.__original

##     def add_header(self, key, val):
##         # useful for something like authentication
##         self.headers[key.capitalize()] = val

## def http_get(fullurl, ranges=None, conditions=None):
##     """HTTP GET, with convenient partial fetches (ranges).

##     XXX conditional fetches?

##     ranges: sequence of pairs of byte ranges (start, end) to fetch;

##     Ranges follow the usual Python rules (the start byte is included,
##     the end byte is not; negative numbers count back from the end of
##     the entity; start None means start of entity; end None means end of
##     entity).  There are restrictions, though: end must not be negative,
##     and if start is negative, end must be None.

##     >>> http_get("http://www.example.com/big.dat",
##                  [(0, 10), (-10, None)])  # first and last 10 bytes
##     >>> http_get("http://www.example.com/big.dat",
##                  [(50000, None)])  # from byte 50000 to the end

##     """
##     if conditions: raise NotImplementedError("conditions not yet implemented")
##     req = self._request(fullurl, data)
##     assert req.get_type() == "http", "http_get for non-HTTP URI"
##     rs = []
##     for start, end in ranges:
##         if start < 0:
##             assert end is None, "invalid range"
##             start = ""
##         else:
##             assert 0 <= start <= end, "invalid range"
##             if start == end: continue
##             end = end - 1
##         rs.append("%s-%s" % range)
##     req.add_header(("Range", "bytes=" % string.join(rs, ", ")))
##     return self.open(req)

## def http_head(self, fullurl):
##     raise NotImplementedError()  # XXX

## def http_put(self, fullurl, data=None):
##     # XXX what about 30x handling?
##     raise NotImplementedError()  # XXX
<commit_msg>Turn on Refresh, Referer and HTTP-Equiv handling by default; Clarify a comment<commit_after>"""Convenient HTTP UserAgent class.

This is a subclass of urllib2.OpenerDirector.

Copyright 2003 John J. Lee <jjl@pobox.com>

This code is free software; you can redistribute it and/or modify it under
the terms of the BSD License (see the file COPYING included with the
distribution).

"""

import sys
import urllib2, httplib
import ClientCookie
if sys.version_info[:2] >= (2, 4):
    from urllib2 import OpenerDirector, BaseHandler, \
         HTTPHandler, HTTPSHandler, HTTPErrorProcessor
    class SaneHTTPCookieProcessor(ClientCookie.HTTPCookieProcessor):
        # Workaround for RFC 2109 bug http://python.org/sf/1157027 (at least if
        # you don't pass your own CookieJar in: if that's the case, you should
        # pass rfc2965=True to the DefaultCookiePolicy constructor yourself, or
        # set the corresponding attribute).
        def __init__(self, cookiejar=None):
            if cookiejar is None:
                cookiejar = CookieJar(DefaultCookiePolicy(rfc2965=True))
            self.cookiejar = cookiejar
    HTTPCookieProcessor = SaneHTTPCookieProcessor
else:
    from ClientCookie import OpenerDirector, BaseHandler, \
         HTTPHandler, HTTPSHandler, HTTPErrorProcessor, HTTPCookieProcessor

class HTTPRefererProcessor(BaseHandler):
    def http_request(self, request):
        # See RFC 2616 14.36.  The only times we know the source of the
        # request URI has a URI associated with it are redirect, and
        # Browser.click() / Browser.submit() / Browser.follow_link().
        # Otherwise, it's the user's job to add any Referer header before
        # .open()ing.
        if hasattr(request, "redirect_dict"):
            request = self.parent._add_referer_header(
                request, origin_request=False)
        return request

    https_request = http_request

class UserAgent(OpenerDirector):
    """Convenient user-agent class.

    Do not use .add_handler() to add a handler for something already dealt with
    by this code.

    Public attributes:

    addheaders: list of (name, value) pairs specifying headers to send with
     every request, unless they are overridden in the Request instance.

     >>> ua = UserAgent()
     >>> ua.addheaders = [
     ...  ("User-agent", "Mozilla/5.0 (compatible)"),
     ...  ("From", "responsible.person@example.com")]

    """

    handler_classes = {
        # scheme handlers
        "http": HTTPHandler,
        "ftp": urllib2.FTPHandler,  # CacheFTPHandler is buggy in 2.3
        "file": urllib2.FileHandler,
        "gopher": urllib2.GopherHandler,
        # XXX etc.

        # other handlers
        "_unknown": urllib2.UnknownHandler,
        # HTTP{S,}Handler depend on HTTPErrorProcessor too
        "_http_error": HTTPErrorProcessor,
        "_http_request_upgrade": ClientCookie.HTTPRequestUpgradeProcessor,
        "_http_default_error": urllib2.HTTPDefaultErrorHandler,

        # feature handlers
        "_authen": urllib2.HTTPBasicAuthHandler,
        # XXX rest of authentication stuff
        "_redirect": ClientCookie.HTTPRedirectHandler,
        "_cookies": HTTPCookieProcessor,
        "_refresh": ClientCookie.HTTPRefreshProcessor,
        "_referer": HTTPRefererProcessor,  # from this module, note
        "_equiv": ClientCookie.HTTPEquivProcessor,
        "_seek": ClientCookie.SeekableProcessor,
        "_proxy": urllib2.ProxyHandler,
        # XXX there's more to proxies, too

        # debug handlers
        "_debug_redirect": ClientCookie.HTTPRedirectDebugProcessor,
        "_debug_response_body": ClientCookie.HTTPResponseDebugProcessor,
        }

    default_schemes = ["http", "ftp", "file", "gopher"]
    default_others = ["_unknown", "_http_error", "_http_request_upgrade",
                      "_http_default_error"]
    default_features = ["_authen", "_redirect", "_cookies", "_refresh",
                        "_referer", "equiv", "_seek", "_proxy"]
    if hasattr(httplib, 'HTTPS'):
        handler_classes["https"] = HTTPSHandler
        default_schemes.append("https")
    if hasattr(ClientCookie, "HTTPRobotRulesProcessor"):
        handler_classes["_robots"] = ClientCookie.HTTPRobotRulesProcessor
        default_features.append("_robots")

    def __init__(self):
        OpenerDirector.__init__(self)

        self._ua_handlers = {}
        for scheme in (self.default_schemes+
                       self.default_others+
                       self.default_features):
            klass = self.handler_classes[scheme]
            self._ua_handlers[scheme] = klass()
        for handler in self._ua_handlers.itervalues():
            self.add_handler(handler)

        # special case, requires extra support from mechanize.Browser
        self._handle_referer = True

    def close(self):
        OpenerDirector.close(self)
        self._ua_handlers = None

    # XXX
##     def set_timeout(self, timeout):
##         self._timeout = timeout
##     def set_http_connection_cache(self, conn_cache):
##         self._http_conn_cache = conn_cache
##     def set_ftp_connection_cache(self, conn_cache):
##         # XXX ATM, FTP has cache as part of handler; should it be separate?
##         self._ftp_conn_cache = conn_cache

    def set_handled_schemes(self, schemes):
        """Set sequence of protocol scheme strings.

        If this fails (with ValueError) because you've passed an unknown
        scheme, the set of handled schemes WILL be updated, but schemes in the
        list that come after the unknown scheme won't be handled.

        """
        want = {}
        for scheme in schemes:
            if scheme.startswith("_"):
                raise ValueError("invalid scheme '%s'" % scheme)
            want[scheme] = None

        # get rid of scheme handlers we don't want
        for scheme, oldhandler in self._ua_handlers.items():
            if scheme.startswith("_"): continue  # not a scheme handler
            if scheme not in want:
                self._replace_handler(scheme, None)
            else:
                del want[scheme]  # already got it
        # add the scheme handlers that are missing
        for scheme in want.keys():
            if scheme not in self.handler_classes:
                raise ValueError("unknown scheme '%s'")
            self._set_handler(scheme, True)

    def _add_referer_header(self, request):
        raise NotImplementedError(
            "this class can't do HTTP Referer: use mechanize.Browser instead")

    def set_cookiejar(self, cookiejar):
        """Set a ClientCookie.CookieJar, or None."""
        self._set_handler("_cookies", obj=cookiejar)
    def set_credentials(self, credentials):
        """Set a urllib2.HTTPPasswordMgr, or None."""
        # XXX use Greg Stein's httpx instead?
        self._set_handler("_authen", obj=credentials)

    # these methods all take a boolean parameter
    def set_handle_robots(self, handle):
        """Set whether to observe rules from robots.txt."""
        self._set_handler("_robots", handle)
    def set_handle_redirect(self, handle):
        """Set whether to handle HTTP Refresh headers."""
        self._set_handler("_redirect", handle)
    def set_handle_refresh(self, handle):
        """Set whether to handle HTTP Refresh headers."""
        self._set_handler("_refresh", handle)
    def set_handle_equiv(self, handle):
        """Set whether to treat HTML http-equiv headers like HTTP headers.

        Response objects will be .seek()able if this is set.

        """
        self._set_handler("_equiv", handle)
    def set_handle_referer(self, handle):
        """Set whether to add Referer header to each request.

        This base class does not implement this feature (so don't turn this on
        if you're using this base class directly), but the subclass
        mechanize.Browser does.

        """
        self._set_handler("_referer", handle)
        self._handle_referer = True
    def set_seekable_responses(self, handle):
        """Make response objects .seek()able."""
        self._set_handler("_seek", handle)
    def set_debug_redirects(self, handle):
        """Print information about HTTP redirects.

        This includes refreshes, which show up as faked 302 redirections at the
        moment.

        """
        self._set_handler("_debug_redirect", handle)
    def set_debug_responses(self, handle):
        """Print HTTP response bodies."""
        self._set_handler("_debug_response_body", handle)
    def set_debug_http(self, handle):
        """Print HTTP headers."""
        level = int(bool(handle))
        for scheme in "http", "https":
            h = self._ua_handlers.get(scheme)
            if h is not None:
                h.set_http_debuglevel(level)

    def _set_handler(self, name, handle=None, obj=None):
        if handle is None:
            handle = obj is not None
        if handle:
            handler_class = self.handler_classes[name]
            if obj is not None:
                newhandler = handler_class(obj)
            else:
                newhandler = handler_class()
        else:
            newhandler = None
        self._replace_handler(name, newhandler)

    # XXXX I'd *really* rather get rid of this and just rebuild every time.
    #  This is fragile to base class changes, and hard to understand.
    #  Have to make sure there's no state directly stored in handlers, though,
    #  and have appropriate methods for adding state back to the cookie etc.
    #  handlers known to this class (only the ones in urllib2 / ClientCookie --
    #  no need to care about other peoples' as long as it's documented that
    #  calling the set_* methods will in general clobber handler state).
    def _replace_handler(self, name, newhandler=None):
        # first, if handler was previously added, remove it
        if name is not None:
            try:
                handler = self._ua_handlers[name]
            except:
                pass
            else:
                for table in (
                    [self.handle_open,
                     self.process_request, self.process_response]+
                    self.handle_error.values()):
                    for handlers in table.values():
                        remove(handlers, handler)
                    remove(self.handlers, handler)
        # then add the replacement, if any
        if newhandler is not None:
            self.add_handler(newhandler)
            self._ua_handlers[name] = newhandler

def remove(sequence, obj):
    # for use when can't use .remove() because of obj.__cmp__ :-(
    # (ClientCookie only requires Python 2.0, which doesn't have __lt__)
    i = 0
    while i < len(sequence):
        if sequence[i] is obj:
            del sequence[i]
        else:
            i += 1

# XXX
# This is urllib2.Request with a new .set_method() method,
# for HTTP HEAD / PUT -- move into ClientCookie if/when need it.
# Maybe it should have a constructor arg, too.
## class Request:

##     def __init__(self, url, data=None, headers={}):
##         # unwrap('<URL:type://host/path>') --> 'type://host/path'
##         self.__original = unwrap(url)
##         self.type = None
##         # self.__r_type is what's left after doing the splittype
##         self.host = None
##         self.port = None
##         self.data = data
##         self.headers = {}
##         for key, value in headers.items():
##             self.add_header(key, value)
##         if data is None:
##             self._method = "GET"
##         else:
##             self._method = "POST"

##     def __getattr__(self, attr):
##         # XXX this is a fallback mechanism to guard against these
##         # methods getting called in a non-standard order.  this may be
##         # too complicated and/or unnecessary.
##         # XXX should the __r_XXX attributes be public?
##         if attr[:12] == '_Request__r_':
##             name = attr[12:]
##             if hasattr(Request, 'get_' + name):
##                 getattr(self, 'get_' + name)()
##                 return getattr(self, attr)
##         raise AttributeError, attr

##     def get_method(self):
##         return self._method

##     def set_method(self, method):
##         if method == "POST":
##             if data is None:
##                 data = ""
##         else:
##             self.data = None
##         self._method == method

##     def add_data(self, data):
##         self.data = data

##     def has_data(self):
##         return self.data is not None

##     def get_data(self):
##         return self.data

##     def get_full_url(self):
##         return self.__original

##     def get_type(self):
##         if self.type is None:
##             self.type, self.__r_type = splittype(self.__original)
##             if self.type is None:
##                 raise ValueError, "unknown url type: %s" % self.__original
##         return self.type

##     def get_host(self):
##         if self.host is None:
##             self.host, self.__r_host = splithost(self.__r_type)
##             if self.host:
##                 self.host = unquote(self.host)
##         return self.host

##     def get_selector(self):
##         return self.__r_host

##     def set_proxy(self, host, type):
##         self.host, self.type = host, type
##         self.__r_host = self.__original

##     def add_header(self, key, val):
##         # useful for something like authentication
##         self.headers[key.capitalize()] = val

## def http_get(fullurl, ranges=None, conditions=None):
##     """HTTP GET, with convenient partial fetches (ranges).

##     XXX conditional fetches?

##     ranges: sequence of pairs of byte ranges (start, end) to fetch;

##     Ranges follow the usual Python rules (the start byte is included,
##     the end byte is not; negative numbers count back from the end of
##     the entity; start None means start of entity; end None means end of
##     entity).  There are restrictions, though: end must not be negative,
##     and if start is negative, end must be None.

##     >>> http_get("http://www.example.com/big.dat",
##                  [(0, 10), (-10, None)])  # first and last 10 bytes
##     >>> http_get("http://www.example.com/big.dat",
##                  [(50000, None)])  # from byte 50000 to the end

##     """
##     if conditions: raise NotImplementedError("conditions not yet implemented")
##     req = self._request(fullurl, data)
##     assert req.get_type() == "http", "http_get for non-HTTP URI"
##     rs = []
##     for start, end in ranges:
##         if start < 0:
##             assert end is None, "invalid range"
##             start = ""
##         else:
##             assert 0 <= start <= end, "invalid range"
##             if start == end: continue
##             end = end - 1
##         rs.append("%s-%s" % range)
##     req.add_header(("Range", "bytes=" % string.join(rs, ", ")))
##     return self.open(req)

## def http_head(self, fullurl):
##     raise NotImplementedError()  # XXX

## def http_put(self, fullurl, data=None):
##     # XXX what about 30x handling?
##     raise NotImplementedError()  # XXX
<|endoftext|>