<commit_before>"""Radial Velocity calculations.

Goals
-----
To calculate when the radial velocity is different by a certain value.
Plot radial velocity phase curves. Indicating obtained measurement locations.

"""
import argparse
import logging
from datetime import datetime
from typing import Dict, List, Any, Union

import astropy.units as u
import ephem
import matplotlib.pyplot as plt
import numpy as np
from astropy.constants import c
from mpl_toolkits.axes_grid1 import host_subplot

from utils.parse import parse_paramfile
from utils.rv_utils import RV, JulianDate, prepare_mass_params
from utils.rv_utils import strtimes2jd, join_times, check_core_parameters

c_km_s = c.to(u.kilometer / u.second)  # Speed of light in km/s

def _parser():
    # type: () -> argparse.Namespace
    """RV Argparse parser."""
    parser = argparse.ArgumentParser(description='Radial velocity plotting')
    parser.add_argument('params', help='RV parameters filename', type=str)
    parser.add_argument('-d', '--date', default=None,
                        help='Reference date in format YYYY-MM-DD [HH:MM:SS]. Default=None uses time of now.')
    # pre-predicting parated lines
    # parser.add_argument('-r', '--rv_diff', help='RV difference threshold to find')
    parser.add_argument('-o', '--obs_times', help='Times of previous observations YYYY-MM-DD format',
                        nargs='+', default=None)
    parser.add_argument('-l', '--obs_list', help='File with list of obs times.', type=str, default=None)
    # parser.add_argument('-f', '--files', help='Params and obs-times are file'
    #                    ' names to open', action='store_true')
    parser.add_argument('-m', '--mode', help='Display mode '
                                             ' e.g. phase or time plot. Default="phase"',
                        choices=['phase', 'time'], default='phase', type=str)
    parser.add_argument("--save_only", help="Only save the figure, do not show it.", action="store_true")
    parser.add_argument("--debug", help="Turning on debug output", action='store_true', default=False)
    return parser.parse_args()

def main(params, mode="phase", obs_times=None, obs_list=None, date=None,
         save_only=False):  # obs_times=None, mode='phase', rv_diff=None
    # type: (Dict[str, Union[str, float]], str, List[str], str, str, bool) -> None
    r"""Radial velocity displays.

    Parameters
    ----------
    params: str
        Filename for text file containing the rv parameters. Format of 'param = value\n'.
    mode: str
        Mode for script to use. Phase, time, future.
    obs_times: list of str
        Dates of observations added manually at command line of format YYYY-MM-DD.
    obs_list: str
        Filename for list which contains obs_times (YYY-MM-DD HH:MM:SS).
    date: str
        Reference date for some modes. Defaults=None)
    """

    # Load in params and store as a dictionary
    parameters = parse_paramfile(params)

    parameters = prepare_mass_params(parameters, only_msini=True)

    parameters = check_core_parameters(parameters)

    # combine obs_times and obs_list and turn into jd.
    if obs_times:
        if (".txt" in obs_times) or (".dat" in obs_times):
            raise ValueError("Filename given instead of dates for obs_times.")

    obs_times = join_times(obs_times, obs_list)
    obs_jd = strtimes2jd(obs_times)

    # Calculate companion semi-major axis
    if mode in ("error", "indiv"):
        pass
    else:
        pass

    host_orbit = RV.from_dict(parameters)
    companion_orbit = host_orbit.create_companion()

    if mode == "phase":
        fig = binary_phase_curve(host_orbit, companion_orbit, t_past=obs_jd)
    elif mode == "time":
        if date is not None:
            print("Date doing into ehem.julian_date", date)
            date = ephem.julian_date(date)
        fig = binary_time_curve(host_orbit, companion_orbit, t_past=obs_jd, start_day=date)
    else:
        raise NotImplementedError("Other modes are not Implemented yet.")
    if not save_only:
        fig.show()

    return fig

def binary_phase_curve(host, companion, cycle_fraction=1, ignore_mean=False, t_past=False, t_future=False):
    # type: (RV, RV, float, bool, Any, Any) -> int
    """Plot RV phase curve centered on zero.

    Parameters
    ----------
    host: RV
        RV class for systems host.
    cycle_fraction: float
        Fraction of phase space to plot. (Default=1)
    ignore_mean: bool
        Remove the contribution from the systems mean velocity. (Default=False)
    t_past: float, array-like
        Times of past observations.
    t_future: float, array-like
        Times of future observations.

    Returns
    -------
    fig: matplotlib.Figure
        Returns figure object.

    """
    host.ignore_mean = ignore_mean
    companion.ignore_mean = ignore_mean

    phase = np.linspace(-0.5, 0.5, 100) * cycle_fraction
    t = host.tau + phase * host.period

    host_rvs = host.rv_at_phase(phase)
    companion_rvs = companion.rv_at_phase(phase)

    fig = plt.figure(figsize=(10, 7))
    fig.subplots_adjust()
    ax1 = host_subplot(111)
    ax1.plot(phase, host_rvs, label="Host", lw=2, color="k")
    ax1.set_xlabel("Orbital Phase")
    ax1.set_ylabel("Host RV (km/s)")

    ax2 = ax1.twinx()
    companion_label = generate_companion_label(companion)
    ax2.plot(phase, companion_rvs, '--', label=companion_label, lw=2)
    ax2.set_ylabel("Companion RV (km/s)")

    if 'name' in host._params.keys():
        if "companion" in host._params.keys():
            plt.title("RV Phase Curve for {} {}".format(host._params['name'].upper(), host._params['companion']))
        else:
            plt.title("RV Phase Curve for {}".format(host._params['name'].upper()))
    else:
        plt.title("RV Phase Curve")

    if t_past:
        for t_num, t_val in enumerate(t_past):
            phi = ((t_val - host.tau) / host.period + 0.5) % 1 - 0.5
            rv_star = host.rv_at_phase(phi)
            rv_planet = companion.rv_at_phase(phi)
            ax1.plot(phi, rv_star, ".", markersize=10, markeredgewidth=2)
            ax2.plot(phi, rv_planet, "+", markersize=10, markeredgewidth=2)

    if t_future:
        raise NotImplementedError("Adding future observations times is not implemented yet.")
    # if t_future:
    #     for t_num, t_val in enumerate(t_future):
    #         phi = ((t_val - params[4])/params[5] - 0.5) % 1 + 0.5
    #         rv_star = RV_from_params(t_val, params, ignore_mean=False)
    #         rv_planet = RV_from_params(t_val, params, ignore_mean=False, companion=True)
    #         ax1.plot(phi, rv_star, "+", markersize=12, markeredgewidth=3)
    #         ax2.plot(phi, rv_planet, "+", markersize=12, markeredgewidth=3)

    # Determine rv max amplitudes.
    host_delta_y = host.max_amp() * 1.1
    comp_delta_y = companion.max_amp() * 1.1

    ax1.set_ylim(host.gamma - host_delta_y, host.gamma + host_delta_y)
    ax2.set_ylim(companion.gamma - comp_delta_y, companion.gamma + comp_delta_y)

    ax1.axhline(host.gamma, color="black", linestyle="-.", alpha=0.5)
    ax2.axhline(companion.gamma, color="black", linestyle="-.", alpha=0.5)

    plt.legend(loc=0)
    return fig

# Lots of duplication - could be improved

def binary_time_curve(host, companion, cycle_fraction=1, ignore_mean=False, t_past=False, t_future=False,
                      start_day=None):
    # type: (RV, RV, float, bool, Any, Any, Any) -> int
    """Plot RV phase curve centered on zero.

    Parameters
    ----------
    host: RV
        RV class for system.
    cycle_fraction: float
        Fraction of phase space to plot. (Default=1)
    ignore_mean: bool
        Remove the contribution from the systems mean velocity. (Default=False)
    t_past: float, array-like
        Times of past observations in julian days.
    t_future: float, array-like
        Times of future observations  in julian days.
    start_day: str
        Day to being RV curve from in julian days. The Default=None sets the start time to ephem.now().

    Returns
    -------
    exit_status: bool
        Returns 0 if successful.

        Displays matplotlib figure.
    """
    host.ignore_mean = ignore_mean
    companion.ignore_mean = ignore_mean

    if start_day is None:
        ephem_now = ephem.now()
        t_start = ephem.julian_date(ephem_now)
    else:
        t_start = start_day

    # Make curve from start of t_past
    if isinstance(t_past, float):
        obs_start = t_past
    elif t_past is not None:
        obs_start = np.min(t_past)
    else:
        obs_start = t_start
    # Specify 100 points per period
    num_cycles = ((t_start + host.period * cycle_fraction) - np.min([t_start, obs_start])) / host.period
    num_points = np.ceil(500 * num_cycles)
    if num_points > 10000:
        logging.debug("num points = {}".format(num_points))
        raise ValueError("num_points is to large")

    t_space = np.linspace(min([t_start, obs_start]), t_start + host.period * cycle_fraction, num_points)

    # host_rvs = RV_from_params(t_space, params, ignore_mean=ignore_mean)
    host_rvs = host.rv_at_times(t_space)
    companion_rvs = companion.rv_at_times(t_space)

    fig = plt.figure(figsize=(10, 7))
    fig.subplots_adjust()
    ax1 = host_subplot(111)
    ax1.plot(t_space - t_start, host_rvs, label="Host", lw=2, color="k")

    start_dt = JulianDate(t_start).to_datetime()
    if (start_dt.hour == 0) and (start_dt.minute == 0) and (start_dt.second == 0):
        start_string = datetime.strftime(start_dt, "%Y-%m-%d")
    else:
        start_string = datetime.strftime(start_dt, "%Y-%m-%d %H:%M:%S")  # Issue with 00:00:01 not appearing
    ax1.set_xlabel("Days from {!s}".format(start_string))

    ax1.set_ylabel("Host RV (km/s)")

    ax2 = ax1.twinx()
    companion_label = generate_companion_label(companion)
    ax2.plot(t_space - t_start, companion_rvs, '--', label=companion_label, lw=2)
    ax2.set_ylabel("Companion RV (km/s)")

    if 'name' in host._params.keys():
        if "companion" in host._params.keys():
            plt.title("Radial Velocity Curve for {} {}".format(host._params['name'].upper(), host._params['companion']))
        else:
            plt.title("Radial Velocity Curve for {}".format(host._params['name'].upper()))
    else:
        plt.title("Radial Velocity Curve")
    if t_past is not None:
        t_past = np.asarray(t_past)
        rv_star = host.rv_at_times(t_past)
        rv_planet = companion.rv_at_times(t_past)
        ax1.plot(t_past - t_start, rv_star, "b.", markersize=10, markeredgewidth=2, label="Host Obs")
        ax2.plot(t_past - t_start, rv_planet, "r+", markersize=10, markeredgewidth=2, label="Comp Obs")

    if t_future:
        t_past = np.asarray(t_future)
        rv_star = host.rv_at_times(t_future)
        rv_planet = companion.rv_at_times(t_future)
        ax1.plot(t_future - t_start, rv_star, "ko", markersize=10, markeredgewidth=2, label="Host Obs")
        ax2.plot(t_future - t_start, rv_planet, "g*", markersize=10, markeredgewidth=2, label="Comp Obs")
    # raise NotImplementedError("Adding future observations times is not implemented yet.")

    # Determine rv max amplitudes.
    amp1 = host.max_amp()
    amp2 = companion.max_amp()

    # Adjust axis limits
    ax1.set_ylim(host.gamma - (amp1 * 1.1), host.gamma + (amp1 * 1.1))
    ax2.set_ylim(companion.gamma - (amp2 * 1.1), companion.gamma + (amp2 * 1.1))

    ax1.axhline(host.gamma, color="black", linestyle="-.", alpha=0.5)
    ax2.axhline(companion.gamma, color="black", linestyle="-.", alpha=0.5)

    plt.legend(loc=0)

    return fig

if __name__ == '__main__':
    args = vars(_parser())
    debug = args.pop('debug')
    opts = {k: args[k] for k in args}

    if debug:
        logging.basicConfig(level=logging.DEBUG,
                            format='%(asctime)s %(levelname)s %(message)s')
    else:
        logging.basicConfig(level=logging.WARNING,
                            format='%(asctime)s %(levelname)s %(message)s')

    fig = main(**opts)
<commit_msg>add import of label<commit_after>"""Radial Velocity calculations.

Goals
-----
To calculate when the radial velocity is different by a certain value.
Plot radial velocity phase curves. Indicating obtained measurement locations.

"""
import argparse
import logging
from datetime import datetime
from typing import Dict, List, Any, Union

import astropy.units as u
import ephem
import matplotlib.pyplot as plt
import numpy as np
from astropy.constants import c
from mpl_toolkits.axes_grid1 import host_subplot

from utils.parse import parse_paramfile
from utils.rv_utils import RV, JulianDate, prepare_mass_params, generate_companion_label
from utils.rv_utils import strtimes2jd, join_times, check_core_parameters

c_km_s = c.to(u.kilometer / u.second)  # Speed of light in km/s

def _parser():
    # type: () -> argparse.Namespace
    """RV Argparse parser."""
    parser = argparse.ArgumentParser(description='Radial velocity plotting')
    parser.add_argument('params', help='RV parameters filename', type=str)
    parser.add_argument('-d', '--date', default=None,
                        help='Reference date in format YYYY-MM-DD [HH:MM:SS]. Default=None uses time of now.')
    # pre-predicting parated lines
    # parser.add_argument('-r', '--rv_diff', help='RV difference threshold to find')
    parser.add_argument('-o', '--obs_times', help='Times of previous observations YYYY-MM-DD format',
                        nargs='+', default=None)
    parser.add_argument('-l', '--obs_list', help='File with list of obs times.', type=str, default=None)
    # parser.add_argument('-f', '--files', help='Params and obs-times are file'
    #                    ' names to open', action='store_true')
    parser.add_argument('-m', '--mode', help='Display mode '
                                             ' e.g. phase or time plot. Default="phase"',
                        choices=['phase', 'time'], default='phase', type=str)
    parser.add_argument("--save_only", help="Only save the figure, do not show it.", action="store_true")
    parser.add_argument("--debug", help="Turning on debug output", action='store_true', default=False)
    return parser.parse_args()

def main(params, mode="phase", obs_times=None, obs_list=None, date=None,
         save_only=False):  # obs_times=None, mode='phase', rv_diff=None
    # type: (Dict[str, Union[str, float]], str, List[str], str, str, bool) -> None
    r"""Radial velocity displays.

    Parameters
    ----------
    params: str
        Filename for text file containing the rv parameters. Format of 'param = value\n'.
    mode: str
        Mode for script to use. Phase, time, future.
    obs_times: list of str
        Dates of observations added manually at command line of format YYYY-MM-DD.
    obs_list: str
        Filename for list which contains obs_times (YYY-MM-DD HH:MM:SS).
    date: str
        Reference date for some modes. Defaults=None)
    """

    # Load in params and store as a dictionary
    parameters = parse_paramfile(params)

    parameters = prepare_mass_params(parameters, only_msini=True)

    parameters = check_core_parameters(parameters)

    # combine obs_times and obs_list and turn into jd.
    if obs_times:
        if (".txt" in obs_times) or (".dat" in obs_times):
            raise ValueError("Filename given instead of dates for obs_times.")

    obs_times = join_times(obs_times, obs_list)
    obs_jd = strtimes2jd(obs_times)

    # Calculate companion semi-major axis
    if mode in ("error", "indiv"):
        pass
    else:
        pass

    host_orbit = RV.from_dict(parameters)
    companion_orbit = host_orbit.create_companion()

    if mode == "phase":
        fig = binary_phase_curve(host_orbit, companion_orbit, t_past=obs_jd)
    elif mode == "time":
        if date is not None:
            print("Date doing into ehem.julian_date", date)
            date = ephem.julian_date(date)
        fig = binary_time_curve(host_orbit, companion_orbit, t_past=obs_jd, start_day=date)
    else:
        raise NotImplementedError("Other modes are not Implemented yet.")
    if not save_only:
        fig.show()

    return fig

def binary_phase_curve(host, companion, cycle_fraction=1, ignore_mean=False, t_past=False, t_future=False):
    # type: (RV, RV, float, bool, Any, Any) -> int
    """Plot RV phase curve centered on zero.

    Parameters
    ----------
    host: RV
        RV class for systems host.
    cycle_fraction: float
        Fraction of phase space to plot. (Default=1)
    ignore_mean: bool
        Remove the contribution from the systems mean velocity. (Default=False)
    t_past: float, array-like
        Times of past observations.
    t_future: float, array-like
        Times of future observations.

    Returns
    -------
    fig: matplotlib.Figure
        Returns figure object.

    """
    host.ignore_mean = ignore_mean
    companion.ignore_mean = ignore_mean

    phase = np.linspace(-0.5, 0.5, 100) * cycle_fraction
    t = host.tau + phase * host.period

    host_rvs = host.rv_at_phase(phase)
    companion_rvs = companion.rv_at_phase(phase)

    fig = plt.figure(figsize=(10, 7))
    fig.subplots_adjust()
    ax1 = host_subplot(111)
    ax1.plot(phase, host_rvs, label="Host", lw=2, color="k")
    ax1.set_xlabel("Orbital Phase")
    ax1.set_ylabel("Host RV (km/s)")

    ax2 = ax1.twinx()
    companion_label = generate_companion_label(companion)
    ax2.plot(phase, companion_rvs, '--', label=companion_label, lw=2)
    ax2.set_ylabel("Companion RV (km/s)")

    if 'name' in host._params.keys():
        if "companion" in host._params.keys():
            plt.title("RV Phase Curve for {} {}".format(host._params['name'].upper(), host._params['companion']))
        else:
            plt.title("RV Phase Curve for {}".format(host._params['name'].upper()))
    else:
        plt.title("RV Phase Curve")

    if t_past:
        for t_num, t_val in enumerate(t_past):
            phi = ((t_val - host.tau) / host.period + 0.5) % 1 - 0.5
            rv_star = host.rv_at_phase(phi)
            rv_planet = companion.rv_at_phase(phi)
            ax1.plot(phi, rv_star, ".", markersize=10, markeredgewidth=2)
            ax2.plot(phi, rv_planet, "+", markersize=10, markeredgewidth=2)

    if t_future:
        raise NotImplementedError("Adding future observations times is not implemented yet.")
    # if t_future:
    #     for t_num, t_val in enumerate(t_future):
    #         phi = ((t_val - params[4])/params[5] - 0.5) % 1 + 0.5
    #         rv_star = RV_from_params(t_val, params, ignore_mean=False)
    #         rv_planet = RV_from_params(t_val, params, ignore_mean=False, companion=True)
    #         ax1.plot(phi, rv_star, "+", markersize=12, markeredgewidth=3)
    #         ax2.plot(phi, rv_planet, "+", markersize=12, markeredgewidth=3)

    # Determine rv max amplitudes.
    host_delta_y = host.max_amp() * 1.1
    comp_delta_y = companion.max_amp() * 1.1

    ax1.set_ylim(host.gamma - host_delta_y, host.gamma + host_delta_y)
    ax2.set_ylim(companion.gamma - comp_delta_y, companion.gamma + comp_delta_y)

    ax1.axhline(host.gamma, color="black", linestyle="-.", alpha=0.5)
    ax2.axhline(companion.gamma, color="black", linestyle="-.", alpha=0.5)

    plt.legend(loc=0)
    return fig

# Lots of duplication - could be improved

def binary_time_curve(host, companion, cycle_fraction=1, ignore_mean=False, t_past=False, t_future=False,
                      start_day=None):
    # type: (RV, RV, float, bool, Any, Any, Any) -> int
    """Plot RV phase curve centered on zero.

    Parameters
    ----------
    host: RV
        RV class for system.
    cycle_fraction: float
        Fraction of phase space to plot. (Default=1)
    ignore_mean: bool
        Remove the contribution from the systems mean velocity. (Default=False)
    t_past: float, array-like
        Times of past observations in julian days.
    t_future: float, array-like
        Times of future observations  in julian days.
    start_day: str
        Day to being RV curve from in julian days. The Default=None sets the start time to ephem.now().

    Returns
    -------
    exit_status: bool
        Returns 0 if successful.

        Displays matplotlib figure.
    """
    host.ignore_mean = ignore_mean
    companion.ignore_mean = ignore_mean

    if start_day is None:
        ephem_now = ephem.now()
        t_start = ephem.julian_date(ephem_now)
    else:
        t_start = start_day

    # Make curve from start of t_past
    if isinstance(t_past, float):
        obs_start = t_past
    elif t_past is not None:
        obs_start = np.min(t_past)
    else:
        obs_start = t_start
    # Specify 100 points per period
    num_cycles = ((t_start + host.period * cycle_fraction) - np.min([t_start, obs_start])) / host.period
    num_points = np.ceil(500 * num_cycles)
    if num_points > 10000:
        logging.debug("num points = {}".format(num_points))
        raise ValueError("num_points is to large")

    t_space = np.linspace(min([t_start, obs_start]), t_start + host.period * cycle_fraction, num_points)

    # host_rvs = RV_from_params(t_space, params, ignore_mean=ignore_mean)
    host_rvs = host.rv_at_times(t_space)
    companion_rvs = companion.rv_at_times(t_space)

    fig = plt.figure(figsize=(10, 7))
    fig.subplots_adjust()
    ax1 = host_subplot(111)
    ax1.plot(t_space - t_start, host_rvs, label="Host", lw=2, color="k")

    start_dt = JulianDate(t_start).to_datetime()
    if (start_dt.hour == 0) and (start_dt.minute == 0) and (start_dt.second == 0):
        start_string = datetime.strftime(start_dt, "%Y-%m-%d")
    else:
        start_string = datetime.strftime(start_dt, "%Y-%m-%d %H:%M:%S")  # Issue with 00:00:01 not appearing
    ax1.set_xlabel("Days from {!s}".format(start_string))

    ax1.set_ylabel("Host RV (km/s)")

    ax2 = ax1.twinx()
    companion_label = generate_companion_label(companion)
    ax2.plot(t_space - t_start, companion_rvs, '--', label=companion_label, lw=2)
    ax2.set_ylabel("Companion RV (km/s)")

    if 'name' in host._params.keys():
        if "companion" in host._params.keys():
            plt.title("Radial Velocity Curve for {} {}".format(host._params['name'].upper(), host._params['companion']))
        else:
            plt.title("Radial Velocity Curve for {}".format(host._params['name'].upper()))
    else:
        plt.title("Radial Velocity Curve")
    if t_past is not None:
        t_past = np.asarray(t_past)
        rv_star = host.rv_at_times(t_past)
        rv_planet = companion.rv_at_times(t_past)
        ax1.plot(t_past - t_start, rv_star, "b.", markersize=10, markeredgewidth=2, label="Host Obs")
        ax2.plot(t_past - t_start, rv_planet, "r+", markersize=10, markeredgewidth=2, label="Comp Obs")

    if t_future:
        t_past = np.asarray(t_future)
        rv_star = host.rv_at_times(t_future)
        rv_planet = companion.rv_at_times(t_future)
        ax1.plot(t_future - t_start, rv_star, "ko", markersize=10, markeredgewidth=2, label="Host Obs")
        ax2.plot(t_future - t_start, rv_planet, "g*", markersize=10, markeredgewidth=2, label="Comp Obs")
    # raise NotImplementedError("Adding future observations times is not implemented yet.")

    # Determine rv max amplitudes.
    amp1 = host.max_amp()
    amp2 = companion.max_amp()

    # Adjust axis limits
    ax1.set_ylim(host.gamma - (amp1 * 1.1), host.gamma + (amp1 * 1.1))
    ax2.set_ylim(companion.gamma - (amp2 * 1.1), companion.gamma + (amp2 * 1.1))

    ax1.axhline(host.gamma, color="black", linestyle="-.", alpha=0.5)
    ax2.axhline(companion.gamma, color="black", linestyle="-.", alpha=0.5)

    plt.legend(loc=0)

    return fig

if __name__ == '__main__':
    args = vars(_parser())
    debug = args.pop('debug')
    opts = {k: args[k] for k in args}

    if debug:
        logging.basicConfig(level=logging.DEBUG,
                            format='%(asctime)s %(levelname)s %(message)s')
    else:
        logging.basicConfig(level=logging.WARNING,
                            format='%(asctime)s %(levelname)s %(message)s')

    fig = main(**opts)
<|endoftext|>


<commit_before>"""
Yum plugin for Amazon S3 access.

This plugin provides access to a protected Amazon S3 bucket using either boto
or Amazon's REST authentication scheme.

On CentOS this file goes into /usr/lib/yum-plugins/s3.py

You will also need two configuration files.   See s3.conf and s3test.repo for
examples on how to deploy those.

"""

#   Copyright 2011, Robert Mela
#   Copyright 2011, Jens Braeuer
#
#   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.

def createBotoGrabber():
    import boto
    from urlparse import urlparse
    import sys
    import re
    from urlgrabber.grabber import URLGrabber

    class BotoGrabber(URLGrabber):
        def __init__(self, awsAccessKey, awsSecretKey, baseurl):
            self.verbose_logger.log(logginglevels.DEBUG_4, "creating empty URLGrabber instance")
            URLGrabber.__init__(self)
            self.verbose_logger.log(logginglevels.DEBUG_4, "BotoGrabber init BASE_URL=%s" % baseurl)
            if not baseurl: raise Exception("BotoGrabberInit got blank baseurl")
            try: baseurl = baseurl[0]
            except: pass
            self.s3 = boto.connect_s3(awsAccessKey, awsSecretKey)
            self.baseurl = urlparse(baseurl)
            if hasattr(self.baseurl, 'netloc'):
                self.bucket_name = self.baseurl.netloc
                self.key_prefix = self.baseurl.path[1:]
            else:
                self.bucket_name = self.baseurl[1]
                self.key_prefix = self.baseurl[2]
            m = re.match('(.*)\.s3.*\.amazonaws\.com', self.bucket_name)
            if (m):
                self.bucket_name = m.group(1)
            if sys.stdout.isatty():
                print "%s - %s" % (self.bucket_name, self.key_prefix)

        def _key_name(self,url):
            self.verbose_logger.log(logginglevels.DEBUG_4, "BotoGrabber _key_name url=%s, key_prefix=%s" % ( url, self.key_prefix ))
            if not url.startswith("http://"):
                return "%s%s" % ( self.key_prefix, url )
            if sys.stdout.isatty():
                print "Notice: extracting path from url (%s) instead of using prefix (%s)" % (url,self.key_prefix)
            result = urlparse(url)[2]
            if sys.stdout.isatty():
                print "Notice: extracted path is: %s" % result
            return result

        def _key(self, key_name):
            bucket = self.s3.get_bucket(self.bucket_name)
            self.verbose_logger.log(logginglevels.DEBUG_4,
                                    "BotoGrabber _key for bucket_name=%s, key_name=%s" % ( self.bucket_name, key_name ))
            return bucket.get_key(key_name)

        def urlgrab(self, url, filename=None, **kwargs):
            """urlgrab(url) copy the file to the local filesystem"""
            self.verbose_logger.log(logginglevels.DEBUG_4, "BotoGrabber urlgrab url=%s filename=%s" % ( url, filename ))
            key_name = self._key_name(url)
            self.verbose_logger.log(logginglevels.DEBUG_4,
                                    "BotoGrabber urlgrab url=%s key_name=%s filename=%s" % ( url, key_name, filename ))
            key = self._key(key_name)
            if not key: raise Exception("Can not get key for key=%s" % key_name )
            if not filename: filename = key.key
            key.get_contents_to_filename(filename)
            return filename
            # zzz - does this return a value or something?

        def urlopen(self, url, **kwargs):
            """urlopen(url) open the remote file and return a file object"""
            self.verbose_logger.log(logginglevels.DEBUG_4, "BotoGrabber urlopen url=%s" % url)
            return self._key(url)

        def urlread(self, url, limit=None, **kwargs):
            """urlread(url) return the contents of the file as a string"""
            self.verbose_logger.log(logginglevels.DEBUG_4, "BotoGrabber urlread url=%s" % url)
            return self._key(url).read()

    return BotoGrabber

def createUrllibGrabber():

    import os
    import sys
    import urllib2
    import time, sha, hmac, base64

    class UrllibGrabber:
        @classmethod

        def s3sign(cls,request, secret_key, key_id, date=None):
                date=time.strftime("%a, %d %b %Y %H:%M:%S +0000", date or time.gmtime() )
                host = request.get_host()
                bucket = host.split('.')[0]
                request.add_header( 'Date', date)
                resource = "/%s%s" % ( bucket, request.get_selector() )
                sigstring = """%(method)s\n\n\n%(date)s\n%(canon_amzn_resource)s""" % {
                                           'method':request.get_method(),
                                           #'content_md5':'',
                                           #'content_type':'', # only for PUT
                                           'date':request.headers.get('Date'),
                                           #'canon_amzn_headers':'',
                                           'canon_amzn_resource':resource }
                digest = hmac.new(secret_key, sigstring, sha ).digest()
                digest = base64.b64encode(digest)
                request.add_header('Authorization', "AWS %s:%s" % ( key_id,  digest ))

        def __init__(self, awsAccessKey, awsSecretKey, baseurl ):
            try: baseurl = baseurl[0]
            except: pass
            self.baseurl = baseurl
            self.awsAccessKey = awsAccessKey
            self.awsSecretKey = awsSecretKey

        def _request(self,url):
            req = urllib2.Request("%s%s" % (self.baseurl, url))
            UrllibGrabber.s3sign(req, self.awsSecretKey, self.awsAccessKey )
            return req

        def urlgrab(self, url, filename=None, **kwargs):
            """urlgrab(url) copy the file to the local filesystem"""
            self.verbose_logger.log(logginglevels.DEBUG_4, "UrlLibGrabber urlgrab url=%s filename=%s" % ( url, filename ))
            req = self._request(url)
            if not filename:
                filename = req.get_selector()
                if filename[0] == '/': filename = filename[1:]
            out = open(filename, 'w+')
            resp = urllib2.urlopen(req)
            buff = resp.read(8192)
            while buff:
                out.write(buff)
                buff = resp.read(8192)
            return filename
            # zzz - does this return a value or something?

        def urlopen(self, url, **kwargs):
            """urlopen(url) open the remote file and return a file object"""
            return urllib2.urlopen( self._request(url) )

        def urlread(self, url, limit=None, **kwargs):
            """urlread(url) return the contents of the file as a string"""
            return urllib2.urlopen( self._request(url) ).read()

    return UrllibGrabber

def createGrabber():
    try:
        rv = createBotoGrabber()
        self.verbose_logger.log(logginglevels.DEBUG_4, "Created BotoGrabber")
        return rv
    except:
        self.verbose_logger.log(logginglevels.DEBUG_4, "Creating UrllibGrabber")
        return createUrllibGrabber()

AmazonS3Grabber = createGrabber()

import logging
import os
import sys
import urllib

from yum.plugins import TYPE_CORE
from yum.yumRepo import YumRepository
from yum import config
from yum import logginglevels

import yum.Errors

__revision__ = "1.0.0"

requires_api_version = '2.5'
plugin_type = TYPE_CORE
CONDUIT=None

def config_hook(conduit):
    logger = logging.getLogger("yum.verbose.main")
    config.RepoConf.s3_enabled = config.BoolOption(False)
    config.RepoConf.key_id = config.Option() or conduit.confString('main', 'aws_access_key_id')
    config.RepoConf.secret_key = config.Option() or conduit.confString('main', 'aws_secret_access_key')

def init_hook(conduit):
    """
    Plugin initialization hook. Setup the S3 repositories.
    """

    repos = conduit.getRepos()
    for key,repo in repos.repos.iteritems():
        if isinstance(repo, YumRepository) and repo.s3_enabled:
            new_repo = AmazonS3Repo(key)
            new_repo.baseurl = repo.baseurl
            new_repo.mirrorlist = repo.mirrorlist
            new_repo.basecachedir = repo.basecachedir
            new_repo.gpgcheck = repo.gpgcheck
            new_repo.proxy = repo.proxy
            new_repo.enablegroups = repo.enablegroups
            new_repo.key_id = repo.key_id
            new_repo.secret_key = repo.secret_key
            repos.delete(repo.id)
            repos.add(new_repo)

class AmazonS3Repo(YumRepository):
    """
    Repository object for Amazon S3.
    """

    def __init__(self, repoid):
        YumRepository.__init__(self, repoid)
        self.enable()
        self.grabber = None

    def setupGrab(self):
        YumRepository.setupGrab(self)
        self.grabber = AmazonS3Grabber(self.key_id, self.secret_key )

    def _getgrabfunc(self): raise Exception("get grabfunc!")
    def _getgrab(self):
        if not self.grabber:
            self.grabber = AmazonS3Grabber(self.key_id, self.secret_key, baseurl=self.baseurl )
        return self.grabber

    grabfunc = property(lambda self: self._getgrabfunc())
    grab = property(lambda self: self._getgrab())

<commit_msg>Move code to detect interactive invocations into own function.<commit_after>"""
Yum plugin for Amazon S3 access.

This plugin provides access to a protected Amazon S3 bucket using either boto
or Amazon's REST authentication scheme.

On CentOS this file goes into /usr/lib/yum-plugins/s3.py

You will also need two configuration files.   See s3.conf and s3test.repo for
examples on how to deploy those.

"""

#   Copyright 2011, Robert Mela
#   Copyright 2011, Jens Braeuer
#
#   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.

def interactive_notify(msg):
    if sys.stdout.isatty():
        print msg

def createBotoGrabber():
    import boto
    from urlparse import urlparse
    import sys
    import re
    from urlgrabber.grabber import URLGrabber

    class BotoGrabber(URLGrabber):
        def __init__(self, awsAccessKey, awsSecretKey, baseurl):
            self.verbose_logger.log(logginglevels.DEBUG_4, "creating empty URLGrabber instance")
            URLGrabber.__init__(self)
            self.verbose_logger.log(logginglevels.DEBUG_4, "BotoGrabber init BASE_URL=%s" % baseurl)
            if not baseurl: raise Exception("BotoGrabberInit got blank baseurl")
            try: baseurl = baseurl[0]
            except: pass
            self.s3 = boto.connect_s3(awsAccessKey, awsSecretKey)
            self.baseurl = urlparse(baseurl)
            if hasattr(self.baseurl, 'netloc'):
                self.bucket_name = self.baseurl.netloc
                self.key_prefix = self.baseurl.path[1:]
            else:
                self.bucket_name = self.baseurl[1]
                self.key_prefix = self.baseurl[2]
            m = re.match('(.*)\.s3.*\.amazonaws\.com', self.bucket_name)
            if (m):
                self.bucket_name = m.group(1)
            interactive_notify("%s - %s" % (self.bucket_name, self.key_prefix))

        def _key_name(self,url):
            self.verbose_logger.log(logginglevels.DEBUG_4, "BotoGrabber _key_name url=%s, key_prefix=%s" % ( url, self.key_prefix ))
            if not url.startswith("http://"):
                return "%s%s" % ( self.key_prefix, url )
            interactive_notify("Notice: extracting path from url (%s) instead of using prefix (%s)" % (url,self.key_prefix))
            result = urlparse(url)[2]
            interactive_notify("Notice: extracted path is: %s" % result)
            return result

        def _key(self, key_name):
            bucket = self.s3.get_bucket(self.bucket_name)
            self.verbose_logger.log(logginglevels.DEBUG_4,
                                    "BotoGrabber _key for bucket_name=%s, key_name=%s" % ( self.bucket_name, key_name ))
            return bucket.get_key(key_name)

        def urlgrab(self, url, filename=None, **kwargs):
            """urlgrab(url) copy the file to the local filesystem"""
            self.verbose_logger.log(logginglevels.DEBUG_4, "BotoGrabber urlgrab url=%s filename=%s" % ( url, filename ))
            key_name = self._key_name(url)
            self.verbose_logger.log(logginglevels.DEBUG_4,
                                    "BotoGrabber urlgrab url=%s key_name=%s filename=%s" % ( url, key_name, filename ))
            key = self._key(key_name)
            if not key: raise Exception("Can not get key for key=%s" % key_name )
            if not filename: filename = key.key
            key.get_contents_to_filename(filename)
            return filename
            # zzz - does this return a value or something?

        def urlopen(self, url, **kwargs):
            """urlopen(url) open the remote file and return a file object"""
            self.verbose_logger.log(logginglevels.DEBUG_4, "BotoGrabber urlopen url=%s" % url)
            return self._key(url)

        def urlread(self, url, limit=None, **kwargs):
            """urlread(url) return the contents of the file as a string"""
            self.verbose_logger.log(logginglevels.DEBUG_4, "BotoGrabber urlread url=%s" % url)
            return self._key(url).read()

    return BotoGrabber

def createUrllibGrabber():

    import os
    import sys
    import urllib2
    import time, sha, hmac, base64

    class UrllibGrabber:
        @classmethod

        def s3sign(cls,request, secret_key, key_id, date=None):
                date=time.strftime("%a, %d %b %Y %H:%M:%S +0000", date or time.gmtime() )
                host = request.get_host()
                bucket = host.split('.')[0]
                request.add_header( 'Date', date)
                resource = "/%s%s" % ( bucket, request.get_selector() )
                sigstring = """%(method)s\n\n\n%(date)s\n%(canon_amzn_resource)s""" % {
                                           'method':request.get_method(),
                                           #'content_md5':'',
                                           #'content_type':'', # only for PUT
                                           'date':request.headers.get('Date'),
                                           #'canon_amzn_headers':'',
                                           'canon_amzn_resource':resource }
                digest = hmac.new(secret_key, sigstring, sha ).digest()
                digest = base64.b64encode(digest)
                request.add_header('Authorization', "AWS %s:%s" % ( key_id,  digest ))

        def __init__(self, awsAccessKey, awsSecretKey, baseurl ):
            try: baseurl = baseurl[0]
            except: pass
            self.baseurl = baseurl
            self.awsAccessKey = awsAccessKey
            self.awsSecretKey = awsSecretKey

        def _request(self,url):
            req = urllib2.Request("%s%s" % (self.baseurl, url))
            UrllibGrabber.s3sign(req, self.awsSecretKey, self.awsAccessKey )
            return req

        def urlgrab(self, url, filename=None, **kwargs):
            """urlgrab(url) copy the file to the local filesystem"""
            self.verbose_logger.log(logginglevels.DEBUG_4, "UrlLibGrabber urlgrab url=%s filename=%s" % ( url, filename ))
            req = self._request(url)
            if not filename:
                filename = req.get_selector()
                if filename[0] == '/': filename = filename[1:]
            out = open(filename, 'w+')
            resp = urllib2.urlopen(req)
            buff = resp.read(8192)
            while buff:
                out.write(buff)
                buff = resp.read(8192)
            return filename
            # zzz - does this return a value or something?

        def urlopen(self, url, **kwargs):
            """urlopen(url) open the remote file and return a file object"""
            return urllib2.urlopen( self._request(url) )

        def urlread(self, url, limit=None, **kwargs):
            """urlread(url) return the contents of the file as a string"""
            return urllib2.urlopen( self._request(url) ).read()

    return UrllibGrabber

def createGrabber():
    try:
        rv = createBotoGrabber()
        self.verbose_logger.log(logginglevels.DEBUG_4, "Created BotoGrabber")
        return rv
    except:
        self.verbose_logger.log(logginglevels.DEBUG_4, "Creating UrllibGrabber")
        return createUrllibGrabber()

AmazonS3Grabber = createGrabber()

import logging
import os
import sys
import urllib

from yum.plugins import TYPE_CORE
from yum.yumRepo import YumRepository
from yum import config
from yum import logginglevels

import yum.Errors

__revision__ = "1.0.0"

requires_api_version = '2.5'
plugin_type = TYPE_CORE
CONDUIT=None

def config_hook(conduit):
    logger = logging.getLogger("yum.verbose.main")
    config.RepoConf.s3_enabled = config.BoolOption(False)
    config.RepoConf.key_id = config.Option() or conduit.confString('main', 'aws_access_key_id')
    config.RepoConf.secret_key = config.Option() or conduit.confString('main', 'aws_secret_access_key')

def init_hook(conduit):
    """
    Plugin initialization hook. Setup the S3 repositories.
    """

    repos = conduit.getRepos()
    for key,repo in repos.repos.iteritems():
        if isinstance(repo, YumRepository) and repo.s3_enabled:
            new_repo = AmazonS3Repo(key)
            new_repo.baseurl = repo.baseurl
            new_repo.mirrorlist = repo.mirrorlist
            new_repo.basecachedir = repo.basecachedir
            new_repo.gpgcheck = repo.gpgcheck
            new_repo.proxy = repo.proxy
            new_repo.enablegroups = repo.enablegroups
            new_repo.key_id = repo.key_id
            new_repo.secret_key = repo.secret_key
            repos.delete(repo.id)
            repos.add(new_repo)

class AmazonS3Repo(YumRepository):
    """
    Repository object for Amazon S3.
    """

    def __init__(self, repoid):
        YumRepository.__init__(self, repoid)
        self.enable()
        self.grabber = None

    def setupGrab(self):
        YumRepository.setupGrab(self)
        self.grabber = AmazonS3Grabber(self.key_id, self.secret_key )

    def _getgrabfunc(self): raise Exception("get grabfunc!")
    def _getgrab(self):
        if not self.grabber:
            self.grabber = AmazonS3Grabber(self.key_id, self.secret_key, baseurl=self.baseurl )
        return self.grabber

    grabfunc = property(lambda self: self._getgrabfunc())
    grab = property(lambda self: self._getgrab())

<|endoftext|>


<commit_before>#!/usr/bin/python

from time import sleep, strftime
from datetime import datetime
import picamera
#from Adafruit_I2C import Adafruit_I2C
#from Adafruit_MCP230xx import Adafruit_MCP230XX
import Adafruit_CharLCD as LCD

#VIDEO_DAYS = 5
FRAMES_PER_HOUR = 6
#FRAMES = FRAMES_PER_HOUR * 24 * VIDEO_DAYS
ROTATION = 90

lcd = LCD.Adafruit_CharLCDPlate()
lcd.backlight(lcd.ON)

def capture_frame(path, frame):
    lcd.clear()
    lcd.message('Shooting #%03d' % frame)
    with picamera.PiCamera() as cam:
        cam.rotation(ROTATION)
        time.sleep(2)
        cam.capture('%s/frame%03d.jpg' % (path, frame))

#create directory with date
date = datetime.now().strftime('%b %d')
directory = '/home/pi/%s' % date
if not os.path.exists(directory):
    os.makedirs(directory)

# Capture the images
while True:
    # Note the time before the capture
    start = time.time()
    capture_frame(directory, frame)
    # Wait for the next capture. Note that we take into
    # account the length of time it took to capture the
    # image when calculating the delay
    while True:
        seconds_to_next = int(60 * 60 / FRAMES_PER_HOUR) - (time.time() - start)
        if seconds_to_next <= 0:
            break
        lcd.clear()
        lcd.message('taken: %03d' % frame)
        lcd.message('next: %s' % str(datetime.timedelta(seconds=seconds_to_next)))
        time.sleep(1)
<commit_msg>removed backlight<commit_after>#!/usr/bin/python

from time import sleep, strftime
from datetime import datetime
import picamera
#from Adafruit_I2C import Adafruit_I2C
#from Adafruit_MCP230xx import Adafruit_MCP230XX
import Adafruit_CharLCD as LCD

#VIDEO_DAYS = 5
FRAMES_PER_HOUR = 6
#FRAMES = FRAMES_PER_HOUR * 24 * VIDEO_DAYS
ROTATION = 90

lcd = LCD.Adafruit_CharLCDPlate()
#lcd.backlight(lcd.ON)

def capture_frame(path, frame):
    lcd.clear()
    lcd.message('Shooting #%03d' % frame)
    with picamera.PiCamera() as cam:
        cam.rotation(ROTATION)
        time.sleep(2)
        cam.capture('%s/frame%03d.jpg' % (path, frame))

#create directory with date
date = datetime.now().strftime('%b %d')
directory = '/home/pi/%s' % date
if not os.path.exists(directory):
    os.makedirs(directory)

# Capture the images
while True:
    # Note the time before the capture
    start = time.time()
    capture_frame(directory, frame)
    # Wait for the next capture. Note that we take into
    # account the length of time it took to capture the
    # image when calculating the delay
    while True:
        seconds_to_next = int(60 * 60 / FRAMES_PER_HOUR) - (time.time() - start)
        if seconds_to_next <= 0:
            break
        lcd.clear()
        lcd.message('taken: %03d' % frame)
        lcd.message('next: %s' % str(datetime.timedelta(seconds=seconds_to_next)))
        time.sleep(1)
<|endoftext|>


<commit_before>import os
import sys
import re
from datetime import datetime, date, timedelta
import sqlite3 as sql
import urllib2
from xml.dom import minidom
import webbrowser

'''
Constants
'''
MAX_TITLE_LENGTH = 20
MAX_DESC_LENGTH = 65
DB_NAME = 'test.db'

DATABASE_NAME = os.path.join(os.path.dirname(os.path.realpath(__file__)), DB_NAME)

def main():
    initDatabase()

    '''
    Update unwatched episodes list with the episodes from the feed
    '''
    updateSeries(all = True)

    if len(sys.argv) == 2:
        if sys.argv[1] == 'uwcount':
            episodes = getUnwatched()

            episodes = [(ep[1], ep[2], ep[3]) for ep in episodes if toDate(ep[5]) <= date.today()]

            print len(episodes)
            for ep in episodes:
               print "{0} s{1:02} e{2:02}".format(*ep)

        exit()
    '''
    Get user input on what to do
    '''
    print ''
    listUnwatched()
    while True:
        print ''
        print "Enter a command or h for help"
        input = raw_input(":")
        print ''

        if   input == 'h': printHelp()
        elif input == 'l': listUnwatched()
        elif input == 'ls': listMyShows()
        elif input == 'fr': forceUpdate()
        elif input == 'q': exit()
        elif re.match('^o\d+$', input): openLinks(int(re.search('^o(\d+)$', input).group(1)))
        elif re.match('^w\d+$', input): markWatched(int(re.search('^w(\d+)$', input).group(1)))
        elif re.match('^wall\d+$', input): markWatched(int(re.search('^wall(\d+)$', input).group(1)), True)
        elif re.match('^s[\w\s]+$', input): searchForSeries(re.search('^s([\w\s]+)$', input).group(1))
        elif re.match('^a\d+$', input): addSeries(re.search('^a(\d+)$', input).group(1))
        elif re.match('^d\d+$', input): deleteSeries(re.search('^d(\d+)$', input).group(1))
        else: print "Invalid input. Enter h for help."

def listMyShows():
    '''
    Lists all shows in the database along with their ID
    '''
    con = sql.connect(DATABASE_NAME)
    cur = con.cursor()
    cur.execute("SELECT * FROM shows;")
    shows = cur.fetchall();
    con.close()

    print "Found {0} results".format(len(shows))

    if len(shows) > 0:
        print " {0:<5} | {1}".format("ID", "Series Name")
        for show in shows:
            id = show[0]
            name = show[1].encode('utf-8')
            print " {0:<5} | {1}".format(id, name)

def forceUpdate():
    '''
    Force program to refresh all episode data from TVRage
    '''
    con = sql.connect(DATABASE_NAME)
    cur = con.cursor()
    cur.execute("INSERT OR REPLACE INTO settings VALUES('lastUpdate', '');")
    con.commit()
    con.close()
    updateSeries(all = True)

def updateSeries(id = [0,], all = False):
    '''
    Get latest unwatched episode data from TVRage and update database
    Parameters:
    id (list)  - List of TVRage series IDs to update
    all (bool) - True: all shows in database will be updated
                 False: Only shows specified in id list will be updated
    '''
    if all:
        # Get list of shows from database
        con = sql.connect(DATABASE_NAME)
        cur = con.cursor()
        cur.execute("SELECT id FROM shows;")
        result = cur.fetchall();
        con.close()

        # flatten result list
        seriesIds = [e for tup in result for e in tup]
        results = getEpisodes(seriesIds)
    else:
        results = getEpisodes(id, force = True)

    con = sql.connect(DATABASE_NAME)
    cur = con.cursor()
    for result in results:
        cur.execute("INSERT OR IGNORE INTO unwatched VALUES(?, ?, ?, ?, ?, 0)",
                    (result['id'], result['season'], result['episode'], result['description'], result['airDate']))
        con.commit()
    con.close()

def openLinks(index):
    try:
        episode = getUnwatched()[index]

        urls = ("http://www.newshost.co.za/?s={0}+s{1:02d}e{2:02d}", "http://www.nzbclub.com/search.aspx?q={0}+s{1:02d}e{2:02d}")

        print "Opening links for", episode[1]
        for url in urls:
            url = urllib2.quote(url.format(episode[1], episode[2], episode[3]), safe="%/:=&?~#+!$,;'@()*[]")
            webbrowser.open_new_tab(url)

    except IndexError:
        print "Invalid ID"

def searchForSeries(query):
    '''
    Display a list of show names that match the search string
    Parameters:
    query (string) - What to search for
    '''
    url = "http://services.tvrage.com/feeds/search.php?show={0}".format(query)
    url = urllib2.quote(url, safe="%/:=&?~#+!$,;'@()*[]")

    try:
        dom = minidom.parse(urllib2.urlopen(url, timeout=10.0))

        shows = dom.getElementsByTagName('show')

        print "Found {0} results".format(len(shows))
        print " {0:<5} | {1}".format("ID", "Series Name")

        for show in shows:
            id = show.childNodes[1].firstChild.nodeValue
            name = show.childNodes[3].firstChild.nodeValue.encode('utf-8')
            print " {0:<5} | {1}".format(id, name)

    except:
        print "Unexpected error or timeout"

def addSeries(id):
    '''
    Add a series to monitor unwatched episode
    Parameters:
    id (int) - ID of the series from TVRage
    '''
    url = "http://services.tvrage.com/feeds/showinfo.php?sid={0}"
    try:
        dom = minidom.parse(urllib2.urlopen(url.format(id), timeout=10.0))
        showName = dom.getElementsByTagName('showname').item(0).firstChild.nodeValue

        con = sql.connect(DATABASE_NAME)
        cur = con.cursor()
        cur.execute("INSERT OR IGNORE INTO shows VALUES(?, ?);", (id,showName))
        con.commit()
        con.close()
        print "Added", showName

        updateSeries(id = [id,])
    except:
        print "Unexpected error or timeout"

def deleteSeries(id):
    '''
    Delete an entire series from the database
    Parameters:
    id (int) - ID of the series from TVRage
    '''

    con = sql.connect(DATABASE_NAME)
    cur = con.cursor()
    cur.execute("PRAGMA foreign_keys = ON;")
    con.commit()

    cur.execute("SELECT title FROM shows WHERE id = ?;", (id,))
    result = cur.fetchone()
    if result:
        cur.execute("DELETE FROM shows WHERE id = ?;", (id,))
        con.commit()
        print "Deleted", result[0]
    else:
        print "Show not found. Use ls to look up the ID"
    con.close()

def markWatched(index, allPrevious = False):
    '''
    Mark an episode as watched so it no longer appears in the list
    Parameters:
    index (int)         - Zero based index for episode to mark as watched from the list
    allPrevious (bool)  - If true all episodes in this series up to and including the one specified in index will be marked as watched
    '''
    try:
        episode = getUnwatched()[index]

        if toDate(episode[5]) > date.today():
            print "This episode has not aired yet."
            return

        # Mark episode watched in database
        con = sql.connect(DATABASE_NAME)
        cur = con.cursor()

        if allPrevious:
            cur.execute("""UPDATE unwatched SET watched = 1
                        WHERE id = ?
                        AND ((season < ?)
                        OR (season = ? AND episode <= ?))
                        """, (episode[0], episode[2], episode[2], episode[3]))
            print "Marked {0} Season {1} Episode {2} and all previous episodes in the series as watched".format(episode[1], episode[2], episode[3])

        else:
            cur.execute("""UPDATE unwatched SET watched = 1
                        WHERE id = ?
                        AND   season = ?
                        AND   episode = ?
                        """, (episode[0], episode[2], episode[3]))
            print "Marked {0} Season {1} Episode {2} as watched".format(episode[1], episode[2], episode[3])
        con.commit()
        con.close()

    except IndexError:
        print "Invalid ID"

def getUnwatched():
    '''
    Get unwatched episodes

    Returns:
    (list) - list of unwatched episodes
    '''
    con = sql.connect(DATABASE_NAME)
    cur = con.cursor()
    cur.execute("SELECT shows.id, shows.title, season, episode, description, airDate FROM unwatched INNER JOIN shows on shows.id = unwatched.id WHERE watched = 0 ORDER BY airDate ASC, shows.title ASC;")
    data = cur.fetchall();
    con.close()

    return data

def listUnwatched():
    '''
    Display a list of unwatched episodes
    '''
    episodes = getUnwatched()

    if len(episodes) == 0:
        print "No unwatched episodes!"
    else:
        print " {0:3} | {1:20} | {2} | {3:10} | {4}".format("ID", "Show", "Episode", "Date", "Description")

        for idx, row in enumerate(episodes):
            title = (row[1][:MAX_TITLE_LENGTH-3] + '...') if len(row[1]) > MAX_TITLE_LENGTH else row[1]
            desc = (row[4][:MAX_DESC_LENGTH-3] + '...') if len(row[4]) > MAX_DESC_LENGTH else row[4]

            # Don't show ID for episodes not aired yet.
            epDate = toDate(row[5])
            if epDate is None or epDate > date.today():
                idx = ' - '

            print " {0:<3} | {1:20} | s{2:02} e{3:02} | {4} | {5}".format(idx, title, row[2], row[3], row[5], desc.encode('utf-8'))

            try:
                if toDate(row[5]) <= date.today() and toDate(episodes[idx+1][5]) > date.today():
                    print ''
            except:
                pass

def printHelp():
    '''
    Display help and available commands from help.txt
    '''
    try:
        with open('help.txt', 'r') as fin:
            print fin.read()
    except IOError:
        print "Help file not found :("

def initDatabase():
    '''
    Setup the SQLite database if it doesn't already exist
    '''
    con = sql.connect(DATABASE_NAME)
    with con:
        cur = con.cursor()

        cur.execute("""CREATE TABLE IF NOT EXISTS unwatched (
                    'id' INTEGER NOT NULL,
                    'season' INTEGER NOT NULL,
                    'episode' INTEGER NOT NULL,
                    'description' TEXT,
                    'airDate' TEXT NOT NULL,
                    'watched' INTEGER NOT NULL DEFAULT (0),
                    PRIMARY KEY(id, season, episode)
                    FOREIGN KEY(id) REFERENCES shows(id) ON DELETE CASCADE
                    )""")

        cur.execute("""CREATE TABLE IF NOT EXISTS shows (
                    'id' INTEGER NOT NULL,
                    'title' TEXT NOT NULL,
                    PRIMARY KEY (id)
                    )""")

        cur.execute("""CREATE TABLE IF NOT EXISTS settings (
                    'name' TEXT NOT NULL,
                    'value' TEXT NOT NULL,
                    PRIMARY KEY (name)
                    )""")
        cur.execute("INSERT OR IGNORE INTO settings VALUES('lastUpdate', '');");
        con.commit()

def getEpisodes(seriesIds, force = False):
    '''
    Get the episode list for shows
    Parameters:
    seriesIds - List of series IDs from TVRage

    Returns:
    List of episodes with the following information:
    id          - TVRage ID for the series
    series      - Name of the series
    season      - Season number
    episode     - Episode number in the season
    airDate     - Original air date
    description - Description for the episode
    '''

    # Checks if an update is necessary
    doUpdate = False
    if force:
        doUpdate = True
    else:
        con = sql.connect(DATABASE_NAME)
        cur = con.cursor()
        cur.execute("SELECT value FROM settings WHERE name = 'lastUpdate';")
        data = cur.fetchone()[0];
        con.close()
        if data:
            if toDate(data) < date.today():
                doUpdate = True
        else:
            doUpdate = True

    episodeList = []

    if doUpdate:
        print "Updating... "

        feedUrl = "http://www.tvrage.com/feeds/episode_list.php?sid={0}"

        for id in seriesIds:
            try:
                dom = minidom.parse(urllib2.urlopen("http://www.tvrage.com/feeds/episode_list.php?sid={0}".format(id), timeout=10.0))

                seriesName = dom.childNodes[0].childNodes[1].firstChild.nodeValue
                totalSeasons = int(dom.childNodes[0].childNodes[3].firstChild.nodeValue)

                print seriesName

                for season in dom.getElementsByTagName('Season'):
                    seasonNum = int(season.getAttribute('no'))
                    for episode in season.getElementsByTagName('episode'):

                        episodeNumFull = episode.getElementsByTagName('epnum')[0].firstChild.nodeValue
                        episodeNum = episode.getElementsByTagName('seasonnum')[0].firstChild.nodeValue
                        airdate = episode.getElementsByTagName('airdate')[0].firstChild.nodeValue
                        title = episode.getElementsByTagName('title')[0].firstChild.nodeValue

                        # Only get episodes that have an air date and those airing up until the next week
                        nextWeek = date.today() + timedelta(days=7)
                        if airdate != "0000-00-00" and toDate(airdate) <= nextWeek:
                            episodeList.append({
                                'id': id,
                                'series': seriesName,
                                'season': int(seasonNum),
                                'episode': int(episodeNum),
                                'airDate': airdate,
                                'description': title
                            })

                # Update last update date in database
                con = sql.connect(DATABASE_NAME)
                cur = con.cursor()
                cur.execute("INSERT OR REPLACE INTO settings VALUES('lastUpdate', ?);", (date.today(),))
                con.commit()
                con.close()
            except:
                print "Unexpected error or timeout"

    return episodeList

def toDate(dateStr):
    '''
    Convert a date string to a date object

    Parameters:
    dateStr - date string in yyyy-mm-dd format

    Returns:
    date object
    '''
    try:
        return datetime.strptime(dateStr, '%Y-%m-%d').date()
    except:
        # Sometimes the date is like 2010-00-00.. weird.
        pass
main()<commit_msg>fixed date comparison bug in command line mode<commit_after>import os
import sys
import re
from datetime import datetime, date, timedelta
import sqlite3 as sql
import urllib2
from xml.dom import minidom
import webbrowser

'''
Constants
'''
MAX_TITLE_LENGTH = 20
MAX_DESC_LENGTH = 65
DB_NAME = 'test.db'

DATABASE_NAME = os.path.join(os.path.dirname(os.path.realpath(__file__)), DB_NAME)

def main():
    initDatabase()

    '''
    Update unwatched episodes list with the episodes from the feed
    '''
    updateSeries(all = True)

    if len(sys.argv) == 2:
        if sys.argv[1] == 'uwcount':
            episodes = getUnwatched()

            episodes = [(ep[1], ep[2], ep[3]) for ep in episodes if toDate(ep[5]) is not None and toDate(ep[5]) <= date.today()]

            print len(episodes)
            for ep in episodes:
               print "{0} s{1:02} e{2:02}".format(*ep)

        exit()
    '''
    Get user input on what to do
    '''
    print ''
    listUnwatched()
    while True:
        print ''
        print "Enter a command or h for help"
        input = raw_input(":")
        print ''

        if   input == 'h': printHelp()
        elif input == 'l': listUnwatched()
        elif input == 'ls': listMyShows()
        elif input == 'fr': forceUpdate()
        elif input == 'q': exit()
        elif re.match('^o\d+$', input): openLinks(int(re.search('^o(\d+)$', input).group(1)))
        elif re.match('^w\d+$', input): markWatched(int(re.search('^w(\d+)$', input).group(1)))
        elif re.match('^wall\d+$', input): markWatched(int(re.search('^wall(\d+)$', input).group(1)), True)
        elif re.match('^s[\w\s]+$', input): searchForSeries(re.search('^s([\w\s]+)$', input).group(1))
        elif re.match('^a\d+$', input): addSeries(re.search('^a(\d+)$', input).group(1))
        elif re.match('^d\d+$', input): deleteSeries(re.search('^d(\d+)$', input).group(1))
        else: print "Invalid input. Enter h for help."

def listMyShows():
    '''
    Lists all shows in the database along with their ID
    '''
    con = sql.connect(DATABASE_NAME)
    cur = con.cursor()
    cur.execute("SELECT * FROM shows;")
    shows = cur.fetchall();
    con.close()

    print "Found {0} results".format(len(shows))

    if len(shows) > 0:
        print " {0:<5} | {1}".format("ID", "Series Name")
        for show in shows:
            id = show[0]
            name = show[1].encode('utf-8')
            print " {0:<5} | {1}".format(id, name)

def forceUpdate():
    '''
    Force program to refresh all episode data from TVRage
    '''
    con = sql.connect(DATABASE_NAME)
    cur = con.cursor()
    cur.execute("INSERT OR REPLACE INTO settings VALUES('lastUpdate', '');")
    con.commit()
    con.close()
    updateSeries(all = True)

def updateSeries(id = [0,], all = False):
    '''
    Get latest unwatched episode data from TVRage and update database
    Parameters:
    id (list)  - List of TVRage series IDs to update
    all (bool) - True: all shows in database will be updated
                 False: Only shows specified in id list will be updated
    '''
    if all:
        # Get list of shows from database
        con = sql.connect(DATABASE_NAME)
        cur = con.cursor()
        cur.execute("SELECT id FROM shows;")
        result = cur.fetchall();
        con.close()

        # flatten result list
        seriesIds = [e for tup in result for e in tup]
        results = getEpisodes(seriesIds)
    else:
        results = getEpisodes(id, force = True)

    con = sql.connect(DATABASE_NAME)
    cur = con.cursor()
    for result in results:
        cur.execute("INSERT OR IGNORE INTO unwatched VALUES(?, ?, ?, ?, ?, 0)",
                    (result['id'], result['season'], result['episode'], result['description'], result['airDate']))
        con.commit()
    con.close()

def openLinks(index):
    try:
        episode = getUnwatched()[index]

        urls = ("http://www.newshost.co.za/?s={0}+s{1:02d}e{2:02d}", "http://www.nzbclub.com/search.aspx?q={0}+s{1:02d}e{2:02d}")

        print "Opening links for", episode[1]
        for url in urls:
            url = urllib2.quote(url.format(episode[1], episode[2], episode[3]), safe="%/:=&?~#+!$,;'@()*[]")
            webbrowser.open_new_tab(url)

    except IndexError:
        print "Invalid ID"

def searchForSeries(query):
    '''
    Display a list of show names that match the search string
    Parameters:
    query (string) - What to search for
    '''
    url = "http://services.tvrage.com/feeds/search.php?show={0}".format(query)
    url = urllib2.quote(url, safe="%/:=&?~#+!$,;'@()*[]")

    try:
        dom = minidom.parse(urllib2.urlopen(url, timeout=10.0))

        shows = dom.getElementsByTagName('show')

        print "Found {0} results".format(len(shows))
        print " {0:<5} | {1}".format("ID", "Series Name")

        for show in shows:
            id = show.childNodes[1].firstChild.nodeValue
            name = show.childNodes[3].firstChild.nodeValue.encode('utf-8')
            print " {0:<5} | {1}".format(id, name)

    except:
        print "Unexpected error or timeout"

def addSeries(id):
    '''
    Add a series to monitor unwatched episode
    Parameters:
    id (int) - ID of the series from TVRage
    '''
    url = "http://services.tvrage.com/feeds/showinfo.php?sid={0}"
    try:
        dom = minidom.parse(urllib2.urlopen(url.format(id), timeout=10.0))
        showName = dom.getElementsByTagName('showname').item(0).firstChild.nodeValue

        con = sql.connect(DATABASE_NAME)
        cur = con.cursor()
        cur.execute("INSERT OR IGNORE INTO shows VALUES(?, ?);", (id,showName))
        con.commit()
        con.close()
        print "Added", showName

        updateSeries(id = [id,])
    except:
        print "Unexpected error or timeout"

def deleteSeries(id):
    '''
    Delete an entire series from the database
    Parameters:
    id (int) - ID of the series from TVRage
    '''

    con = sql.connect(DATABASE_NAME)
    cur = con.cursor()
    cur.execute("PRAGMA foreign_keys = ON;")
    con.commit()

    cur.execute("SELECT title FROM shows WHERE id = ?;", (id,))
    result = cur.fetchone()
    if result:
        cur.execute("DELETE FROM shows WHERE id = ?;", (id,))
        con.commit()
        print "Deleted", result[0]
    else:
        print "Show not found. Use ls to look up the ID"
    con.close()

def markWatched(index, allPrevious = False):
    '''
    Mark an episode as watched so it no longer appears in the list
    Parameters:
    index (int)         - Zero based index for episode to mark as watched from the list
    allPrevious (bool)  - If true all episodes in this series up to and including the one specified in index will be marked as watched
    '''
    try:
        episode = getUnwatched()[index]

        if toDate(episode[5]) > date.today():
            print "This episode has not aired yet."
            return

        # Mark episode watched in database
        con = sql.connect(DATABASE_NAME)
        cur = con.cursor()

        if allPrevious:
            cur.execute("""UPDATE unwatched SET watched = 1
                        WHERE id = ?
                        AND ((season < ?)
                        OR (season = ? AND episode <= ?))
                        """, (episode[0], episode[2], episode[2], episode[3]))
            print "Marked {0} Season {1} Episode {2} and all previous episodes in the series as watched".format(episode[1], episode[2], episode[3])

        else:
            cur.execute("""UPDATE unwatched SET watched = 1
                        WHERE id = ?
                        AND   season = ?
                        AND   episode = ?
                        """, (episode[0], episode[2], episode[3]))
            print "Marked {0} Season {1} Episode {2} as watched".format(episode[1], episode[2], episode[3])
        con.commit()
        con.close()

    except IndexError:
        print "Invalid ID"

def getUnwatched():
    '''
    Get unwatched episodes

    Returns:
    (list) - list of unwatched episodes
    '''
    con = sql.connect(DATABASE_NAME)
    cur = con.cursor()
    cur.execute("SELECT shows.id, shows.title, season, episode, description, airDate FROM unwatched INNER JOIN shows on shows.id = unwatched.id WHERE watched = 0 ORDER BY airDate ASC, shows.title ASC;")
    data = cur.fetchall();
    con.close()

    return data

def listUnwatched():
    '''
    Display a list of unwatched episodes
    '''
    episodes = getUnwatched()

    if len(episodes) == 0:
        print "No unwatched episodes!"
    else:
        print " {0:3} | {1:20} | {2} | {3:10} | {4}".format("ID", "Show", "Episode", "Date", "Description")

        for idx, row in enumerate(episodes):
            title = (row[1][:MAX_TITLE_LENGTH-3] + '...') if len(row[1]) > MAX_TITLE_LENGTH else row[1]
            desc = (row[4][:MAX_DESC_LENGTH-3] + '...') if len(row[4]) > MAX_DESC_LENGTH else row[4]

            # Don't show ID for episodes not aired yet.
            epDate = toDate(row[5])
            if epDate is None or epDate > date.today():
                idx = ' - '

            print " {0:<3} | {1:20} | s{2:02} e{3:02} | {4} | {5}".format(idx, title, row[2], row[3], row[5], desc.encode('utf-8'))

            try:
                if toDate(row[5]) <= date.today() and toDate(episodes[idx+1][5]) > date.today():
                    print ''
            except:
                pass

def printHelp():
    '''
    Display help and available commands from help.txt
    '''
    try:
        with open('help.txt', 'r') as fin:
            print fin.read()
    except IOError:
        print "Help file not found :("

def initDatabase():
    '''
    Setup the SQLite database if it doesn't already exist
    '''
    con = sql.connect(DATABASE_NAME)
    with con:
        cur = con.cursor()

        cur.execute("""CREATE TABLE IF NOT EXISTS unwatched (
                    'id' INTEGER NOT NULL,
                    'season' INTEGER NOT NULL,
                    'episode' INTEGER NOT NULL,
                    'description' TEXT,
                    'airDate' TEXT NOT NULL,
                    'watched' INTEGER NOT NULL DEFAULT (0),
                    PRIMARY KEY(id, season, episode)
                    FOREIGN KEY(id) REFERENCES shows(id) ON DELETE CASCADE
                    )""")

        cur.execute("""CREATE TABLE IF NOT EXISTS shows (
                    'id' INTEGER NOT NULL,
                    'title' TEXT NOT NULL,
                    PRIMARY KEY (id)
                    )""")

        cur.execute("""CREATE TABLE IF NOT EXISTS settings (
                    'name' TEXT NOT NULL,
                    'value' TEXT NOT NULL,
                    PRIMARY KEY (name)
                    )""")
        cur.execute("INSERT OR IGNORE INTO settings VALUES('lastUpdate', '');");
        con.commit()

def getEpisodes(seriesIds, force = False):
    '''
    Get the episode list for shows
    Parameters:
    seriesIds - List of series IDs from TVRage

    Returns:
    List of episodes with the following information:
    id          - TVRage ID for the series
    series      - Name of the series
    season      - Season number
    episode     - Episode number in the season
    airDate     - Original air date
    description - Description for the episode
    '''

    # Checks if an update is necessary
    doUpdate = False
    if force:
        doUpdate = True
    else:
        con = sql.connect(DATABASE_NAME)
        cur = con.cursor()
        cur.execute("SELECT value FROM settings WHERE name = 'lastUpdate';")
        data = cur.fetchone()[0];
        con.close()
        if data:
            if toDate(data) < date.today():
                doUpdate = True
        else:
            doUpdate = True

    episodeList = []

    if doUpdate:
        print "Updating... "

        feedUrl = "http://www.tvrage.com/feeds/episode_list.php?sid={0}"

        for id in seriesIds:
            try:
                dom = minidom.parse(urllib2.urlopen("http://www.tvrage.com/feeds/episode_list.php?sid={0}".format(id), timeout=10.0))

                seriesName = dom.childNodes[0].childNodes[1].firstChild.nodeValue
                totalSeasons = int(dom.childNodes[0].childNodes[3].firstChild.nodeValue)

                print seriesName

                for season in dom.getElementsByTagName('Season'):
                    seasonNum = int(season.getAttribute('no'))
                    for episode in season.getElementsByTagName('episode'):

                        episodeNumFull = episode.getElementsByTagName('epnum')[0].firstChild.nodeValue
                        episodeNum = episode.getElementsByTagName('seasonnum')[0].firstChild.nodeValue
                        airdate = episode.getElementsByTagName('airdate')[0].firstChild.nodeValue
                        title = episode.getElementsByTagName('title')[0].firstChild.nodeValue

                        # Only get episodes that have an air date and those airing up until the next week
                        nextWeek = date.today() + timedelta(days=7)
                        if airdate != "0000-00-00" and toDate(airdate) <= nextWeek:
                            episodeList.append({
                                'id': id,
                                'series': seriesName,
                                'season': int(seasonNum),
                                'episode': int(episodeNum),
                                'airDate': airdate,
                                'description': title
                            })

                # Update last update date in database
                con = sql.connect(DATABASE_NAME)
                cur = con.cursor()
                cur.execute("INSERT OR REPLACE INTO settings VALUES('lastUpdate', ?);", (date.today(),))
                con.commit()
                con.close()
            except:
                print "Unexpected error or timeout"

    return episodeList

def toDate(dateStr):
    '''
    Convert a date string to a date object

    Parameters:
    dateStr - date string in yyyy-mm-dd format

    Returns:
    date object
    '''
    try:
        return datetime.strptime(dateStr, '%Y-%m-%d').date()
    except:
        # Sometimes the date is like 2010-00-00.. weird.
        pass
main()<|endoftext|>


<commit_before>import curses
import locale

# hack to get curses+unicode working
locale.setlocale(locale.LC_ALL, '')

################################################################################
# constants

ELLIPSE = u'\u2026'

# attributes
# why not bits? better tooling around strings
ATTRIBUTES = {
    'B': curses.A_BOLD,
    'S': curses.A_STANDOUT,
    'U': curses.A_UNDERLINE,
    'R': curses.A_REVERSE,  # reverse the colors
}

class Exit(Exception):
    pass

################################################################################
# utilities

def log(stuff):
    f = open('log', 'a')
    print >>f, stuff

################################################################################
# class hierarchy

class Application(object):
    current_window = None
    windows = {}
    modals = {}

    screen = None
    window = None

    colors = [('black', 'white')]

    def run(self):
        curses.wrapper(self.run_curses)

    def run_curses(self, screen):
        # init some curses geometry
        self.screen = screen

        # set up curses colors
        curses.use_default_colors()
        curses.curs_set(0)
        assert len(self.colors) <= 4, 'Too many colors defined'
        assert len(self.colors) > 0, 'Not enough colors defined'
        for i,c in enumerate(self.colors):
            background = getattr(curses, 'COLOR_' + c[0].upper())
            foreground = getattr(curses, 'COLOR_' + c[1].upper())
            curses.init_pair(i, background, foreground)

        # wait for a character for only 0.1s
        curses.halfdelay(1)

        # use the first window as the default window
        self.current_window = self.windows['default']
        self.render()

        while 1:
            input_value = screen.getch()
            # handle non-character input
            try:
                key = chr(input_value)
            except ValueError:
                key = input_value
            # pass the character input to the event loop
            try:
                self.handle(key)
            except Exit:
                return

    def get_size(self):
        return self.screen.getmaxyx()[1], self.screen.getmaxyx()[0]

    def render(self):
        self.screen.clear()
        size = self.get_size()
        w,h = size
        rlines, rstyles = self.current_window.render(size)
        for i,res in enumerate(zip(rlines, rstyles)):
            line, styles = res
            # convert to a unified representation
            style_chars = ['' for j in range(w)]
            for style, start, end in styles:
                for j in range(start, end):
                    style_chars[j] += style
            # convert back to ranges
            style_ranges = []
            current = None
            current_start = None
            for j,s in enumerate(style_chars):
                if s != current:
                    if current_start is not None:
                        style_ranges.append((current, current_start, j))
                    current_start = j
                    current = s
            style_ranges.append((current, current_start, j+1))
            # output the strings with the right styles
            for style, start, end in style_ranges:
                style_attr = 0
                for char in style:
                    style_attr |= ATTRIBUTES.get(char)
                # catch scrolling-past the end of the screen error
                try:
                    self.screen.addstr(i, start,
                                       line[start:end].encode('utf8'),
                                       style_attr)
                except curses.error as e:
                    if i < h - 1 and end == w:
                        raise e
        self.screen.refresh()

    def handle(self, key):
        # see if the application needs to run it's own handler
        signal = self.current_window.handle(key)
        if signal is None:
            # default behavior is to just quit
            if key == 'q':
                raise Exit
            elif key == curses.KEY_RESIZE:
                self.render()
        elif signal == 'redraw':
            self.render()

################################################################################

class Window(object):
    def handle(self, key):
        return

class Box(Window):
    title_parts = []
    option_parts = []
    spacing = 2

    window = None

    def __init__(self, window, title_parts=[], option_parts=[], spacing=2):
        self.window = window
        self.title_parts = title_parts
        self.option_parts = option_parts
        self.spacing = spacing

    def render(self, size):
        w,h = size
        rlines, rstyles = self.window.render((max(w-2,0),
                                              max(h-2,0) if h else h))
        # make decorative lines
        top_line = self.render_inset(self.title_parts, w-2)
        bot_line = self.render_inset(self.option_parts, w-2, align='right')
        top_line = u'\u250c' + top_line + u'\u2510'
        bot_line = u'\u2514' + bot_line + u'\u2518'
        # wrap content
        lines = [(u'\u2502' + line + u'\u2502') for line in rlines]
        lines = [top_line] + lines + [bot_line]
        # wrap up styles
        styles = [[(s[0], s[1]+1, s[2]+1) for s in line] for line in rstyles]
        styles = [[]] + styles + [[]]
        return lines, styles

    def render_inset(self, parts, width, align='left'):
        spacer = u'\u2500' * self.spacing
        accum = spacer
        for part in parts:
            accum += part
            accum += spacer
        # is it too long?
        if len(accum) > width:
            # can we just chop back the spacers?
            if len(accum.rstrip(u'\u2500')) <= width:
                accum = accum.rstrip(u'\u2500')
            else:
                # just lop it off, add an ellipse at the end (-1 for ellipse)
                accum = accum[:width - 1]
                accum += ELLIPSE
        # fill in extra, if needed
        lines = max(width - len(accum), 0) * u'\u2500'
        if align == 'right':  # assume it's left aligned by default
            accum = lines + accum
        else:
            accum += lines
        return accum

class VerticalPile(Window):
    windows = []
    current_window = None

    def __init__(self, *args, **kwargs):
        self.windows = args
        if kwargs.get('index'):
            index = kwargs.get('index')
        else:
            index = 0
        self.current_window = args[index]

    def render(self, size):
        w,h = size
        lines = []
        styles = []
        for win in self.windows:
            wlines, wstyles = win.render((w,h-len(lines)))
            lines.extend(wlines)
            styles.extend(wstyles)
            if len(lines) > h:
                lines = lines[:h]
                styles = styles[:h]
                break
        return lines, styles

    def handle(self, key):
        # don't do anything interesting with the signal
        signal = self.current_window.handle(key)
        return signal

class HorizontalPile(Window):
    pass

class List(Window):
    windows = []
    current_window = None
    # which window is current
    _index = None
    # how far the windows are, in terms of rows
    scroll = 0

    def __init__(self, *args):
        self.windows = list(args)
        if self._index:
            self.current_window = args[self._index]

    @property
    def index(self):
        return self._index

    @index.setter
    def index(self, index):
        # do bounds checking
        if index >= len(self.windows):
            index = len(self.windows) - 1
        elif index < 0:
            index = 0
        self._index = index
        self.current_window = self.windows[index]

    def add(self, window):
        self.windows.append(window)

    def render(self, size):
        w,h = size
        lines, styles, current_row, current_end = self.render_list((w-1,None))
        # just re-render short enough content
        if len(lines) < h:
            lines, styles, current_row, current_end = self.render_list((w,None))
            # double check
            if len(lines) < h:
                return lines, styles
        # draw a persistent scroll bar
        lines = [l + '|' for l in lines]
        percentage_scrolled = float(current_row) / len(lines)
        current_size = max(int(h * (float(current_end - current_row) /
                                    len(lines))),
                           1)
        # do current bounds checking
        if current_end > self.scroll + h:
            self.scroll = current_end - h
        if current_row < self.scroll:
            self.scroll = current_row
        # slice the right lines
        lines = lines[self.scroll:self.scroll + h]
        styles = styles[self.scroll:self.scroll + h]
        # place the scroll bar in the right place
        location = int(len(lines) * percentage_scrolled)
        for i in range(current_size):
            lines[location + i] = lines[location + i][:-1] + '#'

        return lines, styles

    def render_list(self, size):
        w,h = size
        lines = []
        styles = []
        current_row = 0
        current_end = 0
        for win in self.windows:
            wlines, wstyles = win.render((w,None))
            if win is self.current_window:
                current_row = len(lines)
                current_end = len(lines) + len(wlines)
                wstyles = [s + [('R', 0, w)] for s in wstyles]
            lines.extend(wlines)
            styles.extend(wstyles)
        return lines, styles, current_row, current_end

    def handle(self, key):
        if self.current_window is not None:
            signal = self.current_window.handle(key)
        else:
            signal = None
        if signal is None:
            if key in ('n', curses.KEY_UP, 'p', curses.KEY_DOWN):
                if self.index is None:
                    self.index = 0
                elif key in ('n', curses.KEY_DOWN):
                    self.index += 1
                elif key in ('p', curses.KEY_UP):
                    self.index -= 1
                return 'redraw'

################################################################################

class Text(Window):
    text = ''
    style = ''

    def __init__(self, text, style=''):
        self.text = text
        self.style = style

    def render(self, size):
        text = self.text
        texts = text.split('\n')
        w,h = size
        lines = []
        while texts and (len(lines) < h if h else True):
            lines.append(texts[0][:w])
            texts[0] = texts[0][w:]
            if not texts[0]:
                texts.pop(0)
        # pad everything out
        lines = [line + ' ' * (w - len(line)) for line in lines]
        # pad out styles
        styles = [[(self.style, 0, w)] for i in range(len(lines))]
        return lines, styles

class TextNoWrap(Window):
    text = ''
    style = ''

    def __init__(self, text, style=''):
        self.text = text
        self.style = style

    def render(self, size):
        text = self.text
        texts = text.split('\n')
        w,h = size
        lines = []
        while texts and (len(lines) < h if h else True):
            if len(texts[0]) > w:
                lines.append(texts[0][:w-1] + ELLIPSE)
            else:
                lines.append(texts[0])
            texts.pop(0)
        # pad everything out
        lines = [line + ' ' * (w - len(line)) for line in lines]
        # pad out styles
        styles = [[(self.style, 0, w)] for i in range(len(lines))]
        return lines, styles
<commit_msg>Fix the size of the scroll bar (too small)<commit_after>import curses
import locale

# hack to get curses+unicode working
locale.setlocale(locale.LC_ALL, '')

################################################################################
# constants

ELLIPSE = u'\u2026'

# attributes
# why not bits? better tooling around strings
ATTRIBUTES = {
    'B': curses.A_BOLD,
    'S': curses.A_STANDOUT,
    'U': curses.A_UNDERLINE,
    'R': curses.A_REVERSE,  # reverse the colors
}

class Exit(Exception):
    pass

################################################################################
# utilities

def log(stuff):
    f = open('log', 'a')
    print >>f, stuff

################################################################################
# class hierarchy

class Application(object):
    current_window = None
    windows = {}
    modals = {}

    screen = None
    window = None

    colors = [('black', 'white')]

    def run(self):
        curses.wrapper(self.run_curses)

    def run_curses(self, screen):
        # init some curses geometry
        self.screen = screen

        # set up curses colors
        curses.use_default_colors()
        curses.curs_set(0)
        assert len(self.colors) <= 4, 'Too many colors defined'
        assert len(self.colors) > 0, 'Not enough colors defined'
        for i,c in enumerate(self.colors):
            background = getattr(curses, 'COLOR_' + c[0].upper())
            foreground = getattr(curses, 'COLOR_' + c[1].upper())
            curses.init_pair(i, background, foreground)

        # wait for a character for only 0.1s
        curses.halfdelay(1)

        # use the first window as the default window
        self.current_window = self.windows['default']
        self.render()

        while 1:
            input_value = screen.getch()
            # handle non-character input
            try:
                key = chr(input_value)
            except ValueError:
                key = input_value
            # pass the character input to the event loop
            try:
                self.handle(key)
            except Exit:
                return

    def get_size(self):
        return self.screen.getmaxyx()[1], self.screen.getmaxyx()[0]

    def render(self):
        self.screen.clear()
        size = self.get_size()
        w,h = size
        rlines, rstyles = self.current_window.render(size)
        for i,res in enumerate(zip(rlines, rstyles)):
            line, styles = res
            # convert to a unified representation
            style_chars = ['' for j in range(w)]
            for style, start, end in styles:
                for j in range(start, end):
                    style_chars[j] += style
            # convert back to ranges
            style_ranges = []
            current = None
            current_start = None
            for j,s in enumerate(style_chars):
                if s != current:
                    if current_start is not None:
                        style_ranges.append((current, current_start, j))
                    current_start = j
                    current = s
            style_ranges.append((current, current_start, j+1))
            # output the strings with the right styles
            for style, start, end in style_ranges:
                style_attr = 0
                for char in style:
                    style_attr |= ATTRIBUTES.get(char)
                # catch scrolling-past the end of the screen error
                try:
                    self.screen.addstr(i, start,
                                       line[start:end].encode('utf8'),
                                       style_attr)
                except curses.error as e:
                    if i < h - 1 and end == w:
                        raise e
        self.screen.refresh()

    def handle(self, key):
        # see if the application needs to run it's own handler
        signal = self.current_window.handle(key)
        if signal is None:
            # default behavior is to just quit
            if key == 'q':
                raise Exit
            elif key == curses.KEY_RESIZE:
                self.render()
        elif signal == 'redraw':
            self.render()

################################################################################

class Window(object):
    def handle(self, key):
        return

class Box(Window):
    title_parts = []
    option_parts = []
    spacing = 2

    window = None

    def __init__(self, window, title_parts=[], option_parts=[], spacing=2):
        self.window = window
        self.title_parts = title_parts
        self.option_parts = option_parts
        self.spacing = spacing

    def render(self, size):
        w,h = size
        rlines, rstyles = self.window.render((max(w-2,0),
                                              max(h-2,0) if h else h))
        # make decorative lines
        top_line = self.render_inset(self.title_parts, w-2)
        bot_line = self.render_inset(self.option_parts, w-2, align='right')
        top_line = u'\u250c' + top_line + u'\u2510'
        bot_line = u'\u2514' + bot_line + u'\u2518'
        # wrap content
        lines = [(u'\u2502' + line + u'\u2502') for line in rlines]
        lines = [top_line] + lines + [bot_line]
        # wrap up styles
        styles = [[(s[0], s[1]+1, s[2]+1) for s in line] for line in rstyles]
        styles = [[]] + styles + [[]]
        return lines, styles

    def render_inset(self, parts, width, align='left'):
        spacer = u'\u2500' * self.spacing
        accum = spacer
        for part in parts:
            accum += part
            accum += spacer
        # is it too long?
        if len(accum) > width:
            # can we just chop back the spacers?
            if len(accum.rstrip(u'\u2500')) <= width:
                accum = accum.rstrip(u'\u2500')
            else:
                # just lop it off, add an ellipse at the end (-1 for ellipse)
                accum = accum[:width - 1]
                accum += ELLIPSE
        # fill in extra, if needed
        lines = max(width - len(accum), 0) * u'\u2500'
        if align == 'right':  # assume it's left aligned by default
            accum = lines + accum
        else:
            accum += lines
        return accum

class VerticalPile(Window):
    windows = []
    current_window = None

    def __init__(self, *args, **kwargs):
        self.windows = args
        if kwargs.get('index'):
            index = kwargs.get('index')
        else:
            index = 0
        self.current_window = args[index]

    def render(self, size):
        w,h = size
        lines = []
        styles = []
        for win in self.windows:
            wlines, wstyles = win.render((w,h-len(lines)))
            lines.extend(wlines)
            styles.extend(wstyles)
            if len(lines) > h:
                lines = lines[:h]
                styles = styles[:h]
                break
        return lines, styles

    def handle(self, key):
        # don't do anything interesting with the signal
        signal = self.current_window.handle(key)
        return signal

class HorizontalPile(Window):
    pass

class List(Window):
    windows = []
    current_window = None
    # which window is current
    _index = None
    # how far the windows are, in terms of rows
    scroll = 0

    def __init__(self, *args):
        self.windows = list(args)
        if self._index:
            self.current_window = args[self._index]

    @property
    def index(self):
        return self._index

    @index.setter
    def index(self, index):
        # do bounds checking
        if index >= len(self.windows):
            index = len(self.windows) - 1
        elif index < 0:
            index = 0
        self._index = index
        self.current_window = self.windows[index]

    def add(self, window):
        self.windows.append(window)

    def render(self, size):
        w,h = size
        lines, styles, current_row, current_end = self.render_list((w-1,None))
        # just re-render short enough content
        if len(lines) < h:
            lines, styles, current_row, current_end = self.render_list((w,None))
            # double check
            if len(lines) < h:
                return lines, styles
        # draw a persistent scroll bar
        lines = [l + '|' for l in lines]
        current_size = max(int(h * (float(h) / len(lines))), 1)
        # do current bounds checking
        if current_end > self.scroll + h:
            self.scroll = current_end - h
        if current_row < self.scroll:
            self.scroll = current_row
        percentage_scrolled = float(max(self.scroll, 0)) / len(lines)
        # slice the right lines
        lines = lines[self.scroll:self.scroll + h]
        styles = styles[self.scroll:self.scroll + h]
        # place the scroll bar in the right place
        location = int(len(lines) * percentage_scrolled)
        for i in range(location, min(location + current_size, h)):
            # can't assign to tuple index
            lines[i] = lines[i][:-1] + '#'

        return lines, styles

    def render_list(self, size):
        w,h = size
        lines = []
        styles = []
        current_row = 0
        current_end = 0
        for win in self.windows:
            wlines, wstyles = win.render((w,None))
            if win is self.current_window:
                current_row = len(lines)
                current_end = len(lines) + len(wlines)
                wstyles = [s + [('R', 0, w)] for s in wstyles]
            lines.extend(wlines)
            styles.extend(wstyles)
        return lines, styles, current_row, current_end

    def handle(self, key):
        if self.current_window is not None:
            signal = self.current_window.handle(key)
        else:
            signal = None
        if signal is None:
            if key in ('n', curses.KEY_UP, 'p', curses.KEY_DOWN):
                if self.index is None:
                    self.index = 0
                elif key in ('n', curses.KEY_DOWN):
                    self.index += 1
                elif key in ('p', curses.KEY_UP):
                    self.index -= 1
                return 'redraw'

################################################################################

class Text(Window):
    text = ''
    style = ''

    def __init__(self, text, style=''):
        self.text = text
        self.style = style

    def render(self, size):
        text = self.text
        texts = text.split('\n')
        w,h = size
        lines = []
        while texts and (len(lines) < h if h else True):
            lines.append(texts[0][:w])
            texts[0] = texts[0][w:]
            if not texts[0]:
                texts.pop(0)
        # pad everything out
        lines = [line + ' ' * (w - len(line)) for line in lines]
        # pad out styles
        styles = [[(self.style, 0, w)] for i in range(len(lines))]
        return lines, styles

class TextNoWrap(Window):
    text = ''
    style = ''

    def __init__(self, text, style=''):
        self.text = text
        self.style = style

    def render(self, size):
        text = self.text
        texts = text.split('\n')
        w,h = size
        lines = []
        while texts and (len(lines) < h if h else True):
            if len(texts[0]) > w:
                lines.append(texts[0][:w-1] + ELLIPSE)
            else:
                lines.append(texts[0])
            texts.pop(0)
        # pad everything out
        lines = [line + ' ' * (w - len(line)) for line in lines]
        # pad out styles
        styles = [[(self.style, 0, w)] for i in range(len(lines))]
        return lines, styles
<|endoftext|>