<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
import re
import os
from collections import OrderedDict
import json

import six
from six.moves import zip

import numpy as np
from astropy.io import fits, ascii
from astropy.table import Column
import pyyaks.context
from cxotime import CxoTime
from chandra_aca.aca_image import ACAImage

from chandra_aca.dark_model import dark_temp_scale, DARK_SCALE_4C
from mica.cache import lru_cache
from mica.common import MICA_ARCHIVE_PATH, MissingDataError
from . import file_defs

DARK_CAL = pyyaks.context.ContextDict('dark_cal')

MICA_FILES = pyyaks.context.ContextDict('update_mica_files',
                                        basedir=MICA_ARCHIVE_PATH)
MICA_FILES.update(file_defs.MICA_FILES)

def date_to_dark_id(date):
    """
    Convert ``date`` to the corresponding YYYYDOY format for a dark cal identifiers.

    :param date: any CxoTime compatible format
    :returns: dark id (YYYYDOY)
    """
    t = CxoTime(date)
    shape = np.shape(t)
    if not shape:
        return t.date[:4] + t.date[5:8]
    b = t.date.view((str, 1)).reshape(-1, 21)
    b = np.hstack([b[:, :4], b[:, 5:8]])
    b = np.frombuffer(b.tobytes(), dtype=(str, 7))
    return b.reshape(shape)

def dark_id_to_date(dark_id):
    """
    Convert ``dark_id`` (YYYYDOY) to the corresponding CxoTime 'date' format.

    :param date: dark id (YYYYDOY)
    :returns: str in CxoTime 'date' format
    """
    shape = np.shape(dark_id)
    if not shape:
        return '{}:{}'.format(dark_id[:4], dark_id[4:])
    b = np.atleast_1d(dark_id).view((str, 1)).reshape(-1, 7)
    sep = np.array([':'] * len(b))[:, None]
    b = np.hstack([b[:, :4], sep, b[:, 4:]])
    b = np.frombuffer(b.tobytes(), dtype=(str, 8))
    return b.reshape(shape)

@lru_cache()
def get_dark_cal_dirs(dark_cals_dir=MICA_FILES['dark_cals_dir'].abs):
    """
    Get an ordered dict of directory paths containing dark current calibration files,
    where the key is the dark cal identifier (YYYYDOY) and the value is the path.

    :param source: source of dark cal directories ('mica'|'ska')
    :returns: ordered dict of absolute directory paths
    """
    dark_cal_ids = sorted([fn for fn in os.listdir(dark_cals_dir)
                           if re.match(r'[12]\d{6}$', fn)])
    dark_cal_dirs = [os.path.join(dark_cals_dir, id_)
                     for id_ in dark_cal_ids]
    return OrderedDict(zip(dark_cal_ids, dark_cal_dirs))

def get_dark_cal_id(date, select='before'):
    """
    Return the dark calibration id corresponding to ``date``.

    If ``select`` is ``'before'`` (default) then use the first calibration which
    occurs before ``date``.  Other valid options are ``'after'`` and ``'nearest'``.

    :param date: date in any CxoTime format
    :param select: method to select dark cal (before|nearest|after)

    :returns: dark cal id string (YYYYDOY)
    """
    return _get_dark_cal_id_vector(date, select=select)

def _get_dark_cal_id_scalar(date, select='before'):
    dark_cals = get_dark_cal_dirs()
    dark_id = date_to_dark_id(date)

    # Special case if dark_id is exactly an existing dark cal then return that dark
    # cal regardless of the select method.
    if dark_id in dark_cals:
        return dark_id

    dark_cal_ids = list(dark_cals.keys())
    date_secs = CxoTime(date).secs
    dark_cal_secs = CxoTime(np.array([dark_id_to_date(id_) for id_ in dark_cal_ids])).secs

    if select == 'nearest':
        ii = np.argmin(np.abs(dark_cal_secs - date_secs))
    elif select in ('before', 'after'):
        ii = np.searchsorted(dark_cal_secs, date_secs)
        if select == 'before':
            ii -= 1
    else:
        raise ValueError('select arg must be one of "nearest", "before", or "after"')

    try:
        out_dark_id = dark_cal_ids[ii]
    except IndexError:
        raise MissingDataError('No dark cal found {} {}'.format(select, date))

    return out_dark_id

_get_dark_cal_id_vector = np.vectorize(_get_dark_cal_id_scalar, excluded=['select'])

@DARK_CAL.cache
def _get_dark_cal_image_props(date, select='before', t_ccd_ref=None, aca_image=False,
                              allow_negative=False):
    """
    Return the dark calibration image (e-/s) nearest to ``date`` and the corresponding
    dark_props file.

    :param date: date in any CxoTime format
    :param select: method to select dark cal (before|nearest|after)
    :param t_ccd_ref: rescale dark map to temperature (degC, default=no scaling)
    :param aca_image: return an AcaImage instance
    :param allow_negative: allow negative values in raw dark map (default=False)

    :returns: 1024 x 1024 ndarray with dark cal image in e-/s, props dict
    """
    DARK_CAL['id'] = get_dark_cal_id(date, select)

    with fits.open(MICA_FILES['dark_image.fits'].abs, memmap=False) as hdus:
        dark = hdus[0].data
        # Recast as native byte ordering since FITS is typically not. This
        # statement is normally the same as dark.astype(np.float32).
        dark = dark.astype(dark.dtype.type)

    with open(MICA_FILES['dark_props.json'].abs, 'r') as fh:
        props = json.load(fh)

    # Change unicode to ascii at top level
    if six.PY2:
        keys = list(props.keys())
        asciiprops = {}
        for key in keys:
            asciiprops[str(key)] = props[key]
        props = asciiprops

    if t_ccd_ref is not None:
        # Scale factor to adjust data to an effective temperature of t_ccd_ref.
        # For t_ccds warmer than t_ccd_ref this scale factor is < 1, i.e. the
        # observed dark current is made smaller to match what it would be at the
        # lower reference temperature.
        t_ccd = props['ccd_temp']
        dark *= dark_temp_scale(t_ccd, t_ccd_ref)

    if not allow_negative:
        np.clip(dark, a_min=0, a_max=None, out=dark)

    if aca_image:
        dark = ACAImage(dark, row0=-512, col0=-512)

    return dark, props

def get_dark_cal_image(date, select='before', t_ccd_ref=None, aca_image=False,
                       allow_negative=False):
    """
    Return the dark calibration image (e-/s) nearest to ``date``.

    If ``select`` is ``'before'`` (default) then use the first calibration which
    occurs before ``date``.  Other valid options are ``'after'`` and ``'nearest'``.

    :param date: date in any CxoTime format
    :param select: method to select dark cal (before|nearest|after)
    :param t_ccd_ref: rescale dark map to temperature (degC, default=no scaling)
    :param aca_image: return an ACAImage instance instead of ndarray
    :param allow_negative: allow negative values in raw dark map (default=False)

    :returns: 1024 x 1024 ndarray with dark cal image in e-/s
    """
    dark, props = _get_dark_cal_image_props(date, select=select, t_ccd_ref=t_ccd_ref,
                                            aca_image=aca_image,
                                            allow_negative=allow_negative)
    return dark

def get_dark_cal_props(date, select='before', include_image=False, t_ccd_ref=None,
                       aca_image=False, allow_negative=False):
    """
    Return a dark calibration properties structure for ``date``

    If ``select`` is ``'before'`` (default) then use the first calibration which
    occurs before ``date``.  Other valid options are ``'after'`` and ``'nearest'``.

    If ``include_image`` is True then an additional column or key ``image`` is
    defined which contains the corresponding 1024x1024 dark cal image.

    :param date: date in any CxoTime format
    :param select: method to select dark cal (before|nearest|after)
    :param include_image: include the dark cal images in output (default=False)
    :param t_ccd_ref: rescale dark map to temperature (degC, default=no scaling)
    :param aca_image: return an ACAImage instance instead of ndarray
    :param allow_negative: allow negative values in raw dark map (default=False)

    :returns: dict of dark calibration properties
    """
    dark, props = _get_dark_cal_image_props(date, select=select, t_ccd_ref=None,
                                            aca_image=aca_image,
                                            allow_negative=allow_negative)

    if include_image:
        props['image'] = dark

    return props

def get_dark_cal_props_table(start=None, stop=None, include_image=False, as_table=True):
    """
    Return a table of dark calibration properties between ``start`` and ``stop``.

    If ``include_image`` is True then an additional column or key ``image`` is
    defined which contains the corresponding 1024x1024 dark cal image.

    If ``as_table`` is True (default) then the result is an astropy Table object.
    If False then a list of dicts is returned.  In this case the full contents
    of the properties file including replica properties is available.

    :param start: start time (default=beginning of mission)
    :param stop: stop time (default=now)
    :param include_image: include the dark cal images in output (default=False)
    :param as_table: return a Table instead of a list (default=True)

    :returns: astropy Table or list of dark calibration properties
    """
    start_id = date_to_dark_id('1999:001:12:00:00' if start is None else start)
    stop_id = date_to_dark_id(stop)
    dark_dirs = [dark_id for dark_id in get_dark_cal_dirs()
                 if dark_id >= start_id and dark_id <= stop_id]

    # Get the list of properties structures
    props = [get_dark_cal_props(dark_id, include_image=include_image) for dark_id in dark_dirs]

    if as_table:
        # Get rid of a non-scalar data structures and collect col names
        names = []
        for prop in props:
            for key, val in prop.copy().items():
                if isinstance(val, (list, tuple, dict)):
                    del prop[key]
                elif key not in names:
                    names.append(key)

        lines = []
        lines.append(','.join(names))
        for prop in props:
            vals = [str(prop.get(name, '')) for name in names]
            lines.append(','.join(vals))
        table_props = ascii.read(lines, format='csv')

        if include_image:
            x = np.vstack([prop['image'][np.newaxis, :] for prop in props])
            images = Column(x)
            table_props['image'] = images
        props = table_props

    return props
<commit_msg>flake<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
import re
import os
from collections import OrderedDict
import json

import six
from six.moves import zip

import numpy as np
from astropy.io import fits, ascii
from astropy.table import Column
import pyyaks.context
from cxotime import CxoTime
from chandra_aca.aca_image import ACAImage

from chandra_aca.dark_model import dark_temp_scale
from chandra_aca.dark_model import DARK_SCALE_4C  # noqa
from mica.cache import lru_cache
from mica.common import MICA_ARCHIVE_PATH, MissingDataError
from . import file_defs

DARK_CAL = pyyaks.context.ContextDict('dark_cal')

MICA_FILES = pyyaks.context.ContextDict('update_mica_files',
                                        basedir=MICA_ARCHIVE_PATH)
MICA_FILES.update(file_defs.MICA_FILES)

def date_to_dark_id(date):
    """
    Convert ``date`` to the corresponding YYYYDOY format for a dark cal identifiers.

    :param date: any CxoTime compatible format
    :returns: dark id (YYYYDOY)
    """
    t = CxoTime(date)
    shape = np.shape(t)
    if not shape:
        return t.date[:4] + t.date[5:8]
    b = t.date.view((str, 1)).reshape(-1, 21)
    b = np.hstack([b[:, :4], b[:, 5:8]])
    b = np.frombuffer(b.tobytes(), dtype=(str, 7))
    return b.reshape(shape)

def dark_id_to_date(dark_id):
    """
    Convert ``dark_id`` (YYYYDOY) to the corresponding CxoTime 'date' format.

    :param date: dark id (YYYYDOY)
    :returns: str in CxoTime 'date' format
    """
    shape = np.shape(dark_id)
    if not shape:
        return '{}:{}'.format(dark_id[:4], dark_id[4:])
    b = np.atleast_1d(dark_id).view((str, 1)).reshape(-1, 7)
    sep = np.array([':'] * len(b))[:, None]
    b = np.hstack([b[:, :4], sep, b[:, 4:]])
    b = np.frombuffer(b.tobytes(), dtype=(str, 8))
    return b.reshape(shape)

@lru_cache()
def get_dark_cal_dirs(dark_cals_dir=MICA_FILES['dark_cals_dir'].abs):
    """
    Get an ordered dict of directory paths containing dark current calibration files,
    where the key is the dark cal identifier (YYYYDOY) and the value is the path.

    :param source: source of dark cal directories ('mica'|'ska')
    :returns: ordered dict of absolute directory paths
    """
    dark_cal_ids = sorted([fn for fn in os.listdir(dark_cals_dir)
                           if re.match(r'[12]\d{6}$', fn)])
    dark_cal_dirs = [os.path.join(dark_cals_dir, id_)
                     for id_ in dark_cal_ids]
    return OrderedDict(zip(dark_cal_ids, dark_cal_dirs))

def get_dark_cal_id(date, select='before'):
    """
    Return the dark calibration id corresponding to ``date``.

    If ``select`` is ``'before'`` (default) then use the first calibration which
    occurs before ``date``.  Other valid options are ``'after'`` and ``'nearest'``.

    :param date: date in any CxoTime format
    :param select: method to select dark cal (before|nearest|after)

    :returns: dark cal id string (YYYYDOY)
    """
    return _get_dark_cal_id_vector(date, select=select)

def _get_dark_cal_id_scalar(date, select='before'):
    dark_cals = get_dark_cal_dirs()
    dark_id = date_to_dark_id(date)

    # Special case if dark_id is exactly an existing dark cal then return that dark
    # cal regardless of the select method.
    if dark_id in dark_cals:
        return dark_id

    dark_cal_ids = list(dark_cals.keys())
    date_secs = CxoTime(date).secs
    dark_cal_secs = CxoTime(np.array([dark_id_to_date(id_) for id_ in dark_cal_ids])).secs

    if select == 'nearest':
        ii = np.argmin(np.abs(dark_cal_secs - date_secs))
    elif select in ('before', 'after'):
        ii = np.searchsorted(dark_cal_secs, date_secs)
        if select == 'before':
            ii -= 1
    else:
        raise ValueError('select arg must be one of "nearest", "before", or "after"')

    try:
        out_dark_id = dark_cal_ids[ii]
    except IndexError:
        raise MissingDataError('No dark cal found {} {}'.format(select, date))

    return out_dark_id

_get_dark_cal_id_vector = np.vectorize(_get_dark_cal_id_scalar, excluded=['select'])

@DARK_CAL.cache
def _get_dark_cal_image_props(date, select='before', t_ccd_ref=None, aca_image=False,
                              allow_negative=False):
    """
    Return the dark calibration image (e-/s) nearest to ``date`` and the corresponding
    dark_props file.

    :param date: date in any CxoTime format
    :param select: method to select dark cal (before|nearest|after)
    :param t_ccd_ref: rescale dark map to temperature (degC, default=no scaling)
    :param aca_image: return an AcaImage instance
    :param allow_negative: allow negative values in raw dark map (default=False)

    :returns: 1024 x 1024 ndarray with dark cal image in e-/s, props dict
    """
    DARK_CAL['id'] = get_dark_cal_id(date, select)

    with fits.open(MICA_FILES['dark_image.fits'].abs, memmap=False) as hdus:
        dark = hdus[0].data
        # Recast as native byte ordering since FITS is typically not. This
        # statement is normally the same as dark.astype(np.float32).
        dark = dark.astype(dark.dtype.type)

    with open(MICA_FILES['dark_props.json'].abs, 'r') as fh:
        props = json.load(fh)

    # Change unicode to ascii at top level
    if six.PY2:
        keys = list(props.keys())
        asciiprops = {}
        for key in keys:
            asciiprops[str(key)] = props[key]
        props = asciiprops

    if t_ccd_ref is not None:
        # Scale factor to adjust data to an effective temperature of t_ccd_ref.
        # For t_ccds warmer than t_ccd_ref this scale factor is < 1, i.e. the
        # observed dark current is made smaller to match what it would be at the
        # lower reference temperature.
        t_ccd = props['ccd_temp']
        dark *= dark_temp_scale(t_ccd, t_ccd_ref)

    if not allow_negative:
        np.clip(dark, a_min=0, a_max=None, out=dark)

    if aca_image:
        dark = ACAImage(dark, row0=-512, col0=-512)

    return dark, props

def get_dark_cal_image(date, select='before', t_ccd_ref=None, aca_image=False,
                       allow_negative=False):
    """
    Return the dark calibration image (e-/s) nearest to ``date``.

    If ``select`` is ``'before'`` (default) then use the first calibration which
    occurs before ``date``.  Other valid options are ``'after'`` and ``'nearest'``.

    :param date: date in any CxoTime format
    :param select: method to select dark cal (before|nearest|after)
    :param t_ccd_ref: rescale dark map to temperature (degC, default=no scaling)
    :param aca_image: return an ACAImage instance instead of ndarray
    :param allow_negative: allow negative values in raw dark map (default=False)

    :returns: 1024 x 1024 ndarray with dark cal image in e-/s
    """
    dark, props = _get_dark_cal_image_props(date, select=select, t_ccd_ref=t_ccd_ref,
                                            aca_image=aca_image,
                                            allow_negative=allow_negative)
    return dark

def get_dark_cal_props(date, select='before', include_image=False, t_ccd_ref=None,
                       aca_image=False, allow_negative=False):
    """
    Return a dark calibration properties structure for ``date``

    If ``select`` is ``'before'`` (default) then use the first calibration which
    occurs before ``date``.  Other valid options are ``'after'`` and ``'nearest'``.

    If ``include_image`` is True then an additional column or key ``image`` is
    defined which contains the corresponding 1024x1024 dark cal image.

    :param date: date in any CxoTime format
    :param select: method to select dark cal (before|nearest|after)
    :param include_image: include the dark cal images in output (default=False)
    :param t_ccd_ref: rescale dark map to temperature (degC, default=no scaling)
    :param aca_image: return an ACAImage instance instead of ndarray
    :param allow_negative: allow negative values in raw dark map (default=False)

    :returns: dict of dark calibration properties
    """
    dark, props = _get_dark_cal_image_props(date, select=select, t_ccd_ref=None,
                                            aca_image=aca_image,
                                            allow_negative=allow_negative)

    if include_image:
        props['image'] = dark

    return props

def get_dark_cal_props_table(start=None, stop=None, include_image=False, as_table=True):
    """
    Return a table of dark calibration properties between ``start`` and ``stop``.

    If ``include_image`` is True then an additional column or key ``image`` is
    defined which contains the corresponding 1024x1024 dark cal image.

    If ``as_table`` is True (default) then the result is an astropy Table object.
    If False then a list of dicts is returned.  In this case the full contents
    of the properties file including replica properties is available.

    :param start: start time (default=beginning of mission)
    :param stop: stop time (default=now)
    :param include_image: include the dark cal images in output (default=False)
    :param as_table: return a Table instead of a list (default=True)

    :returns: astropy Table or list of dark calibration properties
    """
    start_id = date_to_dark_id('1999:001:12:00:00' if start is None else start)
    stop_id = date_to_dark_id(stop)
    dark_dirs = [dark_id for dark_id in get_dark_cal_dirs()
                 if dark_id >= start_id and dark_id <= stop_id]

    # Get the list of properties structures
    props = [get_dark_cal_props(dark_id, include_image=include_image) for dark_id in dark_dirs]

    if as_table:
        # Get rid of a non-scalar data structures and collect col names
        names = []
        for prop in props:
            for key, val in prop.copy().items():
                if isinstance(val, (list, tuple, dict)):
                    del prop[key]
                elif key not in names:
                    names.append(key)

        lines = []
        lines.append(','.join(names))
        for prop in props:
            vals = [str(prop.get(name, '')) for name in names]
            lines.append(','.join(vals))
        table_props = ascii.read(lines, format='csv')

        if include_image:
            x = np.vstack([prop['image'][np.newaxis, :] for prop in props])
            images = Column(x)
            table_props['image'] = images
        props = table_props

    return props
<|endoftext|>


<commit_before>import numpy as np
from Quaternion import Quat, normalize
from Ska.engarchive import fetch
from kadi import events
import pytest

from .. import asp_l1

def compare_obc_and_asol(atts, times, recs, ptol=2, ytol=2, rtol=50):
    """
    Check that obc solution and ground solution have relatively close quats.
    Note that because the ground aspect solution is not continuous (only run over science obsids
    in aspect intervals) that this method only checks those intervals.

    :param atts: attitudes from get_atts
    :param times: times from get_atts
    :param recs: list of dicts of header records from get_atts
    :param ptol: dq quaternion pitch tolerance in arcsecs
    :param ytol: dq quaternion yaw tolerance in arcsecs
    :param rtol: dq quaternion roll tolerance in arcsecs
    """
    for ai in recs:
        telem = fetch.Msidset(['aoattqt*'], ai['TSTART'], ai['TSTOP'])
        obc_atts = np.vstack([telem['aoattqt{}'.format(idx)].vals
                          for idx in [1, 2, 3, 4]]).transpose()
        obc_times = telem['aoattqt1'].times
        # Test that, at the obc times, the onboard solution and the ground solution
        # are reasonably close
        idxs = np.searchsorted(times[:-1], obc_times)
        dps = []
        dys = []
        drs = []
        for att, obc_att in zip(atts[idxs], obc_atts):
            dq = Quat(normalize(att)).dq(Quat(normalize(obc_att)))
            dps.append(dq.pitch * 3600)
            dys.append(dq.yaw * 3600)
            drs.append(dq.roll0 * 3600)
        dps = np.array(dps)
        dys = np.array(dys)
        drs = np.array(drs)
        assert np.all(np.abs(dys) < ytol)
        assert np.all(np.abs(dps) < ptol)
        assert np.all(np.abs(drs) < rtol)

@pytest.mark.parametrize("obsid", [14333, 15175, 5438, 2121])
def test_get_atts_obsid(obsid):
    atts, times, recs = asp_l1.get_atts(obsid=obsid)
    compare_obc_and_asol(atts, times, recs)

def test_get_atts_time():
    start = '2014:001:00:00:00.000'
    stop = '2014:005:00:00:00.000'
    atts, times, recs = asp_l1.get_atts(start=start, stop=stop)
    assert len(atts) == len(times)
    compare_obc_and_asol(atts, times, recs, rtol=51)
    dwells = events.dwells.filter(start, stop)
    for dwell in dwells:
        if dwell.get_obsid() > 40000:
            continue
        # check that more than 90% of the kalman interval is in times fetched from get_atts
        ok = (times < dwell.tstop) & (times > dwell.tstart)
        assert (times[ok][-1] - times[ok][0]) > dwell.dur * .90
        # also assert that the number of ~.25sec samples works out
        assert (len(times[ok]) * .25625) > dwell.dur * .90

def test_get_atts_filter():
    # Obsid 19039 has a momentum dump that shows up in asp_sol_status
    atts, times, recs = asp_l1.get_atts(obsid=19039)
    uf_atts, uf_times, uf_recs = asp_l1.get_atts(obsid=19039, filter=False)
    # Confirm that approximately 212 seconds are filtered
    assert np.abs((len(uf_atts) - len(atts)) * .25625 - 212.2) < 5
<commit_msg>Set OR/ER boundary to 38000 in one aspl1 test<commit_after>import numpy as np
from Quaternion import Quat, normalize
from Ska.engarchive import fetch
from kadi import events
import pytest

from .. import asp_l1

def compare_obc_and_asol(atts, times, recs, ptol=2, ytol=2, rtol=50):
    """
    Check that obc solution and ground solution have relatively close quats.
    Note that because the ground aspect solution is not continuous (only run over science obsids
    in aspect intervals) that this method only checks those intervals.

    :param atts: attitudes from get_atts
    :param times: times from get_atts
    :param recs: list of dicts of header records from get_atts
    :param ptol: dq quaternion pitch tolerance in arcsecs
    :param ytol: dq quaternion yaw tolerance in arcsecs
    :param rtol: dq quaternion roll tolerance in arcsecs
    """
    for ai in recs:
        telem = fetch.Msidset(['aoattqt*'], ai['TSTART'], ai['TSTOP'])
        obc_atts = np.vstack([telem['aoattqt{}'.format(idx)].vals
                          for idx in [1, 2, 3, 4]]).transpose()
        obc_times = telem['aoattqt1'].times
        # Test that, at the obc times, the onboard solution and the ground solution
        # are reasonably close
        idxs = np.searchsorted(times[:-1], obc_times)
        dps = []
        dys = []
        drs = []
        for att, obc_att in zip(atts[idxs], obc_atts):
            dq = Quat(normalize(att)).dq(Quat(normalize(obc_att)))
            dps.append(dq.pitch * 3600)
            dys.append(dq.yaw * 3600)
            drs.append(dq.roll0 * 3600)
        dps = np.array(dps)
        dys = np.array(dys)
        drs = np.array(drs)
        assert np.all(np.abs(dys) < ytol)
        assert np.all(np.abs(dps) < ptol)
        assert np.all(np.abs(drs) < rtol)

@pytest.mark.parametrize("obsid", [14333, 15175, 5438, 2121])
def test_get_atts_obsid(obsid):
    atts, times, recs = asp_l1.get_atts(obsid=obsid)
    compare_obc_and_asol(atts, times, recs)

def test_get_atts_time():
    start = '2014:001:00:00:00.000'
    stop = '2014:005:00:00:00.000'
    atts, times, recs = asp_l1.get_atts(start=start, stop=stop)
    assert len(atts) == len(times)
    compare_obc_and_asol(atts, times, recs, rtol=51)
    dwells = events.dwells.filter(start, stop)
    for dwell in dwells:
        if dwell.get_obsid() > 38000:
            continue
        # check that more than 90% of the kalman interval is in times fetched from get_atts
        ok = (times < dwell.tstop) & (times > dwell.tstart)
        assert (times[ok][-1] - times[ok][0]) > dwell.dur * .90
        # also assert that the number of ~.25sec samples works out
        assert (len(times[ok]) * .25625) > dwell.dur * .90

def test_get_atts_filter():
    # Obsid 19039 has a momentum dump that shows up in asp_sol_status
    atts, times, recs = asp_l1.get_atts(obsid=19039)
    uf_atts, uf_times, uf_recs = asp_l1.get_atts(obsid=19039, filter=False)
    # Confirm that approximately 212 seconds are filtered
    assert np.abs((len(uf_atts) - len(atts)) * .25625 - 212.2) < 5
<|endoftext|>


<commit_before># Miro Community - Easiest way to make a video website
#
# Copyright (C) 2010, 2011, 2012 Participatory Culture Foundation
#
# Miro Community is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# Miro Community is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with Miro Community.  If not, see <http://www.gnu.org/licenses/>.

import datetime
import urllib

from django import forms
from django.conf import settings
from django.contrib.auth.models import User
from django.contrib.sites.models import Site
from django.core.exceptions import ValidationError
from django.core.urlresolvers import reverse, reverse_lazy
from django.forms.models import modelformset_factory
from django.template.defaultfilters import pluralize
from localtv.admin.forms import (EditSettingsForm as _EditSettingsForm,
                                 AuthorForm as _AuthorForm,
                                 BulkEditVideoFormSet as _BulkEditVideoFormSet,
                                 BulkEditVideoForm)
from localtv.models import Video
from paypal.standard.conf import (POSTBACK_ENDPOINT,
                                  SANDBOX_POSTBACK_ENDPOINT,
                                  RECEIVER_EMAIL)
from paypal.standard.forms import PayPalPaymentsForm

from mirocommunity_saas.models import SiteTierInfo, Tier
from mirocommunity_saas.utils.mail import send_welcome_email
from mirocommunity_saas.utils.tiers import (make_tier_change_token,
                                            check_tier_change_token)

class EditSettingsForm(_EditSettingsForm):
    def __init__(self, *args, **kwargs):
        _EditSettingsForm.__init__(self, *args, **kwargs)
        self.tier = SiteTierInfo.objects.get_current().tier
        if not self.tier.custom_css:
            # Uh-oh: custom CSS is not permitted!
            #
            # To handle only letting certain paid users edit CSS,
            # we do two things.
            #
            # 1. Cosmetically, we set the CSS editing box's CSS class
            # to be 'hidden'. (We have some CSS that makes it not show
            # up.) SB: Not sure why we don't just remove the field...
            css_field = self.fields['css']
            css_field.label += ' (upgrade to enable this form field)'
            css_field.widget.attrs['readonly'] = True
            #
            # 2. In validation, we make sure that changing the CSS is
            # rejected as invalid if the site does not have permission
            # to do that.

    def clean_css(self):
        css = self.cleaned_data.get('css')
        # Does the current tier permit custom CSS? If so, return the data the
        # user submitted.
        if self.tier.custom_css:
            return css # no questions asked

        # We permit the value if it's the same as self.instance has:
        # SB: Leaving this for now, for backwards-compatibility, but again:
        # if css isn't allowed, we should probably not be storing it. This is
        # probably only here because we didn't remove the field on __init__.
        if self.instance.css == css:
            return css

        # Otherwise, reject the change.
        self.data['css'] = self.instance.css
        raise ValidationError(
            "To edit CSS for your site, you have to upgrade.")

class AuthorForm(_AuthorForm):
    def __init__(self, *args, **kwargs):
        _AuthorForm.__init__(self, *args, **kwargs)
        ## Add a note to the 'role' help text indicating how many admins
        ## are permitted with this kind of account.
        self.tier = SiteTierInfo.objects.get_current().tier
        if self.tier.admin_limit is not None:
            # For backwards-compatibility, pretend the site owner (a
            # superuser) counts toward the limit.
            limit = self.tier.admin_limit + 1
            message = ('With a {tier_name}, you may have {limit} '
                       'administrator{s}.').format(tier_name=self.tier.name,
                                                   limit=limit,
                                                   s=pluralize(limit))
            self.fields['role'].help_text = message

    def clean_role(self):
        role = self.cleaned_data['role']

        # Nothing to check if they aren't being set to admin.
        if role != 'admin':
            return role

        # Nothing to check if there is no admin limit.
        if self.tier.admin_limit is None:
            return role

        # Nothing to check if the user is already an admin.
        if 'role' not in self.changed_data:
            return role

        # And finally, we're good if there's room for another admin.
        admin_count = self.site_settings.admins.exclude(is_superuser=True
                                              ).exclude(is_active=False
                                              ).count()
        if (admin_count + 1) <= self.tier.admin_limit:
            return role

        # For backwards-compatibility, pretend the site owner (a
        # superuser) counts toward the limit.
        limit = self.tier.admin_limit + 1
        raise ValidationError("You already have {limit} admin{s} in your "
                              "site. Upgrade to have access to more.".format(
                                        limit=limit, s=pluralize(limit)))

AuthorFormSet = modelformset_factory(User,
                                     form=AuthorForm,
                                     can_delete=True,
                                     extra=0)

class BulkEditVideoFormSet(_BulkEditVideoFormSet):
    def clean(self):
        tier = SiteTierInfo.objects.get_current().tier
        self.approval_count = 0
        _BulkEditVideoFormSet.clean(self)
        if tier.video_limit is not None:
            video_count = Video.objects.filter(status=Video.ACTIVE,
                                               site=settings.SITE_ID
                                      ).count()
            remaining = tier.video_limit - video_count
            if remaining < 0:
                raise ValidationError('You already have {0} videos over your '
                                      'limit ({1}). Upgrade to approve '
                                      'more.'.format(-1 * remaining,
                                                     tier.video_limit))
            elif self.approval_count > remaining:
                raise ValidationError('You can only approve {0} videos, '
                                      'but tried to approve {1} instead. '
                                      'Upgrade to approve more.'.format(
                                      remaining, self.approval_count))

    def action_approve(self, form):
        if form.instance.status != Video.ACTIVE:
            self.approval_count += 1
        _BulkEditVideoFormSet.action_approve(self, form)

    def action_feature(self, form):
        if form.instance.status != Video.ACTIVE:
            self.approval_count += 1
        _BulkEditVideoFormSet.action_feature(self, form)

VideoFormSet = modelformset_factory(
    Video,
    form=BulkEditVideoForm,
    formset=BulkEditVideoFormSet,
    can_delete=True,
    extra=1)

class DowngradeConfirmationForm(forms.Form):
    """
    This form is designed to get confirmation of a downgrade.

    """
    action = reverse_lazy('localtv_admin_tier_confirm')
    method = "get"

    tier = forms.SlugField(widget=forms.HiddenInput)

    def __init__(self, tier, *args, **kwargs):
        super(DowngradeConfirmationForm, self).__init__(*args, **kwargs)
        self.tier = tier
        self.initial['tier'] = tier.slug

class TierChangeForm(forms.Form):
    """
    This form can be used to change a tier without going through paypal.

    """
    action = reverse_lazy('localtv_admin_tier_change')
    method = "post"

    tier = forms.models.ModelChoiceField(queryset=Tier.objects.all(),
                                         widget=forms.HiddenInput,
                                         to_field_name='slug')
    token = forms.CharField(widget=forms.HiddenInput)

    def __init__(self, *args, **kwargs):
        super(TierChangeForm, self).__init__(*args, **kwargs)
        self.tier_info = SiteTierInfo.objects.get_current()
        self.fields['tier'].queryset = self.tier_info.available_tiers.all()
        if 'tier' in self.initial:
            self.initial['token'] = make_tier_change_token(
                                                         self.initial['tier'])

    def clean_tier(self):
        tier = self.cleaned_data['tier']
        if tier == self.tier_info.tier:
            raise ValidationError("Selected tier is the current tier.")
        return tier

    def clean(self):
        if self.errors:
            # If there are already errors, we can't continue with this part
            # of the validation.
            return self.cleaned_data
        token = self.cleaned_data['token']
        tier = self.cleaned_data['tier']
        if not check_tier_change_token(tier, token):
            raise ValidationError("Invalid tier change token.")
        return self.cleaned_data

    def save(self):
        self.tier_info.tier = self.cleaned_data['tier']
        self.tier_info.tier_changed = datetime.datetime.now()
        self.tier_info.save()
        # Run send_welcome_email to get it out of the way if they're
        # arriving after paying on paypal. It won't do anything if it was
        # already sent, and the overhead is minimal.
        send_welcome_email()

class PayPalCancellationForm(forms.Form):
    action = (SANDBOX_POSTBACK_ENDPOINT if settings.PAYPAL_TEST else
              POSTBACK_ENDPOINT)
    method = "get"

    cmd = forms.CharField(widget=forms.HiddenInput)
    alias = forms.CharField(widget=forms.HiddenInput)

    def __init__(self, tier, *args, **kwargs):
        super(PayPalCancellationForm, self).__init__(*args, **kwargs)
        self.tier = tier
        tier_info = SiteTierInfo.objects.get_current()
        self.initial.update({
            'cmd': '_subscr-find',
            'alias': tier_info.subscription.payer_email
        })

class PayPalSubscriptionForm(PayPalPaymentsForm):
    action = (SANDBOX_POSTBACK_ENDPOINT if settings.PAYPAL_TEST else
              POSTBACK_ENDPOINT)
    method = "post"
    return_url = 'localtv_admin_tier_change'
    cancel_return = 'localtv_admin_tier'
    notify_url = 'paypal-ipn'

    def __init__(self, tier, *args, **kwargs):
        super(PayPalSubscriptionForm, self).__init__(*args, **kwargs)
        self.tier = tier
        site = Site.objects.get_current()
        tier_info = SiteTierInfo.objects.get_current()
        # Downgrades are delayed until the subscription expires, so we should
        # only use the return_url (which would immediately change the tier) if
        # this is an upgrade.
        cancel_return = 'http://{domain}{url}'.format(
                                    domain=site.domain,
                                    url=reverse(self.cancel_return))
        if tier.price < tier_info.tier.price:
            return_url = cancel_return
        else:
            return_params = TierChangeForm(initial={'tier': tier}).initial
            return_url = 'http://{domain}{url}?{query}'.format(
                                    domain=site.domain,
                                    url=reverse(self.return_url),
                                    query=urllib.urlencode(return_params))
        # This initial data says: A subscription for the ``RECEIVER_EMAIL``
        # business with the tier's price, paid every thirty days. The
        # subscription should recur indefinitely and payments should be
        # retried if they fail. cancel_return is where they will be sent if
        # they decide not to pay, or if this is a downgrade (the root of their
        # site). return_url is where they will be sent if they decide to do
        # the change and it's an upgrade (the TierChangeView). notify_url is
        # where ipn notifications will be sent. For more information, see
        # the IPN documentation:
        # https://cms.paypal.com/cms_content/US/en_US/files/developer/PP_WebsitePaymentsStandard_IntegrationGuide.pdf
        self.initial = {
            'cmd': '_xclick-subscriptions',
            'business': RECEIVER_EMAIL,
            # TODO: Should probably reference a url on the current site.
            'image_url': "http://www.mirocommunity.org/images/mc_logo.png",
            'a3': unicode(self.tier.price),
            'p3': '30',
            't3': 'D',
            'src': '1',
            'sra': '1',
            'cancel_return': cancel_return,
            'notify_url': 'http://{domain}{url}'.format(
                                   domain=site.domain,
                                   url=reverse(self.notify_url)),
            'return_url': return_url,
            'item_name': ("Miro Community subscription ({name} on "
                          "{domain})").format(
                                   name=tier.name,
                                   domain=site.domain),
        }
        if not tier_info.had_subscription:
            # If they've never had a subscription before, we add a thirty-day
            # free trial.
            self.initial.update({
                'a1': '0',
                'p1': '30',
                't1': 'D'
            })
        elif tier_info.subscription is not None:
            # If they are currently subscribed, we only allow them to modify
            # that subscription.
            self.initial['modify'] = '2'
<commit_msg>Import PAYPAL_TEST from paypal conf, like the other paypal settings.<commit_after># Miro Community - Easiest way to make a video website
#
# Copyright (C) 2010, 2011, 2012 Participatory Culture Foundation
#
# Miro Community is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# Miro Community is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with Miro Community.  If not, see <http://www.gnu.org/licenses/>.

import datetime
import urllib

from django import forms
from django.conf import settings
from django.contrib.auth.models import User
from django.contrib.sites.models import Site
from django.core.exceptions import ValidationError
from django.core.urlresolvers import reverse, reverse_lazy
from django.forms.models import modelformset_factory
from django.template.defaultfilters import pluralize
from localtv.admin.forms import (EditSettingsForm as _EditSettingsForm,
                                 AuthorForm as _AuthorForm,
                                 BulkEditVideoFormSet as _BulkEditVideoFormSet,
                                 BulkEditVideoForm)
from localtv.models import Video
from paypal.standard.conf import (POSTBACK_ENDPOINT,
                                  SANDBOX_POSTBACK_ENDPOINT,
                                  RECEIVER_EMAIL,
                                  TEST as PAYPAL_TEST)
from paypal.standard.forms import PayPalPaymentsForm

from mirocommunity_saas.models import SiteTierInfo, Tier
from mirocommunity_saas.utils.mail import send_welcome_email
from mirocommunity_saas.utils.tiers import (make_tier_change_token,
                                            check_tier_change_token)

class EditSettingsForm(_EditSettingsForm):
    def __init__(self, *args, **kwargs):
        _EditSettingsForm.__init__(self, *args, **kwargs)
        self.tier = SiteTierInfo.objects.get_current().tier
        if not self.tier.custom_css:
            # Uh-oh: custom CSS is not permitted!
            #
            # To handle only letting certain paid users edit CSS,
            # we do two things.
            #
            # 1. Cosmetically, we set the CSS editing box's CSS class
            # to be 'hidden'. (We have some CSS that makes it not show
            # up.) SB: Not sure why we don't just remove the field...
            css_field = self.fields['css']
            css_field.label += ' (upgrade to enable this form field)'
            css_field.widget.attrs['readonly'] = True
            #
            # 2. In validation, we make sure that changing the CSS is
            # rejected as invalid if the site does not have permission
            # to do that.

    def clean_css(self):
        css = self.cleaned_data.get('css')
        # Does the current tier permit custom CSS? If so, return the data the
        # user submitted.
        if self.tier.custom_css:
            return css # no questions asked

        # We permit the value if it's the same as self.instance has:
        # SB: Leaving this for now, for backwards-compatibility, but again:
        # if css isn't allowed, we should probably not be storing it. This is
        # probably only here because we didn't remove the field on __init__.
        if self.instance.css == css:
            return css

        # Otherwise, reject the change.
        self.data['css'] = self.instance.css
        raise ValidationError(
            "To edit CSS for your site, you have to upgrade.")

class AuthorForm(_AuthorForm):
    def __init__(self, *args, **kwargs):
        _AuthorForm.__init__(self, *args, **kwargs)
        ## Add a note to the 'role' help text indicating how many admins
        ## are permitted with this kind of account.
        self.tier = SiteTierInfo.objects.get_current().tier
        if self.tier.admin_limit is not None:
            # For backwards-compatibility, pretend the site owner (a
            # superuser) counts toward the limit.
            limit = self.tier.admin_limit + 1
            message = ('With a {tier_name}, you may have {limit} '
                       'administrator{s}.').format(tier_name=self.tier.name,
                                                   limit=limit,
                                                   s=pluralize(limit))
            self.fields['role'].help_text = message

    def clean_role(self):
        role = self.cleaned_data['role']

        # Nothing to check if they aren't being set to admin.
        if role != 'admin':
            return role

        # Nothing to check if there is no admin limit.
        if self.tier.admin_limit is None:
            return role

        # Nothing to check if the user is already an admin.
        if 'role' not in self.changed_data:
            return role

        # And finally, we're good if there's room for another admin.
        admin_count = self.site_settings.admins.exclude(is_superuser=True
                                              ).exclude(is_active=False
                                              ).count()
        if (admin_count + 1) <= self.tier.admin_limit:
            return role

        # For backwards-compatibility, pretend the site owner (a
        # superuser) counts toward the limit.
        limit = self.tier.admin_limit + 1
        raise ValidationError("You already have {limit} admin{s} in your "
                              "site. Upgrade to have access to more.".format(
                                        limit=limit, s=pluralize(limit)))

AuthorFormSet = modelformset_factory(User,
                                     form=AuthorForm,
                                     can_delete=True,
                                     extra=0)

class BulkEditVideoFormSet(_BulkEditVideoFormSet):
    def clean(self):
        tier = SiteTierInfo.objects.get_current().tier
        self.approval_count = 0
        _BulkEditVideoFormSet.clean(self)
        if tier.video_limit is not None:
            video_count = Video.objects.filter(status=Video.ACTIVE,
                                               site=settings.SITE_ID
                                      ).count()
            remaining = tier.video_limit - video_count
            if remaining < 0:
                raise ValidationError('You already have {0} videos over your '
                                      'limit ({1}). Upgrade to approve '
                                      'more.'.format(-1 * remaining,
                                                     tier.video_limit))
            elif self.approval_count > remaining:
                raise ValidationError('You can only approve {0} videos, '
                                      'but tried to approve {1} instead. '
                                      'Upgrade to approve more.'.format(
                                      remaining, self.approval_count))

    def action_approve(self, form):
        if form.instance.status != Video.ACTIVE:
            self.approval_count += 1
        _BulkEditVideoFormSet.action_approve(self, form)

    def action_feature(self, form):
        if form.instance.status != Video.ACTIVE:
            self.approval_count += 1
        _BulkEditVideoFormSet.action_feature(self, form)

VideoFormSet = modelformset_factory(
    Video,
    form=BulkEditVideoForm,
    formset=BulkEditVideoFormSet,
    can_delete=True,
    extra=1)

class DowngradeConfirmationForm(forms.Form):
    """
    This form is designed to get confirmation of a downgrade.

    """
    action = reverse_lazy('localtv_admin_tier_confirm')
    method = "get"

    tier = forms.SlugField(widget=forms.HiddenInput)

    def __init__(self, tier, *args, **kwargs):
        super(DowngradeConfirmationForm, self).__init__(*args, **kwargs)
        self.tier = tier
        self.initial['tier'] = tier.slug

class TierChangeForm(forms.Form):
    """
    This form can be used to change a tier without going through paypal.

    """
    action = reverse_lazy('localtv_admin_tier_change')
    method = "post"

    tier = forms.models.ModelChoiceField(queryset=Tier.objects.all(),
                                         widget=forms.HiddenInput,
                                         to_field_name='slug')
    token = forms.CharField(widget=forms.HiddenInput)

    def __init__(self, *args, **kwargs):
        super(TierChangeForm, self).__init__(*args, **kwargs)
        self.tier_info = SiteTierInfo.objects.get_current()
        self.fields['tier'].queryset = self.tier_info.available_tiers.all()
        if 'tier' in self.initial:
            self.initial['token'] = make_tier_change_token(
                                                         self.initial['tier'])

    def clean_tier(self):
        tier = self.cleaned_data['tier']
        if tier == self.tier_info.tier:
            raise ValidationError("Selected tier is the current tier.")
        return tier

    def clean(self):
        if self.errors:
            # If there are already errors, we can't continue with this part
            # of the validation.
            return self.cleaned_data
        token = self.cleaned_data['token']
        tier = self.cleaned_data['tier']
        if not check_tier_change_token(tier, token):
            raise ValidationError("Invalid tier change token.")
        return self.cleaned_data

    def save(self):
        self.tier_info.tier = self.cleaned_data['tier']
        self.tier_info.tier_changed = datetime.datetime.now()
        self.tier_info.save()
        # Run send_welcome_email to get it out of the way if they're
        # arriving after paying on paypal. It won't do anything if it was
        # already sent, and the overhead is minimal.
        send_welcome_email()

class PayPalCancellationForm(forms.Form):
    action = (SANDBOX_POSTBACK_ENDPOINT if PAYPAL_TEST else POSTBACK_ENDPOINT)
    method = "get"

    cmd = forms.CharField(widget=forms.HiddenInput)
    alias = forms.CharField(widget=forms.HiddenInput)

    def __init__(self, tier, *args, **kwargs):
        super(PayPalCancellationForm, self).__init__(*args, **kwargs)
        self.tier = tier
        tier_info = SiteTierInfo.objects.get_current()
        self.initial.update({
            'cmd': '_subscr-find',
            'alias': tier_info.subscription.payer_email
        })

class PayPalSubscriptionForm(PayPalPaymentsForm):
    action = (SANDBOX_POSTBACK_ENDPOINT if PAYPAL_TEST else POSTBACK_ENDPOINT)
    method = "post"
    return_url = 'localtv_admin_tier_change'
    cancel_return = 'localtv_admin_tier'
    notify_url = 'paypal-ipn'

    def __init__(self, tier, *args, **kwargs):
        super(PayPalSubscriptionForm, self).__init__(*args, **kwargs)
        self.tier = tier
        site = Site.objects.get_current()
        tier_info = SiteTierInfo.objects.get_current()
        # Downgrades are delayed until the subscription expires, so we should
        # only use the return_url (which would immediately change the tier) if
        # this is an upgrade.
        cancel_return = 'http://{domain}{url}'.format(
                                    domain=site.domain,
                                    url=reverse(self.cancel_return))
        if tier.price < tier_info.tier.price:
            return_url = cancel_return
        else:
            return_params = TierChangeForm(initial={'tier': tier}).initial
            return_url = 'http://{domain}{url}?{query}'.format(
                                    domain=site.domain,
                                    url=reverse(self.return_url),
                                    query=urllib.urlencode(return_params))
        # This initial data says: A subscription for the ``RECEIVER_EMAIL``
        # business with the tier's price, paid every thirty days. The
        # subscription should recur indefinitely and payments should be
        # retried if they fail. cancel_return is where they will be sent if
        # they decide not to pay, or if this is a downgrade (the root of their
        # site). return_url is where they will be sent if they decide to do
        # the change and it's an upgrade (the TierChangeView). notify_url is
        # where ipn notifications will be sent. For more information, see
        # the IPN documentation:
        # https://cms.paypal.com/cms_content/US/en_US/files/developer/PP_WebsitePaymentsStandard_IntegrationGuide.pdf
        self.initial = {
            'cmd': '_xclick-subscriptions',
            'business': RECEIVER_EMAIL,
            # TODO: Should probably reference a url on the current site.
            'image_url': "http://www.mirocommunity.org/images/mc_logo.png",
            'a3': unicode(self.tier.price),
            'p3': '30',
            't3': 'D',
            'src': '1',
            'sra': '1',
            'cancel_return': cancel_return,
            'notify_url': 'http://{domain}{url}'.format(
                                   domain=site.domain,
                                   url=reverse(self.notify_url)),
            'return_url': return_url,
            'item_name': ("Miro Community subscription ({name} on "
                          "{domain})").format(
                                   name=tier.name,
                                   domain=site.domain),
        }
        if not tier_info.had_subscription:
            # If they've never had a subscription before, we add a thirty-day
            # free trial.
            self.initial.update({
                'a1': '0',
                'p1': '30',
                't1': 'D'
            })
        elif tier_info.subscription is not None:
            # If they are currently subscribed, we only allow them to modify
            # that subscription.
            self.initial['modify'] = '2'
<|endoftext|>


<commit_before>import mailcap
import mimetypes
import os
import os.path
import shlex
import signal
import stat
import subprocess
import sys
import tempfile
import traceback

import urwid
from typing import Optional

from mitmproxy import addons
from mitmproxy import controller
from mitmproxy import exceptions
from mitmproxy import master
from mitmproxy import io
from mitmproxy import log
from mitmproxy.addons import view
from mitmproxy.addons import intercept
import mitmproxy.options
from mitmproxy.tools.console import flowlist
from mitmproxy.tools.console import flowview
from mitmproxy.tools.console import grideditor
from mitmproxy.tools.console import help
from mitmproxy.tools.console import options
from mitmproxy.tools.console import palettepicker
from mitmproxy.tools.console import palettes
from mitmproxy.tools.console import signals
from mitmproxy.tools.console import statusbar
from mitmproxy.tools.console import window
from mitmproxy.utils import strutils

from mitmproxy.net import tcp

EVENTLOG_SIZE = 500

class Options(mitmproxy.options.Options):
    def __init__(
            self,
            *,  # all args are keyword-only.
            eventlog: bool = False,
            focus_follow: bool = False,
            intercept: Optional[str] = None,
            filter: Optional[str] = None,
            palette: Optional[str] = None,
            palette_transparent: bool = False,
            no_mouse: bool = False,
            order: Optional[str] = None,
            order_reversed: bool = False,
            **kwargs
    ):
        self.eventlog = eventlog
        self.focus_follow = focus_follow
        self.intercept = intercept
        self.filter = filter
        self.palette = palette
        self.palette_transparent = palette_transparent
        self.no_mouse = no_mouse
        self.order = order
        self.order_reversed = order_reversed
        super().__init__(**kwargs)

class ConsoleMaster(master.Master):
    palette = []

    def __init__(self, options, server):
        super().__init__(options, server)
        self.view = view.View()
        self.stream_path = None
        # This line is just for type hinting
        self.options = self.options  # type: Options
        self.options.errored.connect(self.options_error)

        self.palette = options.palette
        self.palette_transparent = options.palette_transparent

        self.logbuffer = urwid.SimpleListWalker([])

        self.view_stack = []

        signals.call_in.connect(self.sig_call_in)
        signals.pop_view_state.connect(self.sig_pop_view_state)
        signals.replace_view_state.connect(self.sig_replace_view_state)
        signals.push_view_state.connect(self.sig_push_view_state)
        signals.sig_add_log.connect(self.sig_add_log)
        self.addons.add(*addons.default_addons())
        self.addons.add(intercept.Intercept(), self.view)

    def __setattr__(self, name, value):
        self.__dict__[name] = value
        signals.update_settings.send(self)

    def options_error(self, opts, exc):
        signals.status_message.send(
            message=str(exc),
            expire=1
        )

    def sig_add_log(self, sender, e, level):
        if self.options.verbosity < log.log_tier(level):
            return

        if level in ("error", "warn"):
            signals.status_message.send(
                message = "{}: {}".format(level.title(), e)
            )
            e = urwid.Text((level, str(e)))
        else:
            e = urwid.Text(str(e))
        self.logbuffer.append(e)
        if len(self.logbuffer) > EVENTLOG_SIZE:
            self.logbuffer.pop(0)
        self.logbuffer.set_focus(len(self.logbuffer) - 1)

    def sig_call_in(self, sender, seconds, callback, args=()):
        def cb(*_):
            return callback(*args)
        self.loop.set_alarm_in(seconds, cb)

    def sig_replace_view_state(self, sender):
        """
            A view has been pushed onto the stack, and is intended to replace
            the current view rather tha creating a new stack entry.
        """
        if len(self.view_stack) > 1:
            del self.view_stack[1]

    def sig_pop_view_state(self, sender):
        """
            Pop the top view off the view stack. If no more views will be left
            after this, prompt for exit.
        """
        if len(self.view_stack) > 1:
            self.view_stack.pop()
            self.loop.widget = self.view_stack[-1]
        else:
            signals.status_prompt_onekey.send(
                self,
                prompt = "Quit",
                keys = (
                    ("yes", "y"),
                    ("no", "n"),
                ),
                callback = self.quit,
            )

    def sig_push_view_state(self, sender, window):
        """
            Push a new view onto the view stack.
        """
        self.view_stack.append(window)
        self.loop.widget = window
        self.loop.draw_screen()

    def run_script_once(self, command, f):
        sc = self.addons.get("scriptloader")
        try:
            with self.handlecontext():
                sc.run_once(command, [f])
        except mitmproxy.exceptions.AddonError as e:
            signals.add_log("Script error: %s" % e, "warn")

    def toggle_eventlog(self):
        self.options.eventlog = not self.options.eventlog
        self.view_flowlist()
        signals.replace_view_state.send(self)

    def _readflows(self, path):
        """
        Utitility function that reads a list of flows
        or prints an error to the UI if that fails.
        Returns
            - None, if there was an error.
            - a list of flows, otherwise.
        """
        try:
            return io.read_flows_from_paths(path)
        except exceptions.FlowReadException as e:
            signals.status_message.send(message=str(e))

    def spawn_editor(self, data):
        text = not isinstance(data, bytes)
        fd, name = tempfile.mkstemp('', "mproxy", text=text)
        with open(fd, "w" if text else "wb") as f:
            f.write(data)
        # if no EDITOR is set, assume 'vi'
        c = os.environ.get("EDITOR") or "vi"
        cmd = shlex.split(c)
        cmd.append(name)
        self.ui.stop()
        try:
            subprocess.call(cmd)
        except:
            signals.status_message.send(
                message="Can't start editor: %s" % " ".join(c)
            )
        else:
            with open(name, "r" if text else "rb") as f:
                data = f.read()
        self.ui.start()
        os.unlink(name)
        return data

    def spawn_external_viewer(self, data, contenttype):
        if contenttype:
            contenttype = contenttype.split(";")[0]
            ext = mimetypes.guess_extension(contenttype) or ""
        else:
            ext = ""
        fd, name = tempfile.mkstemp(ext, "mproxy")
        os.write(fd, data)
        os.close(fd)

        # read-only to remind the user that this is a view function
        os.chmod(name, stat.S_IREAD)

        cmd = None
        shell = False

        if contenttype:
            c = mailcap.getcaps()
            cmd, _ = mailcap.findmatch(c, contenttype, filename=name)
            if cmd:
                shell = True
        if not cmd:
            # hm which one should get priority?
            c = os.environ.get("PAGER") or os.environ.get("EDITOR")
            if not c:
                c = "less"
            cmd = shlex.split(c)
            cmd.append(name)
        self.ui.stop()
        try:
            subprocess.call(cmd, shell=shell)
        except:
            signals.status_message.send(
                message="Can't start external viewer: %s" % " ".join(c)
            )
        self.ui.start()
        os.unlink(name)

    def set_palette(self, name):
        self.palette = name
        self.ui.register_palette(
            palettes.palettes[name].palette(self.palette_transparent)
        )
        self.ui.clear()

    def ticker(self, *userdata):
        changed = self.tick(timeout=0)
        if changed:
            self.loop.draw_screen()
            signals.update_settings.send()
        self.loop.set_alarm_in(0.01, self.ticker)

    def run(self):
        self.ui = urwid.raw_display.Screen()
        self.ui.set_terminal_properties(256)
        self.set_palette(self.palette)
        self.loop = urwid.MainLoop(
            urwid.SolidFill("x"),
            screen = self.ui,
            handle_mouse = not self.options.no_mouse,
        )
        self.ab = statusbar.ActionBar()

        if self.options.rfile:
            ret = self.load_flows_path(self.options.rfile)
            if ret and self.view.store_count():
                signals.add_log(
                    "File truncated or corrupted. "
                    "Loaded as many flows as possible.",
                    "error"
                )
            elif ret and not self.view.store_count():
                self.shutdown()
                print("Could not load file: {}".format(ret), file=sys.stderr)
                sys.exit(1)

        self.loop.set_alarm_in(0.01, self.ticker)
        if self.options.http2 and not tcp.HAS_ALPN:  # pragma: no cover
            def http2err(*args, **kwargs):
                signals.status_message.send(
                    message = "HTTP/2 disabled - OpenSSL 1.0.2+ required."
                              " Use --no-http2 to silence this warning.",
                    expire=5
                )
            self.loop.set_alarm_in(0.01, http2err)

        # It's not clear why we need to handle this explicitly - without this,
        # mitmproxy hangs on keyboard interrupt. Remove if we ever figure it
        # out.
        def exit(s, f):
            raise urwid.ExitMainLoop
        signal.signal(signal.SIGINT, exit)

        self.loop.set_alarm_in(
            0.0001,
            lambda *args: self.view_flowlist()
        )

        self.start()
        try:
            self.loop.run()
        except Exception:
            self.loop.stop()
            sys.stdout.flush()
            print(traceback.format_exc(), file=sys.stderr)
            print("mitmproxy has crashed!", file=sys.stderr)
            print("Please lodge a bug report at:", file=sys.stderr)
            print("\thttps://github.com/mitmproxy/mitmproxy", file=sys.stderr)
            print("Shutting down...", file=sys.stderr)
        sys.stderr.flush()
        self.shutdown()

    def view_help(self, helpctx):
        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                help.HelpView(helpctx),
                None,
                statusbar.StatusBar(self, help.footer),
                None
            )
        )

    def view_options(self):
        for i in self.view_stack:
            if isinstance(i["body"], options.Options):
                return
        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                options.Options(self),
                None,
                statusbar.StatusBar(self, options.footer),
                options.help_context,
            )
        )

    def view_palette_picker(self):
        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                palettepicker.PalettePicker(self),
                None,
                statusbar.StatusBar(self, palettepicker.footer),
                palettepicker.help_context,
            )
        )

    def view_grideditor(self, ge):
        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                ge,
                None,
                statusbar.StatusBar(self, grideditor.base.FOOTER),
                ge.make_help()
            )
        )

    def view_flowlist(self):
        if self.ui.started:
            self.ui.clear()

        if self.options.eventlog:
            body = flowlist.BodyPile(self)
        else:
            body = flowlist.FlowListBox(self)

        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                body,
                None,
                statusbar.StatusBar(self, flowlist.footer),
                flowlist.help_context
            )
        )

    def view_flow(self, flow, tab_offset=0):
        self.view.focus.flow = flow
        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                flowview.FlowView(self, self.view, flow, tab_offset),
                flowview.FlowViewHeader(self, flow),
                statusbar.StatusBar(self, flowview.footer),
                flowview.help_context
            )
        )

    def _write_flows(self, path, flows):
        if not path:
            return
        path = os.path.expanduser(path)
        try:
            f = open(path, "wb")
            fw = io.FlowWriter(f)
            for i in flows:
                fw.add(i)
            f.close()
        except IOError as v:
            signals.status_message.send(message=v.strerror)

    def save_one_flow(self, path, flow):
        return self._write_flows(path, [flow])

    def save_flows(self, path):
        return self._write_flows(path, self.view)

    def load_flows_callback(self, path):
        if not path:
            return
        ret = self.load_flows_path(path)
        return ret or "Flows loaded from %s" % path

    def load_flows_path(self, path):
        reterr = None
        try:
            master.Master.load_flows_file(self, path)
        except exceptions.FlowReadException as e:
            reterr = str(e)
        signals.flowlist_change.send(self)
        return reterr

    def quit(self, a):
        if a != "n":
            raise urwid.ExitMainLoop

    def clear_events(self):
        self.logbuffer[:] = []

    # Handlers
    @controller.handler
    def tcp_message(self, f):
        super().tcp_message(f)
        message = f.messages[-1]
        direction = "->" if message.from_client else "<-"
        self.add_log("{client} {direction} tcp {direction} {server}".format(
            client=repr(f.client_conn.address),
            server=repr(f.server_conn.address),
            direction=direction,
        ), "info")
        self.add_log(strutils.bytes_to_escaped_str(message.content), "debug")

    @controller.handler
    def log(self, evt):
        signals.add_log(evt.msg, evt.level)
<commit_msg>console: Minor cleanups<commit_after>import mailcap
import mimetypes
import os
import os.path
import shlex
import stat
import subprocess
import sys
import tempfile
import traceback

import urwid
from typing import Optional

from mitmproxy import addons
from mitmproxy import controller
from mitmproxy import exceptions
from mitmproxy import master
from mitmproxy import io
from mitmproxy import log
from mitmproxy.addons import view
from mitmproxy.addons import intercept
import mitmproxy.options
from mitmproxy.tools.console import flowlist
from mitmproxy.tools.console import flowview
from mitmproxy.tools.console import grideditor
from mitmproxy.tools.console import help
from mitmproxy.tools.console import options
from mitmproxy.tools.console import palettepicker
from mitmproxy.tools.console import palettes
from mitmproxy.tools.console import signals
from mitmproxy.tools.console import statusbar
from mitmproxy.tools.console import window
from mitmproxy.utils import strutils

from mitmproxy.net import tcp

EVENTLOG_SIZE = 500

class Options(mitmproxy.options.Options):
    def __init__(
            self,
            *,  # all args are keyword-only.
            eventlog: bool = False,
            focus_follow: bool = False,
            intercept: Optional[str] = None,
            filter: Optional[str] = None,
            palette: Optional[str] = None,
            palette_transparent: bool = False,
            no_mouse: bool = False,
            order: Optional[str] = None,
            order_reversed: bool = False,
            **kwargs
    ):
        self.eventlog = eventlog
        self.focus_follow = focus_follow
        self.intercept = intercept
        self.filter = filter
        self.palette = palette
        self.palette_transparent = palette_transparent
        self.no_mouse = no_mouse
        self.order = order
        self.order_reversed = order_reversed
        super().__init__(**kwargs)

class ConsoleMaster(master.Master):
    palette = []

    def __init__(self, options, server):
        super().__init__(options, server)
        self.view = view.View()
        self.stream_path = None
        # This line is just for type hinting
        self.options = self.options  # type: Options
        self.options.errored.connect(self.options_error)

        self.palette = options.palette
        self.palette_transparent = options.palette_transparent

        self.logbuffer = urwid.SimpleListWalker([])

        self.view_stack = []

        signals.call_in.connect(self.sig_call_in)
        signals.pop_view_state.connect(self.sig_pop_view_state)
        signals.replace_view_state.connect(self.sig_replace_view_state)
        signals.push_view_state.connect(self.sig_push_view_state)
        signals.sig_add_log.connect(self.sig_add_log)
        self.addons.add(*addons.default_addons())
        self.addons.add(intercept.Intercept(), self.view)

    def __setattr__(self, name, value):
        self.__dict__[name] = value
        signals.update_settings.send(self)

    def options_error(self, opts, exc):
        signals.status_message.send(
            message=str(exc),
            expire=1
        )

    def sig_add_log(self, sender, e, level):
        if self.options.verbosity < log.log_tier(level):
            return

        if level in ("error", "warn"):
            signals.status_message.send(
                message = "{}: {}".format(level.title(), e)
            )
            e = urwid.Text((level, str(e)))
        else:
            e = urwid.Text(str(e))
        self.logbuffer.append(e)
        if len(self.logbuffer) > EVENTLOG_SIZE:
            self.logbuffer.pop(0)
        self.logbuffer.set_focus(len(self.logbuffer) - 1)

    def sig_call_in(self, sender, seconds, callback, args=()):
        def cb(*_):
            return callback(*args)
        self.loop.set_alarm_in(seconds, cb)

    def sig_replace_view_state(self, sender):
        """
            A view has been pushed onto the stack, and is intended to replace
            the current view rather tha creating a new stack entry.
        """
        if len(self.view_stack) > 1:
            del self.view_stack[1]

    def sig_pop_view_state(self, sender):
        """
            Pop the top view off the view stack. If no more views will be left
            after this, prompt for exit.
        """
        if len(self.view_stack) > 1:
            self.view_stack.pop()
            self.loop.widget = self.view_stack[-1]
        else:
            signals.status_prompt_onekey.send(
                self,
                prompt = "Quit",
                keys = (
                    ("yes", "y"),
                    ("no", "n"),
                ),
                callback = self.quit,
            )

    def sig_push_view_state(self, sender, window):
        """
            Push a new view onto the view stack.
        """
        self.view_stack.append(window)
        self.loop.widget = window
        self.loop.draw_screen()

    def run_script_once(self, command, f):
        sc = self.addons.get("scriptloader")
        try:
            with self.handlecontext():
                sc.run_once(command, [f])
        except mitmproxy.exceptions.AddonError as e:
            signals.add_log("Script error: %s" % e, "warn")

    def toggle_eventlog(self):
        self.options.eventlog = not self.options.eventlog
        self.view_flowlist()
        signals.replace_view_state.send(self)

    def _readflows(self, path):
        """
        Utitility function that reads a list of flows
        or prints an error to the UI if that fails.
        Returns
            - None, if there was an error.
            - a list of flows, otherwise.
        """
        try:
            return io.read_flows_from_paths(path)
        except exceptions.FlowReadException as e:
            signals.status_message.send(message=str(e))

    def spawn_editor(self, data):
        text = not isinstance(data, bytes)
        fd, name = tempfile.mkstemp('', "mproxy", text=text)
        with open(fd, "w" if text else "wb") as f:
            f.write(data)
        # if no EDITOR is set, assume 'vi'
        c = os.environ.get("EDITOR") or "vi"
        cmd = shlex.split(c)
        cmd.append(name)
        self.ui.stop()
        try:
            subprocess.call(cmd)
        except:
            signals.status_message.send(
                message="Can't start editor: %s" % " ".join(c)
            )
        else:
            with open(name, "r" if text else "rb") as f:
                data = f.read()
        self.ui.start()
        os.unlink(name)
        return data

    def spawn_external_viewer(self, data, contenttype):
        if contenttype:
            contenttype = contenttype.split(";")[0]
            ext = mimetypes.guess_extension(contenttype) or ""
        else:
            ext = ""
        fd, name = tempfile.mkstemp(ext, "mproxy")
        os.write(fd, data)
        os.close(fd)

        # read-only to remind the user that this is a view function
        os.chmod(name, stat.S_IREAD)

        cmd = None
        shell = False

        if contenttype:
            c = mailcap.getcaps()
            cmd, _ = mailcap.findmatch(c, contenttype, filename=name)
            if cmd:
                shell = True
        if not cmd:
            # hm which one should get priority?
            c = os.environ.get("PAGER") or os.environ.get("EDITOR")
            if not c:
                c = "less"
            cmd = shlex.split(c)
            cmd.append(name)
        self.ui.stop()
        try:
            subprocess.call(cmd, shell=shell)
        except:
            signals.status_message.send(
                message="Can't start external viewer: %s" % " ".join(c)
            )
        self.ui.start()
        os.unlink(name)

    def set_palette(self, name):
        self.palette = name
        self.ui.register_palette(
            palettes.palettes[name].palette(self.palette_transparent)
        )
        self.ui.clear()

    def ticker(self, *userdata):
        changed = self.tick(timeout=0)
        if changed:
            self.loop.draw_screen()
        self.loop.set_alarm_in(0.01, self.ticker)

    def run(self):
        self.ui = urwid.raw_display.Screen()
        self.ui.set_terminal_properties(256)
        self.set_palette(self.palette)
        self.loop = urwid.MainLoop(
            urwid.SolidFill("x"),
            screen = self.ui,
            handle_mouse = not self.options.no_mouse,
        )
        self.ab = statusbar.ActionBar()

        if self.options.rfile:
            ret = self.load_flows_path(self.options.rfile)
            if ret and self.view.store_count():
                signals.add_log(
                    "File truncated or corrupted. "
                    "Loaded as many flows as possible.",
                    "error"
                )
            elif ret and not self.view.store_count():
                self.shutdown()
                print("Could not load file: {}".format(ret), file=sys.stderr)
                sys.exit(1)

        self.loop.set_alarm_in(0.01, self.ticker)
        if self.options.http2 and not tcp.HAS_ALPN:  # pragma: no cover
            def http2err(*args, **kwargs):
                signals.status_message.send(
                    message = "HTTP/2 disabled - OpenSSL 1.0.2+ required."
                              " Use --no-http2 to silence this warning.",
                    expire=5
                )
            self.loop.set_alarm_in(0.01, http2err)

        self.loop.set_alarm_in(
            0.0001,
            lambda *args: self.view_flowlist()
        )

        self.start()
        try:
            self.loop.run()
        except Exception:
            self.loop.stop()
            sys.stdout.flush()
            print(traceback.format_exc(), file=sys.stderr)
            print("mitmproxy has crashed!", file=sys.stderr)
            print("Please lodge a bug report at:", file=sys.stderr)
            print("\thttps://github.com/mitmproxy/mitmproxy", file=sys.stderr)
            print("Shutting down...", file=sys.stderr)
        finally:
            sys.stderr.flush()
            self.shutdown()

    def view_help(self, helpctx):
        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                help.HelpView(helpctx),
                None,
                statusbar.StatusBar(self, help.footer),
                None
            )
        )

    def view_options(self):
        for i in self.view_stack:
            if isinstance(i["body"], options.Options):
                return
        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                options.Options(self),
                None,
                statusbar.StatusBar(self, options.footer),
                options.help_context,
            )
        )

    def view_palette_picker(self):
        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                palettepicker.PalettePicker(self),
                None,
                statusbar.StatusBar(self, palettepicker.footer),
                palettepicker.help_context,
            )
        )

    def view_grideditor(self, ge):
        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                ge,
                None,
                statusbar.StatusBar(self, grideditor.base.FOOTER),
                ge.make_help()
            )
        )

    def view_flowlist(self):
        if self.ui.started:
            self.ui.clear()

        if self.options.eventlog:
            body = flowlist.BodyPile(self)
        else:
            body = flowlist.FlowListBox(self)

        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                body,
                None,
                statusbar.StatusBar(self, flowlist.footer),
                flowlist.help_context
            )
        )

    def view_flow(self, flow, tab_offset=0):
        self.view.focus.flow = flow
        signals.push_view_state.send(
            self,
            window = window.Window(
                self,
                flowview.FlowView(self, self.view, flow, tab_offset),
                flowview.FlowViewHeader(self, flow),
                statusbar.StatusBar(self, flowview.footer),
                flowview.help_context
            )
        )

    def _write_flows(self, path, flows):
        if not path:
            return
        path = os.path.expanduser(path)
        try:
            f = open(path, "wb")
            fw = io.FlowWriter(f)
            for i in flows:
                fw.add(i)
            f.close()
        except IOError as v:
            signals.status_message.send(message=v.strerror)

    def save_one_flow(self, path, flow):
        return self._write_flows(path, [flow])

    def save_flows(self, path):
        return self._write_flows(path, self.view)

    def load_flows_callback(self, path):
        if not path:
            return
        ret = self.load_flows_path(path)
        return ret or "Flows loaded from %s" % path

    def load_flows_path(self, path):
        reterr = None
        try:
            master.Master.load_flows_file(self, path)
        except exceptions.FlowReadException as e:
            reterr = str(e)
        signals.flowlist_change.send(self)
        return reterr

    def quit(self, a):
        if a != "n":
            raise urwid.ExitMainLoop

    def clear_events(self):
        self.logbuffer[:] = []

    # Handlers
    @controller.handler
    def tcp_message(self, f):
        super().tcp_message(f)
        message = f.messages[-1]
        direction = "->" if message.from_client else "<-"
        self.add_log("{client} {direction} tcp {direction} {server}".format(
            client=repr(f.client_conn.address),
            server=repr(f.server_conn.address),
            direction=direction,
        ), "info")
        self.add_log(strutils.bytes_to_escaped_str(message.content), "debug")

    @controller.handler
    def log(self, evt):
        signals.add_log(evt.msg, evt.level)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import

from .. import backend as K
from .. import activations, initializations, regularizers, constraints
from ..engine import Layer, InputSpec

import tensorflow as tf

class Convolution2D_Transpose(Layer):
    """
        Creates a 2D Convolution Transpose layer (sometimes called "Deconvolution").
        Based on code by Xiaomin Wu in "[fchollet/keras] Anyone implemented a Deconvolutional
              layer combined the Keras and Tensorflow? (#2106)"

        Must be placed in the keras/layers/ directory.

        W_shape --- shape of the weights - should be calculated internally
         from conv_2d:  (self.nb_filter, input_dim, self.filter_length, x)

         [0] ... nb_filter
         [1] ... input_dim
         [2] ... filter_length
         [3] ... 1

        b_shape ... shape of the biases - should be calculated internally
         [0] ... nb_filter

        deconv_shape
         this is output_shape of TF conv2d_transpose
         [ batch_size, input_cols, input_rows, input_depth]

        Input should be optimized:
          W_shape --- shape of the weights - should be calculated internally
          b_shape ... shape of the biases - should be calculated internally

          deconv_shape of the form
              [height, width, output_channels, in_channels]
              this is set to be the Output_shape of the layer

          or output shape of the form  [None, kernel_height, kernel_width, output_depth]

          padding: valid|same

          input_dim, input_length ... Keras input parameters

        Also U can set the output_shape(deconv_shape) according to:

        def conv_transpose_out_length(input_size, filter_size, border_mode, stride):
            if input_size is None:
                return None
            if border_mode == 'valid':
                output_size = (input_size - 1) * stride + filter_size
            elif border_mode == 'same':
                output_size = input_size
            return output_size

         TODO for Serialization:
            The config dictionary returned by your layer's get_config probably included the class name as name.
            Just remove that name entry from get_config.

            If your layer requires more to be instantiated than just calling it with config as kwargs: implement a from_config classmethod.
            For examples of complex from_config methods, see the Lambda layer or Merge layer.
    """

    input_ndim = 4

    def __init__(self,
                 init='glorot_uniform', activation='linear', weights=None,
                 padding='valid', strides=[1,1,1,1], deconv_shape=[], W_shape = [],b_shape=[],
                 W_regularizer=None, b_regularizer=None, activity_regularizer=None,
                 W_constraint=None, b_constraint=None, input_dim=None, input_length=None, **kwargs):

        if padding not in {'valid','same'}:
            raise Exception('Invalid border mode for Convolution2D:', padding)
        self.deconv_shape = deconv_shape
        self.init = initializations.get(init)
        self.activation = activations.get(activation)
        assert padding in {'valid', 'same'}, 'border_mode must be in {valid, same}'
        self.padding = padding
        self.strides = strides

        self.W_regularizer = regularizers.get(W_regularizer)
        self.b_regularizer = regularizers.get(b_regularizer)
        self.activity_regularizer = regularizers.get(activity_regularizer)

        self.W_shape = W_shape
        self.b_shape = b_shape

        self.W_constraint = constraints.get(W_constraint)
        self.b_constraint = constraints.get(b_constraint)
        self.constraints = [self.W_constraint, self.b_constraint]

        self.initial_weights = weights
        #self.input = K.placeholder(ndim=4) # old keras 0.3.x

        # Keras 1.0:
        self.input_spec = [InputSpec(ndim=4)]
        self.initial_weights = weights
        self.input_dim = input_dim
        self.input_length = input_length
        if self.input_dim:
            kwargs['input_shape'] = (self.input_length, self.input_dim)

        super(Convolution2D_Transpose, self).__init__(**kwargs)

    #    OLD keras 0.3.x
    #    def build(self):
    #        self.W = self.init(self.W_shape)
    #        self.b = K.zeros(self.b_shape)
    #        self.trainable_weights = [self.W, self.b]
    #        self.regularizers = []
    #
    #        if self.W_regularizer:
    #            self.W_regularizer.set_param(self.W)
    #            self.regularizers.append(self.W_regularizer)
    #
    #        if self.b_regularizer:
    #            self.b_regularizer.set_param(self.b)
    #            self.regularizers.append(self.b_regularizer)
    #
    #        if self.activity_regularizer:
    #            self.activity_regularizer.set_layer(self)
    #            self.regularizers.append(self.activity_regularizer)
    #
    #        if self.initial_weights is not None:
    #            self.set_weights(self.initial_weights)
    #            del self.initial_weights

    # NEW keras 1.0
    def build(self, input_shape):
        input_dim = input_shape[2]
        #self.W_shape = (self.nb_filter, input_dim, self.filter_length, 1) # goven from outside
        self.W = self.init(self.W_shape, name='{}_W'.format(self.name))
        self.b = K.zeros((self.b_shape), name='{}_b'.format(self.name))
        self.trainable_weights = [self.W, self.b]
        self.regularizers = []

        if self.W_regularizer:
            self.W_regularizer.set_param(self.W)
            self.regularizers.append(self.W_regularizer)

        if self.b_regularizer:
            self.b_regularizer.set_param(self.b)
            self.regularizers.append(self.b_regularizer)

        if self.activity_regularizer:
            self.activity_regularizer.set_layer(self)
            self.regularizers.append(self.activity_regularizer)

        self.constraints = {}
        if self.W_constraint:
            self.constraints[self.W] = self.W_constraint
        if self.b_constraint:
            self.constraints[self.b] = self.b_constraint

        if self.initial_weights is not None:
            self.set_weights(self.initial_weights)
            del self.initial_weights

    @property
    def get_output_shape(self, input_shape):
        return self.deconv_shape

    #def get_output_OLD(self, train=False):
    #    X = self.get_input(train)
    #    X = K.permute_dimensions(X, (0, 2, 3, 1))
    #    conv_out = tf.nn.conv2d_transpose(X, self.W, strides=self.strides,
    #                                      padding=self.padding,
    #                                      output_shape=self.deconv_shape)
    #    output = conv_out + K.reshape(self.b, (1, 1, 1, self.W_shape[2]))
    #    return K.permute_dimensions(output, (0, 3, 1, 2))

    def call(self, X,  mask=None):
        #X = self.get_input(train)
        X = K.permute_dimensions(X, (0, 2, 3, 1))
        conv_out = tf.nn.conv2d_transpose(X, self.W, strides=self.strides,
                                          padding=self.padding.upper(),
                                          output_shape=self.deconv_shape)

        output = conv_out + K.reshape(self.b, (1, 1, 1, self.W_shape[2]))
        return K.permute_dimensions(output, (0, 3, 1, 2))

    def get_config(self):
        config = {
                  #'name': self.__class__.__name__,
                  'init': self.init.__name__,
                  'activation': self.activation.__name__,
                  'padding': self.padding,
                  'strides': self.strides,
                  'W_regularizer': self.W_regularizer.get_config() if self.W_regularizer else None,
                  'b_regularizer': self.b_regularizer.get_config() if self.b_regularizer else None,
                  'activity_regularizer': self.activity_regularizer.get_config() if self.activity_regularizer else None,
                  'W_constraint': self.W_constraint.get_config() if self.W_constraint else None,
                  'b_constraint': self.b_constraint.get_config() if self.b_constraint else None,
                  'W_shape': self.W_shape,
                  'b_shape': self.b_shape,
                  'deconv_shape': self.deconv_shape }
        base_config = super(Convolution2D_Transpose, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def get_output_shape_for(self, input_shape):
        return (self.deconv_shape[0],self.deconv_shape[1],self.deconv_shape[2],self.deconv_shape[3])
        #return (None, 28, 28 ,1)

<commit_msg>Comments are correct, esp. contents of arrays<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import

from .. import backend as K
from .. import activations, initializations, regularizers, constraints
from ..engine import Layer, InputSpec

import tensorflow as tf

class Convolution2D_Transpose(Layer):
    """
        Creates a 2D Convolution Transpose layer (sometimes called "Deconvolution").
        Based on code by Xiaomin Wu in "[fchollet/keras] Anyone implemented a Deconvolutional
              layer combined the Keras and Tensorflow? (#2106)"

        Must be placed in the keras/layers/ directory.

        W_shape --- shape of the weights - should be calculated internally
         [filter_dim_y, filter_dim_x, self.nb_filter(=number of channels in output), number_of_channels_in_input]

        b_shape ... shape of the biases - should be calculated internally
         [0] ... nb_filter

        strides
         Strides of the filters
         [stride_in_batch_size(must be 1), stride_y, stride_x, stride_in_depth (must be 1)]

        deconv_shape
         this is output_shape of TF conv2d_transpose
         deconv_shape = [batch_size, output_size_y, output_size_x, number_of_filters]

        padding: valid|same (small caps)

        input_dim, input_length ... Keras input parameters

        Also U can set the output_shape(deconv_shape) according to:

        def conv_transpose_out_length(input_size, filter_size, border_mode, stride):
            if input_size is None:
                return None
            if border_mode == 'valid':
                output_size = (input_size - 1) * stride + filter_size
            elif border_mode == 'same':
                output_size = input_size
            return output_size

    """

    input_ndim = 4

    def __init__(self,
                 init='glorot_uniform', activation='linear', weights=None,
                 padding='valid', strides=[1,1,1,1], deconv_shape=[], W_shape = [],b_shape=[],
                 W_regularizer=None, b_regularizer=None, activity_regularizer=None,
                 W_constraint=None, b_constraint=None, input_dim=None, input_length=None, **kwargs):

        if padding not in {'valid','same'}:
            raise Exception('Invalid border mode for Convolution2D:', padding)
        self.deconv_shape = deconv_shape
        self.init = initializations.get(init)
        self.activation = activations.get(activation)
        assert padding in {'valid', 'same'}, 'border_mode must be in {valid, same}'
        self.padding = padding
        self.strides = strides

        self.W_regularizer = regularizers.get(W_regularizer)
        self.b_regularizer = regularizers.get(b_regularizer)
        self.activity_regularizer = regularizers.get(activity_regularizer)

        self.W_shape = W_shape
        self.b_shape = b_shape

        self.W_constraint = constraints.get(W_constraint)
        self.b_constraint = constraints.get(b_constraint)
        self.constraints = [self.W_constraint, self.b_constraint]

        self.initial_weights = weights
        #self.input = K.placeholder(ndim=4) # old keras 0.3.x

        # Keras 1.0:
        self.input_spec = [InputSpec(ndim=4)]
        self.initial_weights = weights
        self.input_dim = input_dim
        self.input_length = input_length
        if self.input_dim:
            kwargs['input_shape'] = (self.input_length, self.input_dim)

        super(Convolution2D_Transpose, self).__init__(**kwargs)

    def build(self, input_shape):
        input_dim = input_shape[2]
        #self.W_shape = (self.nb_filter, input_dim, self.filter_length, 1) # goven from outside
        self.W = self.init(self.W_shape, name='{}_W'.format(self.name))
        self.b = K.zeros((self.b_shape), name='{}_b'.format(self.name))
        self.trainable_weights = [self.W, self.b]
        self.regularizers = []

        if self.W_regularizer:
            self.W_regularizer.set_param(self.W)
            self.regularizers.append(self.W_regularizer)

        if self.b_regularizer:
            self.b_regularizer.set_param(self.b)
            self.regularizers.append(self.b_regularizer)

        if self.activity_regularizer:
            self.activity_regularizer.set_layer(self)
            self.regularizers.append(self.activity_regularizer)

        self.constraints = {}
        if self.W_constraint:
            self.constraints[self.W] = self.W_constraint
        if self.b_constraint:
            self.constraints[self.b] = self.b_constraint

        if self.initial_weights is not None:
            self.set_weights(self.initial_weights)
            del self.initial_weights

    @property
    def get_output_shape(self, input_shape):
        return self.deconv_shape

    def call(self, X,  mask=None):
        #X = self.get_input(train)
        X = K.permute_dimensions(X, (0, 2, 3, 1))
        conv_out = tf.nn.conv2d_transpose(X, self.W, strides=self.strides,
                                          padding=self.padding.upper(),
                                          output_shape=self.deconv_shape)

        output = conv_out + K.reshape(self.b, (1, 1, 1, self.W_shape[2]))
        return K.permute_dimensions(output, (0, 3, 1, 2))

    def get_config(self):
        config = {
                  'init': self.init.__name__,
                  'activation': self.activation.__name__,
                  'padding': self.padding,
                  'strides': self.strides,
                  'W_regularizer': self.W_regularizer.get_config() if self.W_regularizer else None,
                  'b_regularizer': self.b_regularizer.get_config() if self.b_regularizer else None,
                  'activity_regularizer': self.activity_regularizer.get_config() if self.activity_regularizer else None,
                  'W_constraint': self.W_constraint.get_config() if self.W_constraint else None,
                  'b_constraint': self.b_constraint.get_config() if self.b_constraint else None,
                  'W_shape': self.W_shape,
                  'b_shape': self.b_shape,
                  'deconv_shape': self.deconv_shape }
        base_config = super(Convolution2D_Transpose, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def get_output_shape_for(self, input_shape):
        return (self.deconv_shape[0],self.deconv_shape[1],self.deconv_shape[2],self.deconv_shape[3])

<|endoftext|>