<commit_before>################################################################################
# Copyright (c) 2017-2018, National Research Foundation (Square Kilometre Array)
#
# Licensed under the BSD 3-Clause License (the "License"); you may not use
# this file except in compliance with the License. You may obtain a copy
# of the License at
#
#   https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################

"""Various sources of correlator data and metadata."""

import urlparse
import os
import logging

import katsdptelstate
import numpy as np
import dask.array as da

from .sensordata import TelstateSensorData
from .chunkstore_s3 import S3ChunkStore
from .chunkstore_npy import NpyFileChunkStore

logger = logging.getLogger(__name__)

class DataSourceNotFound(Exception):
    """File associated with DataSource not found or server not responding."""

class AttrsSensors(object):
    """Metadata in the form of attributes and sensors.

    Parameters
    ----------
    attrs : mapping from string to object
        Metadata attributes
    sensors : mapping from string to :class:`SensorData` objects
        Metadata sensor cache mapping sensor names to raw sensor data
    name : string, optional
        Identifier that describes the origin of the metadata (backend-specific)

    """
    def __init__(self, attrs, sensors, name='custom'):
        self.attrs = attrs
        self.sensors = sensors
        self.name = name

class VisFlagsWeights(object):
    """Correlator data in the form of visibilities, flags and weights.

    Parameters
    ----------
    vis : array-like of complex64, shape (*T*, *F*, *B*)
        Complex visibility data as a function of time, frequency and baseline
    flags : array-like of uint8, shape (*T*, *F*, *B*)
        Flags as a function of time, frequency and baseline
    weights : array-like of float32, shape (*T*, *F*, *B*)
        Visibility weights as a function of time, frequency and baseline
    name : string, optional
        Identifier that describes the origin of the data (backend-specific)

    """
    def __init__(self, vis, flags, weights, name='custom'):
        if not (vis.shape == flags.shape == weights.shape):
            raise ValueError("Shapes of vis %s, flags %s and weights %s differ"
                             % (vis.shape, flags.shape, weights.shape))
        self.vis = vis
        self.flags = flags
        self.weights = weights
        self.name = name

    @property
    def shape(self):
        return self.vis.shape

def _has_chunk_to_flags(has_chunk, block_id, full_chunks):
    """Turn a has_chunk bool into chunk of flags with correct data_lost bit."""
    shape = tuple(chk[idx] for chk, idx in zip(full_chunks, block_id))
    return np.full(shape, 0 if has_chunk else 8, dtype=np.uint8)

def _multi_or_3d(*args):
    """Do bitwise 'or' of two or more 3-D arrays (without modifying them)."""
    args = np.atleast_3d(*args)
    out = args[0] | args[1]
    for arg in args[2:]:
        out |= arg
    return out

class ChunkStoreVisFlagsWeights(VisFlagsWeights):
    """Correlator data stored in a chunk store.

    Parameters
    ----------
    store : :class:`ChunkStore` object
        Chunk store
    base_name : string
        Name of dataset in store, as array name prefix (akin to a filename)
    chunk_info : dict mapping array name to info dict
        Dict specifying dtype, shape and chunks per array
    """
    def __init__(self, store, base_name, chunk_info):
        self.store = store
        darray = {}
        extra_flags = []
        for array, info in chunk_info.iteritems():
            array_name = store.join(base_name, array)
            chunk_args = (array_name, info['chunks'], info['dtype'])
            darray[array] = store.get_dask_array(*chunk_args)
            # Find all missing chunks in array and convert to 'data_lost' flags
            has_array = store.has_dask_array(*chunk_args)
            chunks_lost = da.map_blocks(_has_chunk_to_flags, has_array,
                                        token='missing-chunks-' + array_name,
                                        chunks=info['chunks'], dtype=np.uint8,
                                        full_chunks=info['chunks'])
            extra_flags.append(chunks_lost)
            extra_flags.append('ijk'[:chunks_lost.ndim])
        vis = darray['correlator_data']
        # Combine original L0 flags with extras (missing chunks per array)
        flags = da.atop(_multi_or_3d, 'ijk', darray['flags'], 'ijk',
                        *extra_flags, token=store.join(base_name, 'flags_raw'),
                        dtype=np.uint8)
        # Combine low-resolution weights and high-resolution weights_channel
        weights = darray['weights'] * darray['weights_channel'][..., np.newaxis]
        VisFlagsWeights.__init__(self, vis, flags, weights, base_name)

class DataSource(object):
    """A generic data source presenting both correlator data and metadata.

    Parameters
    ----------
    metadata : :class:`AttrsSensors` object
        Metadata attributes and sensors
    timestamps : array-like of float, length *T*
        Timestamps at centroids of visibilities in UTC seconds since Unix epoch
    data : :class:`VisFlagsWeights` object, optional
        Correlator data (visibilities, flags and weights)

    """
    def __init__(self, metadata, timestamps, data=None):
        self.metadata = metadata
        self.timestamps = timestamps
        self.data = data

    @property
    def name(self):
        name = self.metadata.name
        if self.data and self.data.name != name:
            name += ' | ' + self.data.name
        return name

def view_capture_stream(telstate, capture_block_id=None, stream_name=None,
                        **kwargs):
    """Create telstate view based on capture block ID and stream name.

    This figures out the appropriate capture block ID and L0 stream name from
    a capture-stream specific telstate, or uses the provided ones. It then
    constructs a view on `telstate` with at least the prefixes

      - <capture_block_id>_<stream_name>
      - <capture_block_id>
      - <stream_name>

    Parameters
    ----------
    telstate : :class:`katsdptelstate.TelescopeState` object
        Original telescope state
    capture_block_id : string, optional
        Specify capture block ID explicitly (detected otherwise)
    stream_name : string, optional
        Specify L0 stream name explicitly (detected otherwise)
    kwargs : dict, optional
        Extra keyword arguments, typically meant for other methods and ignored

    Returns
    -------
    telstate : :class:`katsdptelstate.TelescopeState` object
        Telstate with a view that incorporates capture block, stream and combo

    Raises
    ------
    ValueError
        If no capture block or L0 stream could be detected (with no override)
    """
    # Detect the capture block
    if not capture_block_id:
        try:
            capture_block_id = str(telstate['capture_block_id'])
        except KeyError:
            raise ValueError('No capture block ID found in telstate - '
                             'please specify it manually')
    # Detect the captured stream
    if not stream_name:
        try:
            stream_name = str(telstate['stream_name'])
        except KeyError:
            raise ValueError('No captured stream found in telstate - '
                             'please specify the stream manually')
    # Check the stream type
    telstate = telstate.view(stream_name)
    stream_type = telstate.get('stream_type', 'unknown')
    expected_type = 'sdp.vis'
    if stream_type != expected_type:
        raise ValueError("Found stream {!r} but it has the wrong type {!r},"
                         " expected {!r}".format(stream_name, stream_type,
                                                 expected_type))
    logger.info('Using capture block %r and stream %r',
                capture_block_id, stream_name)
    telstate = telstate.view(capture_block_id)
    capture_stream = telstate.SEPARATOR.join((capture_block_id, stream_name))
    telstate = telstate.view(capture_stream)
    return telstate

def _shorten_key(telstate, key):
    """Shorten telstate key by subtracting the first prefix that fits.

    Parameters
    ----------
    telstate : :class:`katsdptelstate.TelescopeState` object
        Telescope state
    key : string
        Telescope state key

    Returns
    -------
    short_key : string
        Suffix of `key` after subtracting first matching prefix, or empty
        string if `key` does not start with any of the prefixes (or exactly
        matches a prefix, which is also considered pathological)

    """
    for prefix in telstate.prefixes:
        if key.startswith(prefix):
            return key[len(prefix):]
    return ''

def _infer_chunk_store(url_parts, telstate, npy_store_path=None,
                       s3_endpoint_url=None, **kwargs):
    """Construct chunk store automatically from dataset URL and telstate.

    Parameters
    ----------
    url_parts : :class:`urlparse.ParseResult` object
        Parsed dataset URL
    telstate : :class:`katsdptelstate.TelescopeState` object
        Telescope state
    npy_store_path : string, optional
        Top-level directory of NpyFileChunkStore (overrides the default)
    s3_endpoint_url : string, optional
        Endpoint of S3 service, e.g. 'http://127.0.0.1:9000' (overrides default)
    kwargs : dict, optional
        Extra keyword arguments, typically meant for other methods and ignored

    Returns
    -------
    store : :class:`katdal.ChunkStore` object
        Chunk store for visibility data

    Raises
    ------
    KeyError
        If telstate lacks critical keys
    :exc:`katdal.chunkstore.StoreUnavailable`
        If the chunk store could not be constructed
    """
    # Use overrides if provided, regardless of URL and telstate (NPY first)
    if npy_store_path:
        return NpyFileChunkStore(npy_store_path)
    if s3_endpoint_url:
        return S3ChunkStore.from_url(s3_endpoint_url, **kwargs)
    # NPY chunk store is an option if the dataset is an RDB file
    if url_parts.scheme == 'file':
        # Look for adjacent data directory (presumably containing NPY files)
        rdb_path = os.path.abspath(url_parts.path)
        store_path = os.path.dirname(os.path.dirname(rdb_path))
        data_path = os.path.join(store_path, telstate['chunk_name'])
        if os.path.isdir(data_path):
            return NpyFileChunkStore(store_path)
    return S3ChunkStore.from_url(telstate['s3_endpoint_url'], **kwargs)

class TelstateDataSource(DataSource):
    """A data source based on :class:`katsdptelstate.TelescopeState`.

    It is assumed that the provided `telstate` already has the appropriate
    views to find observation, stream and chunk store information. It typically
    needs the following prefixes:

      - <capture block ID>_<L0 stream>
      - <capture block ID>
      - <L0 stream>

    Parameters
    ----------
    telstate : :class:`katsdptelstate.TelescopeState` object
        Telescope state with appropriate views
    chunk_store : :class:`katdal.ChunkStore` object, optional
        Chunk store for visibility data (the default is no data - metadata only)
    timestamps : array of float, optional
        Visibility timestamps, overriding (or fixing) the ones found in telstate
    source_name : string, optional
        Name of telstate source (used for metadata name)

    Raises
    ------
    KeyError
        If telstate lacks critical keys
    """
    def __init__(self, telstate, chunk_store=None, timestamps=None,
                 source_name='telstate'):
        self.telstate = telstate
        # Collect sensors
        sensors = {}
        for key in telstate.keys():
            if not telstate.is_immutable(key):
                sensor_name = _shorten_key(telstate, key)
                if sensor_name:
                    sensors[sensor_name] = TelstateSensorData(telstate, key)
        metadata = AttrsSensors(telstate, sensors, name=source_name)
        if timestamps is None:
            # Synthesise timestamps from the relevant telstate bits
            t0 = telstate['sync_time'] + telstate['first_timestamp']
            int_time = telstate['int_time']
            chunk_info = telstate['chunk_info']
            n_dumps = chunk_info['correlator_data']['shape'][0]
            timestamps = t0 + np.arange(n_dumps) * int_time
        if chunk_store is None:
            data = None
        else:
            data = ChunkStoreVisFlagsWeights(
                chunk_store, telstate['chunk_name'], telstate['chunk_info'])
        # Metadata and timestamps with or without data
        DataSource.__init__(self, metadata, timestamps, data)

    @classmethod
    def from_url(cls, url, chunk_store='auto', **kwargs):
        """Construct TelstateDataSource from URL (RDB file / REDIS server).

        Parameters
        ----------
        url : string
            URL serving as entry point to dataset (typically RDB file or REDIS)
        chunk_store : :class:`katdal.ChunkStore` object, optional
            Chunk store for visibility data (obtained automatically by default,
            or set to None for metadata-only dataset)
        kwargs : dict, optional
            Extra keyword arguments passed to telstate view and chunk store init
        """
        url_parts = urlparse.urlparse(url, scheme='file')
        # Merge key-value pairs from URL query with keyword arguments
        # of function (the latter takes precedence)
        url_kwargs = dict(urlparse.parse_qsl(url_parts.query))
        url_kwargs.update(kwargs)
        kwargs = url_kwargs
        # Extract Redis database number if provided
        db = int(kwargs.pop('db', '0'))
        if url_parts.scheme == 'file':
            # RDB dump file
            telstate = katsdptelstate.TelescopeState()
            try:
                telstate.load_from_file(url_parts.path)
            except OSError as err:
                raise DataSourceNotFound(str(err))
        elif url_parts.scheme == 'redis':
            # Redis server
            try:
                telstate = katsdptelstate.TelescopeState(url_parts.netloc, db)
            except katsdptelstate.ConnectionError as e:
                raise DataSourceNotFound(str(e))
        telstate = view_capture_stream(telstate, **kwargs)
        if chunk_store == 'auto':
            chunk_store = _infer_chunk_store(url_parts, telstate, **kwargs)
        return cls(telstate, chunk_store, source_name=url_parts.geturl())

def open_data_source(url, **kwargs):
    """Construct the data source described by the given URL."""
    try:
        return TelstateDataSource.from_url(url, **kwargs)
    except DataSourceNotFound as err:
        # Amend the error message for the case of an IP address without scheme
        url_parts = urlparse.urlparse(url, scheme='file')
        if url_parts.scheme == 'file' and not os.path.isfile(url_parts.path):
            raise DataSourceNotFound(
                '{} (add a URL scheme if {!r} is not meant to be a file)'
                .format(err, url_parts.path))
<commit_msg>Incorporate L1 flags into L0 product<commit_after>################################################################################
# Copyright (c) 2017-2018, National Research Foundation (Square Kilometre Array)
#
# Licensed under the BSD 3-Clause License (the "License"); you may not use
# this file except in compliance with the License. You may obtain a copy
# of the License at
#
#   https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################

"""Various sources of correlator data and metadata."""

import urlparse
import os
import logging

import katsdptelstate
import numpy as np
import dask.array as da

from .sensordata import TelstateSensorData
from .chunkstore_s3 import S3ChunkStore
from .chunkstore_npy import NpyFileChunkStore

logger = logging.getLogger(__name__)

class DataSourceNotFound(Exception):
    """File associated with DataSource not found or server not responding."""

class AttrsSensors(object):
    """Metadata in the form of attributes and sensors.

    Parameters
    ----------
    attrs : mapping from string to object
        Metadata attributes
    sensors : mapping from string to :class:`SensorData` objects
        Metadata sensor cache mapping sensor names to raw sensor data
    name : string, optional
        Identifier that describes the origin of the metadata (backend-specific)

    """
    def __init__(self, attrs, sensors, name='custom'):
        self.attrs = attrs
        self.sensors = sensors
        self.name = name

class VisFlagsWeights(object):
    """Correlator data in the form of visibilities, flags and weights.

    Parameters
    ----------
    vis : array-like of complex64, shape (*T*, *F*, *B*)
        Complex visibility data as a function of time, frequency and baseline
    flags : array-like of uint8, shape (*T*, *F*, *B*)
        Flags as a function of time, frequency and baseline
    weights : array-like of float32, shape (*T*, *F*, *B*)
        Visibility weights as a function of time, frequency and baseline
    name : string, optional
        Identifier that describes the origin of the data (backend-specific)

    """
    def __init__(self, vis, flags, weights, name='custom'):
        if not (vis.shape == flags.shape == weights.shape):
            raise ValueError("Shapes of vis %s, flags %s and weights %s differ"
                             % (vis.shape, flags.shape, weights.shape))
        self.vis = vis
        self.flags = flags
        self.weights = weights
        self.name = name

    @property
    def shape(self):
        return self.vis.shape

def _has_chunk_to_flags(has_chunk, block_id, full_chunks):
    """Turn a has_chunk bool into chunk of flags with correct data_lost bit."""
    shape = tuple(chk[idx] for chk, idx in zip(full_chunks, block_id))
    return np.full(shape, 0 if has_chunk else 8, dtype=np.uint8)

def _multi_or_3d(*args):
    """Do bitwise 'or' of two or more 3-D arrays (without modifying them)."""
    args = np.atleast_3d(*args)
    out = args[0] | args[1]
    for arg in args[2:]:
        out |= arg
    return out

class ChunkStoreVisFlagsWeights(VisFlagsWeights):
    """Correlator data stored in a chunk store.

    Parameters
    ----------
    store : :class:`ChunkStore` object
        Chunk store
    chunk_info : dict mapping array name to info dict
        Dict specifying prefix, dtype, shape and chunks per array
    """
    def __init__(self, store, chunk_info):
        self.store = store
        darray = {}
        extra_flags = []
        for array, info in chunk_info.items():
            array_name = store.join(info['prefix'], array)
            chunk_args = (array_name, info['chunks'], info['dtype'])
            darray[array] = store.get_dask_array(*chunk_args)
            # Find all missing chunks in array and convert to 'data_lost' flags
            has_array = store.has_dask_array(*chunk_args)
            chunks_lost = da.map_blocks(_has_chunk_to_flags, has_array,
                                        token='missing-chunks-' + array_name,
                                        chunks=info['chunks'], dtype=np.uint8,
                                        full_chunks=info['chunks'])
            extra_flags.append(chunks_lost)
            extra_flags.append('ijk'[:chunks_lost.ndim])
        vis = darray['correlator_data']
        base_name = chunk_info['correlator_data']['prefix']
        flags_raw_name = store.join(chunk_info['flags']['prefix'], 'flags_raw')
        # Combine original L0 flags with extras (missing chunks per array)
        flags = da.atop(_multi_or_3d, 'ijk', darray['flags'], 'ijk',
                        *extra_flags, token=flags_raw_name, dtype=np.uint8)
        # Combine low-resolution weights and high-resolution weights_channel
        weights = darray['weights'] * darray['weights_channel'][..., np.newaxis]
        VisFlagsWeights.__init__(self, vis, flags, weights, base_name)

class DataSource(object):
    """A generic data source presenting both correlator data and metadata.

    Parameters
    ----------
    metadata : :class:`AttrsSensors` object
        Metadata attributes and sensors
    timestamps : array-like of float, length *T*
        Timestamps at centroids of visibilities in UTC seconds since Unix epoch
    data : :class:`VisFlagsWeights` object, optional
        Correlator data (visibilities, flags and weights)

    """
    def __init__(self, metadata, timestamps, data=None):
        self.metadata = metadata
        self.timestamps = timestamps
        self.data = data

    @property
    def name(self):
        name = self.metadata.name
        if self.data and self.data.name != name:
            name += ' | ' + self.data.name
        return name

def view_capture_stream(telstate, capture_block_id=None, stream_name=None,
                        **kwargs):
    """Create telstate view based on capture block ID and stream name.

    This figures out the appropriate capture block ID and L0 stream name from
    a capture-stream specific telstate, or uses the provided ones. It then
    constructs a view on `telstate` with at least the prefixes

      - <capture_block_id>_<stream_name>
      - <capture_block_id>
      - <stream_name>

    Parameters
    ----------
    telstate : :class:`katsdptelstate.TelescopeState` object
        Original telescope state
    capture_block_id : string, optional
        Specify capture block ID explicitly (detected otherwise)
    stream_name : string, optional
        Specify L0 stream name explicitly (detected otherwise)
    kwargs : dict, optional
        Extra keyword arguments, typically meant for other methods and ignored

    Returns
    -------
    telstate : :class:`katsdptelstate.TelescopeState` object
        Telstate with a view that incorporates capture block, stream and combo

    Raises
    ------
    ValueError
        If no capture block or L0 stream could be detected (with no override)
    """
    # Detect the capture block
    if not capture_block_id:
        try:
            capture_block_id = str(telstate['capture_block_id'])
        except KeyError:
            raise ValueError('No capture block ID found in telstate - '
                             'please specify it manually')
    # Detect the captured stream
    if not stream_name:
        try:
            stream_name = str(telstate['stream_name'])
        except KeyError:
            raise ValueError('No captured stream found in telstate - '
                             'please specify the stream manually')
    # Check the stream type
    telstate = telstate.view(stream_name)
    stream_type = telstate.get('stream_type', 'unknown')
    expected_type = 'sdp.vis'
    if stream_type != expected_type:
        raise ValueError("Found stream {!r} but it has the wrong type {!r},"
                         " expected {!r}".format(stream_name, stream_type,
                                                 expected_type))
    logger.info('Using capture block %r and stream %r',
                capture_block_id, stream_name)
    telstate = telstate.view(capture_block_id)
    capture_stream = telstate.SEPARATOR.join((capture_block_id, stream_name))
    telstate = telstate.view(capture_stream)
    return telstate

def _shorten_key(telstate, key):
    """Shorten telstate key by subtracting the first prefix that fits.

    Parameters
    ----------
    telstate : :class:`katsdptelstate.TelescopeState` object
        Telescope state
    key : string
        Telescope state key

    Returns
    -------
    short_key : string
        Suffix of `key` after subtracting first matching prefix, or empty
        string if `key` does not start with any of the prefixes (or exactly
        matches a prefix, which is also considered pathological)

    """
    for prefix in telstate.prefixes:
        if key.startswith(prefix):
            return key[len(prefix):]
    return ''

def _infer_chunk_store(url_parts, telstate, npy_store_path=None,
                       s3_endpoint_url=None, **kwargs):
    """Construct chunk store automatically from dataset URL and telstate.

    Parameters
    ----------
    url_parts : :class:`urlparse.ParseResult` object
        Parsed dataset URL
    telstate : :class:`katsdptelstate.TelescopeState` object
        Telescope state
    npy_store_path : string, optional
        Top-level directory of NpyFileChunkStore (overrides the default)
    s3_endpoint_url : string, optional
        Endpoint of S3 service, e.g. 'http://127.0.0.1:9000' (overrides default)
    kwargs : dict, optional
        Extra keyword arguments, typically meant for other methods and ignored

    Returns
    -------
    store : :class:`katdal.ChunkStore` object
        Chunk store for visibility data

    Raises
    ------
    KeyError
        If telstate lacks critical keys
    :exc:`katdal.chunkstore.StoreUnavailable`
        If the chunk store could not be constructed
    """
    # Use overrides if provided, regardless of URL and telstate (NPY first)
    if npy_store_path:
        return NpyFileChunkStore(npy_store_path)
    if s3_endpoint_url:
        return S3ChunkStore.from_url(s3_endpoint_url, **kwargs)
    # NPY chunk store is an option if the dataset is an RDB file
    if url_parts.scheme == 'file':
        # Look for adjacent data directory (presumably containing NPY files)
        rdb_path = os.path.abspath(url_parts.path)
        store_path = os.path.dirname(os.path.dirname(rdb_path))
        data_path = os.path.join(store_path, telstate['chunk_name'])
        if os.path.isdir(data_path):
            return NpyFileChunkStore(store_path)
    return S3ChunkStore.from_url(telstate['s3_endpoint_url'], **kwargs)

class TelstateDataSource(DataSource):
    """A data source based on :class:`katsdptelstate.TelescopeState`.

    It is assumed that the provided `telstate` already has the appropriate
    views to find observation, stream and chunk store information. It typically
    needs the following prefixes:

      - <capture block ID>_<L0 stream>
      - <capture block ID>
      - <L0 stream>

    Parameters
    ----------
    telstate : :class:`katsdptelstate.TelescopeState` object
        Telescope state with appropriate views
    chunk_store : :class:`katdal.ChunkStore` object, optional
        Chunk store for visibility data (the default is no data - metadata only)
    timestamps : array of float, optional
        Visibility timestamps, overriding (or fixing) the ones found in telstate
    source_name : string, optional
        Name of telstate source (used for metadata name)

    Raises
    ------
    KeyError
        If telstate lacks critical keys
    """
    def __init__(self, telstate, chunk_store=None, timestamps=None,
                 source_name='telstate'):
        self.telstate = telstate
        # Collect sensors
        sensors = {}
        for key in telstate.keys():
            if not telstate.is_immutable(key):
                sensor_name = _shorten_key(telstate, key)
                if sensor_name:
                    sensors[sensor_name] = TelstateSensorData(telstate, key)
        metadata = AttrsSensors(telstate, sensors, name=source_name)
        if timestamps is None:
            # Synthesise timestamps from the relevant telstate bits
            t0 = telstate['sync_time'] + telstate['first_timestamp']
            int_time = telstate['int_time']
            chunk_info = telstate['chunk_info']
            n_dumps = chunk_info['correlator_data']['shape'][0]
            timestamps = t0 + np.arange(n_dumps) * int_time
        if chunk_store is None:
            data = None
        else:
            # Augment chunk_info with chunk name prefix of current stream
            prefix = telstate['chunk_name']
            chunk_info = telstate['chunk_info']
            for info in chunk_info.values():
                info['prefix'] = prefix
            # Look for auxiliary streams associated with current dataset
            try:
                all_streams = telstate['sdp_archived_streams']
                capture_block_id = str(telstate['capture_block_id'])
                stream_name = str(telstate['stream_name'])
            except KeyError as e:
                logger.debug('No auxiliary capture streams found: %s', e)
            else:
                aux_streams = [str(s) for s in all_streams if s != stream_name]
                for aux_s in aux_streams:
                    # Look for chunk metadata in appropriate telstate view
                    aux_cs = telstate.SEPARATOR.join((capture_block_id, aux_s))
                    aux_telstate = telstate.root().view(aux_cs)
                    aux_prefix = aux_telstate['chunk_name']
                    aux_info = aux_telstate['chunk_info']
                    # Override available array info
                    for array, info in aux_info.items():
                        info['prefix'] = aux_prefix
                        chunk_info[array] = info
            data = ChunkStoreVisFlagsWeights(chunk_store, chunk_info)
        # Metadata and timestamps with or without data
        DataSource.__init__(self, metadata, timestamps, data)

    @classmethod
    def from_url(cls, url, chunk_store='auto', **kwargs):
        """Construct TelstateDataSource from URL (RDB file / REDIS server).

        Parameters
        ----------
        url : string
            URL serving as entry point to dataset (typically RDB file or REDIS)
        chunk_store : :class:`katdal.ChunkStore` object, optional
            Chunk store for visibility data (obtained automatically by default,
            or set to None for metadata-only dataset)
        kwargs : dict, optional
            Extra keyword arguments passed to telstate view and chunk store init
        """
        url_parts = urlparse.urlparse(url, scheme='file')
        # Merge key-value pairs from URL query with keyword arguments
        # of function (the latter takes precedence)
        url_kwargs = dict(urlparse.parse_qsl(url_parts.query))
        url_kwargs.update(kwargs)
        kwargs = url_kwargs
        # Extract Redis database number if provided
        db = int(kwargs.pop('db', '0'))
        if url_parts.scheme == 'file':
            # RDB dump file
            telstate = katsdptelstate.TelescopeState()
            try:
                telstate.load_from_file(url_parts.path)
            except OSError as err:
                raise DataSourceNotFound(str(err))
        elif url_parts.scheme == 'redis':
            # Redis server
            try:
                telstate = katsdptelstate.TelescopeState(url_parts.netloc, db)
            except katsdptelstate.ConnectionError as e:
                raise DataSourceNotFound(str(e))
        telstate = view_capture_stream(telstate, **kwargs)
        if chunk_store == 'auto':
            chunk_store = _infer_chunk_store(url_parts, telstate, **kwargs)
        return cls(telstate, chunk_store, source_name=url_parts.geturl())

def open_data_source(url, **kwargs):
    """Construct the data source described by the given URL."""
    try:
        return TelstateDataSource.from_url(url, **kwargs)
    except DataSourceNotFound as err:
        # Amend the error message for the case of an IP address without scheme
        url_parts = urlparse.urlparse(url, scheme='file')
        if url_parts.scheme == 'file' and not os.path.isfile(url_parts.path):
            raise DataSourceNotFound(
                '{} (add a URL scheme if {!r} is not meant to be a file)'
                .format(err, url_parts.path))
<|endoftext|>


<commit_before>from __future__ import division, print_function, absolute_import
import six
range = six.moves.range

import collections

import numpy as np
import networkx as nx
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib import cm
from matplotlib.patches import Ellipse
from scipy.spatial.distance import cdist

def plot_graph(G, pos=None, colors=None, nodelabels=None, nodesize=0.05,
  edgescale=1.0, nodeopts={}, labelopts={}, arrowopts={}, cmap='Paired'):
  """Plot a graphs.  Supports both directed and undirected graphs.

  For example:

  .. plot::
      :include-source:

      >>> from graphy import plotting
      >>> import networkx as nx
      >>> G=nx.karate_club_graph()
      >>> plotting.plot_graph(G, pos=nx.spring_layout(G), colors=range(G.number_of_nodes()))
      ...

  Parameters
  ----------
  G : networkx Graph object
      Graph to plot.
  pos : list of tuples
      List of (x,y) positions of nodes.  If not specified, nodes
      are arranged along a circle.
  colors : list of ints (default None)
      Color(s) to use for node faces, if desired.
  nodelabels : list of strings (default None)
      Labels to use for node labels, if desired.
  nodesize : float (default 0.05)
      Size of nodes.
  edgescale : float (default 1.0)
      Controls thickness of edges between nodes.
  nodeopts : dict (default {})
      Extra options to pass into plt.Circle call for plotting nodes.
  labelopts : dict (default {})
      Extra options to pass into plt.text call for plotting labels.
      Could be used to specify fontsize, for example.
  arrowopts : dict (default {})
      Extra options to pass into plt.arrow call for plotting edges.
  cmap : string (default 'Paired')
      Name of colormap to use.
  ax : matplotlib axis (default plt.gca())
      Axis to use for plotting graph.

  """

  class MplColorHelper:
    """Class that helps pick colors from specified colormaps.
    """
    def __init__(self, cmap_name, start_val, stop_val):
      self.cmap_name = cmap_name
      self.cmap = plt.get_cmap(cmap_name)
      self.norm = mpl.colors.Normalize(vmin=start_val, vmax=stop_val)
      self.scalarMap = cm.ScalarMappable(norm=self.norm, cmap=self.cmap)

    def get_rgb(self, val):
      return self.scalarMap.to_rgba(val)

  def intersect_two_circles(xy1, r1, xy2, r2):
    d = np.linalg.norm(xy2-xy1)

    a = (r1**2 - r2**2 + d**2) / (2 * d)
    b = d - a
    h = np.sqrt(r1**2 - a**2)
    xy3 = xy1 + a * (xy2 - xy1) / d

    ix = xy3[0] + h * (xy2[1] - xy1[1]) / d
    iy = xy3[1] - h * (xy2[0] - xy1[0]) / d

    return np.array([ix, iy])

  if pos is None:
    pos = nx.circular_layout(G)

  if colors is None:
    colors = 1

  if not isinstance(colors, collections.Iterable):
    colors = [colors,] * G.number_of_nodes()

  colors = np.asarray(colors)

  cmap_helper = MplColorHelper(cmap, colors.min(), colors.max())

  bbox = plt.gca().get_window_extent()
  asp_ratio = bbox.width/bbox.height

  xys = np.array([pos[ndx] for ndx in range(len(pos))])
  xys -= xys.min(axis=0)
  xys /= xys.max(axis=0)
  xys[:,0] *= asp_ratio

  plt.xlim([-(3*nodesize)*asp_ratio,(1+(3*nodesize))*asp_ratio])
  plt.ylim([-(3*nodesize),1+(3*nodesize)])
  plt.axis('off')

  try:
    edge_weights = nx.get_edge_attributes(G, 'weight')
  except KeyError:
    edge_weights = {}

  for edge in G.edges():
      startxy = xys[edge[0],:].copy()
      endxy   = xys[edge[1],:].copy()
      arrowdict = dict(ec='k', fc='k',
          lw=edge_weights.get(edge,1.0)*edgescale,
          length_includes_head=True,
          shape='full',
          head_length=nodesize*0.5 if G.is_directed() else 0,
          head_width=nodesize*0.5)
      for k, v in arrowopts.iteritems():
          arrowdict[k] = v

      if edge[0] == edge[1]:
          loopoffset = np.sign(startxy - xys.mean(axis=0)) * nodesize * 1.05
          cloop = plt.Circle(startxy + loopoffset, radius=nodesize*0.65, ec='k', fill=False, lw=edge_weights[edge]*edgescale)
          plt.gca().add_artist(cloop)
          arrowloc = intersect_two_circles(startxy, nodesize, startxy+loopoffset, nodesize*0.7)
          arrowlocstart = arrowloc + (arrowloc - startxy)*1e-5
          plt.arrow( arrowlocstart[0], arrowlocstart[1],  arrowloc[0]-arrowlocstart[0], arrowloc[1]-arrowlocstart[1], **arrowdict)

      else:
          angle   = np.arctan2(endxy[1]-startxy[1], endxy[0]-startxy[0])
          offset  = np.array([np.cos(angle),np.sin(angle)])*nodesize
          startxy += offset
          endxy   -= offset

          if not nx.is_directed(G) or G.has_edge(edge[1], edge[0]):
              if edge[0] > edge[1]: # will be drawn in the other direction
                  continue
              else:
                  midxy = (startxy + endxy) / 2.0
                  plt.arrow(midxy[0], midxy[1], endxy[0]-midxy[0], endxy[1]-midxy[1], **arrowdict)
                  plt.arrow(midxy[0], midxy[1], startxy[0]-midxy[0], startxy[1]-midxy[1], **arrowdict)
          else:
              plt.arrow(startxy[0], startxy[1], endxy[0]-startxy[0], endxy[1]-startxy[1], **arrowdict)

          #frac = 0.05 / np.linalg.norm(np.array([ex-sx,ey-sy]))
          #plt.annotate('', xytext=(sx,sy), xy=(ex+1e-4, ey), width=2, arrowprops=dict(arrowstyle=arrowstyle, frac=frac, facecolor='black'))

  for ndx, xy in enumerate(xys):  # Plot nodes
      cnode = plt.Circle((xy[0],xy[1]), ec='none',radius=nodesize, color=cmap_helper.get_rgb(colors[ndx]), **nodeopts)
      plt.gca().add_artist(cnode)
      if nodelabels is not None:
          plt.text(xy[0],xy[1], nodelabels[ndx], ha='center', va='center', **labelopts)

def DrawModularityFigure(mod_ts, optmod_ts=None, data_ts=None, time=None,
                         change_points=None, vis_change_points=None,
                         node_size=100,
                         y1label='Perturbation \n Modularity',
                         y2label='Change in Phase', x1label='Time',
                         state_linewidth=2,
                         changepoint_linewidth=1,
                         network_axis_size=0.2
                         ):
  """Make the Perturbation Modularity plots.

  Parameters
  ----------
  mod_ts : array, num_time x num_partitions
      Modularity over time for partitions of interest (blue dashed).
  optmod_ts : array, num_time x 1
      Modularity over time for optimal partiton (red solid).
  data_ts : array, num_time x SystemSize
      Data timeseries for second plot (phase or change in phase, ect...)
  time : array
      Time points for plots
  change_points : dict ({ time : membership_vector })
      Dictionary of partitions keyed by the first time they become optimal
  vis_change_points : list of time (default None)
      Which change_points to visualize using network plots.  By default, all
      change points will be visualized.
  node_size : int (default 100)
      node size for partition graphs
  state_linewidth : int (default 2)
      Linewidth to use for the state plots.
  changepoint_linewidth : int (default 1)
      Linewidth to use for the changepoint lines.
  network_axis_size : float (default 0.2)
      Size of the axis object for visualizing network plots.

  """

  black_color = '#262626'
  blue_color = '#779ECB'
  red_color = '#C23B22'

  if time is None:
      time = range(len(mod_ts))

  fig, (ax1, ax2) = plt.subplots(2, sharex=True, sharey=False, figsize=(6,6))
  #plt.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=0.1, hspace=0.1)
  fig.subplots_adjust(top=0.75, hspace=0.0, right=0.95, left=0.15)

  # play with the axis
  spine_list = ['top', 'right', 'left', 'bottom']
  for spine in spine_list:
      ax1.spines[spine].set_color(black_color)
      ax2.spines[spine].set_color(black_color)

  ax1.xaxis.label.set_color(black_color)
  ax1.yaxis.label.set_color(black_color)

  ax2.xaxis.label.set_color(black_color)
  ax2.yaxis.label.set_color(black_color)

  for xlabel in ax2.axes.get_xticklabels():
      xlabel.set_fontsize(12)

  for ylabel in ax1.axes.get_yticklabels():
      ylabel.set_fontsize(12)

  for ylabel in ax2.axes.get_yticklabels():
      ylabel.set_fontsize(12)

  #make last y2 label hidden
  ax2.yaxis.set_major_locator(MaxNLocator(nbins=len(ax2.get_yticklabels()), prune='upper')) # added

  ax1.set_ylabel(y1label, fontsize=14)
  ax2.set_xlabel(x1label, fontsize=14)
  ax2.set_ylabel(y2label, fontsize=14)

  # plot the data timeseries in the bottom axis
  color_idx = np.linspace(0, 1, data_ts.shape[1])
  for i, y in zip(color_idx, data_ts.T):
    plt.plot(time, y, color=plt.cm.Paired(i), lw=state_linewidth)
  plt.xlim([time.min(), time.max()])

  # plot the background modularity timeseries
  print(mod_ts.shape)
  print(time.shape)
  ax1.plot(time, mod_ts, color=blue_color, ls = '--', lw = 1)

  transFigure = fig.transFigure.inverted()

  cplist = sorted(change_points.keys())
  # get the mean position for a partition interval
  mean_cp = np.convolve(cplist + [time[-1]], 0.5*np.ones(2), 'same')

  if vis_change_points is None:
    vis_change_points = change_points.keys()

  vis_change_points = sorted(vis_change_points)
  num_vis_change_points = float(len(vis_change_points))

  for cp, mp in zip(cplist, mean_cp[1::]):
    coord3 = transFigure.transform(ax1.transData.transform([cp, ax1.get_ylim()[1]]))
    coord4 = transFigure.transform(ax2.transData.transform([cp, ax2.get_ylim()[0]]))
    fig.lines.append(mpl.lines.Line2D((coord3[0],coord4[0]),(coord3[1],coord4[1]),color='gray',
                               transform=fig.transFigure, lw=changepoint_linewidth))

  for cp, mp in zip(cplist, mean_cp[1::]):
    if cp not in vis_change_points:
      continue

    # this is an inset axes over the main axes
    a = plt.axes([0.1 + 0.85 * vis_change_points.index(cp) / num_vis_change_points, .8, network_axis_size, network_axis_size])
    plot_membership(change_points[cp], ax = a, colormap_name='Paired', node_size=node_size)
    a.text(0.0, 0.0, str(len(set(change_points[cp]))), fontsize = 20, ha='center', va='center', color = black_color)
    plt.setp(a, xticks=[], yticks=[])

    coord1 = transFigure.transform(a.transData.transform([0, -1.3]))
    coord2 = transFigure.transform(ax1.transData.transform([mp, ax1.get_ylim()[1]]))

    fig.lines.append(mpl.lines.Line2D((coord1[0],coord2[0]),(coord1[1],coord2[1]),color=black_color,
                               transform=fig.transFigure, lw=2))

  ax1.plot(time, optmod_ts, color = red_color, lw=3)<commit_msg>Updates<commit_after>from __future__ import division, print_function, absolute_import
import six
range = six.moves.range

import collections

import numpy as np
import networkx as nx
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib import cm
from matplotlib.patches import Ellipse
from scipy.spatial.distance import cdist

def plot_graph(G, pos=None, colors=None, nodelabels=None, nodesize=0.05,
  edgescale=1.0, nodeopts={}, labelopts={}, arrowopts={}, cmap='Paired'):
  """Plot a graphs.  Supports both directed and undirected graphs.
  Undirected edges are drawn as lines while directed edges are drawn
  as arrows.

  For example:

  .. plot::
      :include-source:

      >>> from graphy import plotting
      >>> import networkx as nx
      >>> G=nx.karate_club_graph()
      >>> plotting.plot_graph(G, pos=nx.spring_layout(G), colors=range(G.number_of_nodes()))
      ...

  Parameters
  ----------
  G : networkx Graph object
      Graph to plot.
  pos : list of tuples
      List of (x,y) positions of nodes.  If not specified, nodes
      are arranged along a circle.
  colors : list of ints (default None)
      Color(s) to use for node faces, if desired.
  nodelabels : list of strings (default None)
      Labels to use for node labels, if desired.
  nodesize : float (default 0.05)
      Size of nodes.
  edgescale : float (default 1.0)
      Controls thickness of edges between nodes.
  nodeopts : dict (default {})
      Extra options to pass into plt.Circle call for plotting nodes.
  labelopts : dict (default {})
      Extra options to pass into plt.text call for plotting labels.
      Could be used to specify fontsize, for example.
  arrowopts : dict (default {})
      Extra options to pass into plt.arrow call for plotting edges.
  cmap : string (default 'Paired')
      Name of colormap to use.
  ax : matplotlib axis (default plt.gca())
      Axis to use for plotting graph.

  """

  class MplColorHelper:
    """Class that helps pick colors from specified colormaps.
    """
    def __init__(self, cmap_name, start_val, stop_val):
      self.cmap_name = cmap_name
      self.cmap = plt.get_cmap(cmap_name)
      self.norm = mpl.colors.Normalize(vmin=start_val, vmax=stop_val)
      self.scalarMap = cm.ScalarMappable(norm=self.norm, cmap=self.cmap)

    def get_rgb(self, val):
      return self.scalarMap.to_rgba(val)

  def intersect_two_circles(xy1, r1, xy2, r2):
    d = np.linalg.norm(xy2-xy1)

    a = (r1**2 - r2**2 + d**2) / (2 * d)
    b = d - a
    h = np.sqrt(r1**2 - a**2)
    xy3 = xy1 + a * (xy2 - xy1) / d

    ix = xy3[0] + h * (xy2[1] - xy1[1]) / d
    iy = xy3[1] - h * (xy2[0] - xy1[0]) / d

    return np.array([ix, iy])

  if pos is None:
    pos = nx.circular_layout(G)

  if colors is None:
    colors = 1

  if not isinstance(colors, collections.Iterable):
    colors = [colors,] * G.number_of_nodes()

  colors = np.asarray(colors)

  cmap_helper = MplColorHelper(cmap, colors.min(), colors.max())

  bbox = plt.gca().get_window_extent()
  asp_ratio = bbox.width/bbox.height

  xys = np.array([pos[ndx] for ndx in range(len(pos))])
  xys -= xys.min(axis=0)
  xys /= xys.max(axis=0)
  xys[:,0] *= asp_ratio

  plt.xlim([-(3*nodesize)*asp_ratio,(1+(3*nodesize))*asp_ratio])
  plt.ylim([-(3*nodesize),1+(3*nodesize)])
  plt.axis('off')

  try:
    edge_weights = nx.get_edge_attributes(G, 'weight')
  except KeyError:
    edge_weights = {}

    arrowdict = dict(ec='k', fc='k',
        length_includes_head=True,
        shape='full',
        head_length=nodesize*0.5 if G.is_directed() else 0,
        head_width =nodesize*0.5 if G.is_directed() else 0)
    for k, v in arrowopts.iteritems():
        arrowdict[k] = v

    for edge in G.edges():
      startxy = xys[edge[0],:].copy()
      endxy   = xys[edge[1],:].copy()

      edgesize = edge_weights.get(edge,1.0)*edgescale,
      arrowdict['lw'] = edgesize

      if edge[0] == edge[1]:
          loopoffset = np.sign(startxy - xys.mean(axis=0)) * nodesize * 1.05
          cloop = plt.Circle(startxy + loopoffset, radius=nodesize*0.65, ec='k', fill=False, lw=edgesize)
          plt.gca().add_artist(cloop)
          arrowloc = intersect_two_circles(startxy, nodesize, startxy+loopoffset, nodesize*0.7)
          arrowlocstart = arrowloc + (arrowloc - startxy)*1e-5
          plt.arrow( arrowlocstart[0], arrowlocstart[1],  arrowloc[0]-arrowlocstart[0], arrowloc[1]-arrowlocstart[1], **arrowdict)

      else:
          angle   = np.arctan2(endxy[1]-startxy[1], endxy[0]-startxy[0])
          offset  = np.array([np.cos(angle),np.sin(angle)])*nodesize
          startxy += offset
          endxy   -= offset

          if not nx.is_directed(G) or G.has_edge(edge[1], edge[0]):
              if edge[0] > edge[1]: # will be drawn in the other direction
                  continue
              else:
                  midxy = (startxy + endxy) / 2.0
                  plt.arrow(midxy[0], midxy[1], endxy[0]-midxy[0], endxy[1]-midxy[1], **arrowdict)
                  plt.arrow(midxy[0], midxy[1], startxy[0]-midxy[0], startxy[1]-midxy[1], **arrowdict)
          else:
              plt.arrow(startxy[0], startxy[1], endxy[0]-startxy[0], endxy[1]-startxy[1], **arrowdict)

          #frac = 0.05 / np.linalg.norm(np.array([ex-sx,ey-sy]))
          #plt.annotate('', xytext=(sx,sy), xy=(ex+1e-4, ey), width=2, arrowprops=dict(arrowstyle=arrowstyle, frac=frac, facecolor='black'))

  for ndx, xy in enumerate(xys):  # Plot nodes
      cnode = plt.Circle((xy[0],xy[1]), ec='none',radius=nodesize, color=cmap_helper.get_rgb(colors[ndx]), **nodeopts)
      plt.gca().add_artist(cnode)
      if nodelabels is not None:
          plt.text(xy[0],xy[1], nodelabels[ndx], ha='center', va='center', **labelopts)

def DrawModularityFigure(mod_ts, optmod_ts=None, data_ts=None, time=None,
                         change_points=None, vis_change_points=None,
                         node_size=100,
                         y1label='Perturbation \n Modularity',
                         y2label='Change in Phase', x1label='Time',
                         state_linewidth=2,
                         changepoint_linewidth=1,
                         network_axis_size=0.2
                         ):
  """Make the Perturbation Modularity plots.

  Parameters
  ----------
  mod_ts : array, num_time x num_partitions
      Modularity over time for partitions of interest (blue dashed).
  optmod_ts : array, num_time x 1
      Modularity over time for optimal partiton (red solid).
  data_ts : array, num_time x SystemSize
      Data timeseries for second plot (phase or change in phase, ect...)
  time : array
      Time points for plots
  change_points : dict ({ time : membership_vector })
      Dictionary of partitions keyed by the first time they become optimal
  vis_change_points : list of time (default None)
      Which change_points to visualize using network plots.  By default, all
      change points will be visualized.
  node_size : int (default 100)
      node size for partition graphs
  state_linewidth : int (default 2)
      Linewidth to use for the state plots.
  changepoint_linewidth : int (default 1)
      Linewidth to use for the changepoint lines.
  network_axis_size : float (default 0.2)
      Size of the axis object for visualizing network plots.

  """

  black_color = '#262626'
  blue_color = '#779ECB'
  red_color = '#C23B22'

  if time is None:
      time = range(len(mod_ts))

  fig, (ax1, ax2) = plt.subplots(2, sharex=True, sharey=False, figsize=(6,6))
  #plt.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=0.1, hspace=0.1)
  fig.subplots_adjust(top=0.75, hspace=0.0, right=0.95, left=0.15)

  # play with the axis
  spine_list = ['top', 'right', 'left', 'bottom']
  for spine in spine_list:
      ax1.spines[spine].set_color(black_color)
      ax2.spines[spine].set_color(black_color)

  ax1.xaxis.label.set_color(black_color)
  ax1.yaxis.label.set_color(black_color)

  ax2.xaxis.label.set_color(black_color)
  ax2.yaxis.label.set_color(black_color)

  for xlabel in ax2.axes.get_xticklabels():
      xlabel.set_fontsize(12)

  for ylabel in ax1.axes.get_yticklabels():
      ylabel.set_fontsize(12)

  for ylabel in ax2.axes.get_yticklabels():
      ylabel.set_fontsize(12)

  #make last y2 label hidden
  ax2.yaxis.set_major_locator(MaxNLocator(nbins=len(ax2.get_yticklabels()), prune='upper')) # added

  ax1.set_ylabel(y1label, fontsize=14)
  ax2.set_xlabel(x1label, fontsize=14)
  ax2.set_ylabel(y2label, fontsize=14)

  # plot the data timeseries in the bottom axis
  color_idx = np.linspace(0, 1, data_ts.shape[1])
  for i, y in zip(color_idx, data_ts.T):
    plt.plot(time, y, color=plt.cm.Paired(i), lw=state_linewidth)
  plt.xlim([time.min(), time.max()])

  # plot the background modularity timeseries
  print(mod_ts.shape)
  print(time.shape)
  ax1.plot(time, mod_ts, color=blue_color, ls = '--', lw = 1)

  transFigure = fig.transFigure.inverted()

  cplist = sorted(change_points.keys())
  # get the mean position for a partition interval
  mean_cp = np.convolve(cplist + [time[-1]], 0.5*np.ones(2), 'same')

  if vis_change_points is None:
    vis_change_points = change_points.keys()

  vis_change_points = sorted(vis_change_points)
  num_vis_change_points = float(len(vis_change_points))

  for cp, mp in zip(cplist, mean_cp[1::]):
    coord3 = transFigure.transform(ax1.transData.transform([cp, ax1.get_ylim()[1]]))
    coord4 = transFigure.transform(ax2.transData.transform([cp, ax2.get_ylim()[0]]))
    fig.lines.append(mpl.lines.Line2D((coord3[0],coord4[0]),(coord3[1],coord4[1]),color='gray',
                               transform=fig.transFigure, lw=changepoint_linewidth))

  for cp, mp in zip(cplist, mean_cp[1::]):
    if cp not in vis_change_points:
      continue

    # this is an inset axes over the main axes
    a = plt.axes([0.1 + 0.85 * vis_change_points.index(cp) / num_vis_change_points, .8, network_axis_size, network_axis_size])
    plot_membership(change_points[cp], ax = a, colormap_name='Paired', node_size=node_size)
    a.text(0.0, 0.0, str(len(set(change_points[cp]))), fontsize = 20, ha='center', va='center', color = black_color)
    plt.setp(a, xticks=[], yticks=[])

    coord1 = transFigure.transform(a.transData.transform([0, -1.3]))
    coord2 = transFigure.transform(ax1.transData.transform([mp, ax1.get_ylim()[1]]))

    fig.lines.append(mpl.lines.Line2D((coord1[0],coord2[0]),(coord1[1],coord2[1]),color=black_color,
                               transform=fig.transFigure, lw=2))

  ax1.plot(time, optmod_ts, color = red_color, lw=3)<|endoftext|>


<commit_before>#!/usr/bin/python
#
# Copyright (c) 2005-2006 rPath, Inc.
#
# This program is distributed under the terms of the Common Public License,
# version 1.0. A copy of this license should have been distributed with this
# source file in a file called LICENSE. If it is not present, the license
# is always available at http://www.opensource.org/licenses/cpl.php.
#
# This program is distributed in the hope that it will be useful, but
# without any waranty; without even the implied warranty of merchantability
# or fitness for a particular purpose. See the Common Public License for
# full details.
#

import sys
import os
if 'CONARY_PATH' in os.environ:
    sys.path.insert(0, os.environ['CONARY_PATH'])
    sys.path.insert(0, os.environ['CONARY_PATH']+"/conary/scripts")

import time
import types

from conary import dbstore
from conary.dbstore import sqlerrors
from conary.repository.netrepos.schema import VERSION
from schema import getTables, getIndexes

if len(sys.argv) != 3:
    print "Usage: migrate <sqlite_path> <pgsql_spec>"

sqlite = dbstore.connect(sys.argv[1], driver = "sqlite")
cs = sqlite.cursor()
pgsql = dbstore.connect(sys.argv[2], driver = "postgresql")
cp = pgsql.cursor()

# create the tables, avoid the indexes
for stmt in getTables("postgresql"):
    print stmt
    cp.execute(stmt)
pgsql.loadSchema()

for t in sqlite.tables.keys():
    if t in pgsql.tables:
        continue
    print "Only in sqlite:", t
for t in pgsql.tables.keys():
    if t in sqlite.tables:
        continue
    print "Only in pgsql:", t

tList = [
    'Branches',
    'Items',
    'Versions',
    'Labels',
    'LabelMap',
    'Flavors',
    'FlavorMap',
    'FlavorScores',
    'Users',
    'UserGroups',
    'UserGroupMembers',
    'EntitlementGroups',
    'Entitlements',
    'EntitlementOwners',
    'Permissions',
    'Instances',
    'Dependencies',
    'Latest',
    'Metadata',
    'MetadataItems',
    'Nodes',
    'ChangeLogs',
    'PGPKeys',
    'PGPFingerprints',
    'Provides',
    'Requires',
    'TroveTroves',
    'TroveInfo',
    'FileStreams',
    'TroveFiles',
    ]

def timings(current, total, tstart):
    tnow = time.time()
    tpassed = max(tnow-tstart,1)
    speed = max(current/tpassed,1)
    tremaining=(total-current)/speed
    return "%d/%d %02d%% (%d rec/sec, %d:%02d passed, %d:%02d remaining)" % (
        current, total, (current*100)/max(total,1),
        speed,
        tpassed/60, tpassed % 60,
        tremaining/60, tremaining % 60)

# escape the row data based on the sqlite column types
def escape(cu, data):
    row = list(data)
    for i in range(len(data.data)):
        if data.description[i][1] == 8:
            row[i] = cu.binary(data.data[i])
    return tuple(row)

for t in tList:
    count = cs.execute("SELECT COUNT(ROWID) FROM %s" % t).fetchone()[0]
    i = 0
    cs.execute("SELECT * FROM %s" % t)
    t1 = time.time()
    while True:
        row = cs.fetchone()
        if row is None:
            break
        fields = cs.fields()
        if t.lower() == "permissions":
            if "write" in fields:
                fields[fields.index("write")] = "canWrite"
        if t.lower() == "users":
            if "user" in fields:
                fields[fields.index("user")] = "userName"
        if t.lower() == "trovetroves":
            if "byDefault" in fields:
                fields[fields.index("byDefault")] = "flags"
        if t.lower() == "trovefiles":
            # versionId was declared as a binary string in sqlite instead of integer
            row[2] = int(row[2])
        sql = "INSERT INTO %s (%s) VALUES (%s)" % (
            t, ", ".join(fields), ",".join(["?"]*len(fields)))
        i += 1
        try:
            cp.execute(sql, escape(cp, row))
        except sqlerrors.ColumnNotUnique:
            print "\r%s: DUPLICATE" % t, row
            pgsql.commit()
        except sqlerrors.ConstraintViolation, e:
            print
            print "%s: SKIPPED CONSTRAINT VIOLATION: %s" % (t, sql)
            print row, e.msg
            print
            pgsql.commit()
        except:
            print "ERROR - SQL", sql
            raise
        else:
            if i % 1000 == 0:
                t2 = time.time()
                sys.stdout.write("\r%s: %s" % (t, timings(i, count, t1)))
                sys.stdout.flush()
            if i % 10000 == 0:
                pgsql.commit()
    print "\r%s: %s %s" % (t, timings(count, count, t1), " "*10)
    pgsql.commit()

# and now create the indexes
for stmt in getIndexes("postgresql"):
    print stmt
    cp.execute(stmt)
pgsql.setVersion(VERSION)
pgsql.commit()
<commit_msg>handle boolean type columns from slite as well<commit_after>#!/usr/bin/python
#
# Copyright (c) 2005-2006 rPath, Inc.
#
# This program is distributed under the terms of the Common Public License,
# version 1.0. A copy of this license should have been distributed with this
# source file in a file called LICENSE. If it is not present, the license
# is always available at http://www.opensource.org/licenses/cpl.php.
#
# This program is distributed in the hope that it will be useful, but
# without any waranty; without even the implied warranty of merchantability
# or fitness for a particular purpose. See the Common Public License for
# full details.
#

import sys
import os
if 'CONARY_PATH' in os.environ:
    sys.path.insert(0, os.environ['CONARY_PATH'])
    sys.path.insert(0, os.environ['CONARY_PATH']+"/conary/scripts")

import time
import types

from conary import dbstore
from conary.dbstore import sqlerrors
from conary.repository.netrepos.schema import VERSION
from schema import getTables, getIndexes

if len(sys.argv) != 3:
    print "Usage: migrate <sqlite_path> <pgsql_spec>"

sqlite = dbstore.connect(sys.argv[1], driver = "sqlite")
cs = sqlite.cursor()
pgsql = dbstore.connect(sys.argv[2], driver = "postgresql")
cp = pgsql.cursor()

# create the tables, avoid the indexes
for stmt in getTables("postgresql"):
    print stmt
    cp.execute(stmt)
pgsql.loadSchema()

for t in sqlite.tables.keys():
    if t in pgsql.tables:
        continue
    print "Only in sqlite:", t
for t in pgsql.tables.keys():
    if t in sqlite.tables:
        continue
    print "Only in pgsql:", t

tList = [
    'Branches',
    'Items',
    'Versions',
    'Labels',
    'LabelMap',
    'Flavors',
    'FlavorMap',
    'FlavorScores',
    'Users',
    'UserGroups',
    'UserGroupMembers',
    'EntitlementGroups',
    'Entitlements',
    'EntitlementOwners',
    'Permissions',
    'Instances',
    'Dependencies',
    'Latest',
    'Metadata',
    'MetadataItems',
    'Nodes',
    'ChangeLogs',
    'PGPKeys',
    'PGPFingerprints',
    'Provides',
    'Requires',
    'TroveTroves',
    'TroveInfo',
    'FileStreams',
    'TroveFiles',
    ]

def timings(current, total, tstart):
    tnow = time.time()
    tpassed = max(tnow-tstart,1)
    speed = max(current/tpassed,1)
    tremaining=(total-current)/speed
    return "%d/%d %02d%% (%d rec/sec, %d:%02d passed, %d:%02d remaining)" % (
        current, total, (current*100)/max(total,1),
        speed,
        tpassed/60, tpassed % 60,
        tremaining/60, tremaining % 60)

# escape the row data based on the sqlite column types
def escape(cu, data):
    row = list(data)
    for i in range(len(data.data)):
        assert(data.description[i][1] in (0,1,6,8,9))
        if data.description[i][1] == 8:
            row[i] = cu.binary(data.data[i])
        elif data.description[i][1] == 9:
            row[i] = int(data.data[i])
    return tuple(row)

for t in tList:
    count = cs.execute("SELECT COUNT(ROWID) FROM %s" % t).fetchone()[0]
    i = 0
    cs.execute("SELECT * FROM %s" % t)
    t1 = time.time()
    while True:
        row = cs.fetchone()
        if row is None:
            break
        fields = cs.fields()
        if t.lower() == "permissions":
            if "write" in fields:
                fields[fields.index("write")] = "canWrite"
        if t.lower() == "users":
            if "user" in fields:
                fields[fields.index("user")] = "userName"
        if t.lower() == "trovetroves":
            if "byDefault" in fields:
                fields[fields.index("byDefault")] = "flags"
        if t.lower() == "trovefiles":
            # versionId was declared as a binary string in sqlite instead of integer
            row[2] = int(row[2])
        sql = "INSERT INTO %s (%s) VALUES (%s)" % (
            t, ", ".join(fields), ",".join(["?"]*len(fields)))
        i += 1
        try:
            cp.execute(sql, escape(cp, row))
        except sqlerrors.ColumnNotUnique:
            print "\r%s: DUPLICATE" % t, row
            pgsql.commit()
        except sqlerrors.ConstraintViolation, e:
            print
            print "%s: SKIPPED CONSTRAINT VIOLATION: %s" % (t, sql)
            print row, e.msg
            print
            pgsql.commit()
        except:
            print "ERROR - SQL", sql
            raise
        else:
            if i % 1000 == 0:
                t2 = time.time()
                sys.stdout.write("\r%s: %s" % (t, timings(i, count, t1)))
                sys.stdout.flush()
            if i % 10000 == 0:
                pgsql.commit()
    print "\r%s: %s %s" % (t, timings(count, count, t1), " "*10)
    pgsql.commit()

# and now create the indexes
for stmt in getIndexes("postgresql"):
    print stmt
    cp.execute(stmt)
pgsql.setVersion(VERSION)
pgsql.commit()
<|endoftext|>


<commit_before>from django import forms
from django.contrib.auth.models import User
from django.contrib.localflavor.us.forms import USPhoneNumberField
from django.utils.translation import ugettext_lazy as _
from django.forms.util import ErrorList
from django.utils.safestring import mark_safe
from django.utils.html import conditional_escape
from django.utils.encoding import force_unicode

from crispy_forms.helper import FormHelper
from crispy_forms.layout import Layout, Fieldset, ButtonHolder, Submit
from crispy_forms.bootstrap import FormActions

from userena.forms import SignupFormOnlyEmail, AuthenticationForm

from gratitude.gratitude.models import UserDetail, Gratitude

PROFILE_PLACEHOLDER = 'I am grateful for...'

class AlertErrorList(ErrorList):
    def as_ul(self):
        if not self: return u''
        return mark_safe(u'
%s
'
                % ''.join([u'
%s
' % conditional_escape(force_unicode(e)) for e in self]))

def get_datetime_field():
   return StrippingDateTimeField(required=False, widget=forms.TextInput(attrs={'class':'jquery-datetime'}))

def get_time_field():
   return StrippingTimeField(required=False, widget=forms.TimeInput(format="%H:%M", attrs={'class':'jquery-time'}))

class StrippingTimeField(forms.TimeField):
   def to_python(self, value):
      return super(StrippingTimeField, self).to_python(value.strip())

class StrippingDateTimeField(forms.DateTimeField):
   def to_python(self, value):
      return super(StrippingDateTimeField, self).to_python(value.strip())

class EmailForm(forms.Form):
   message = forms.CharField(max_length=160,widget=forms.Textarea)
   email_address = forms.CharField()

class MessagingForm(forms.Form):
   user_id = forms.IntegerField(widget=forms.HiddenInput, required=True)
   no_messages = forms.BooleanField(required=False, label=_('Do not send me any more emails'), help_text='If you want to stop getting email messages from us, check this box.')

   def save(self):
      userDetail = UserDetail.objects.filter(user_id__exact=self.cleaned_data['user_id'])[:1].get()
      userDetail.no_messages = self.cleaned_data['no_messages']
      userDetail.save()

class ProfileForm(forms.Form):
   entry0 = forms.CharField(required=False,max_length=5000,widget=forms.TextInput(attrs={'placeholder':PROFILE_PLACEHOLDER, 'autofocus':'autofocus', 'tabindex': 1}))
   entry1 = forms.CharField(required=False,max_length=5000,widget=forms.TextInput(attrs={'placeholder':PROFILE_PLACEHOLDER, 'tabindex': 2}))
   entry2 = forms.CharField(required=False,max_length=5000,widget=forms.TextInput(attrs={'placeholder':PROFILE_PLACEHOLDER, 'tabindex': 3}))
   stashed = forms.BooleanField(required=False,initial=False,widget=forms.HiddenInput())
   stash_id = forms.BooleanField(required=False,initial="",widget=forms.HiddenInput())

   def __init__(self, *args, **kwargs):
      self.user = kwargs.pop("user")
      super(ProfileForm, self).__init__(*args, **kwargs)

   def save(self):
      for entry in ['entry0', 'entry1', 'entry2']:
         cleanEntry = self.cleaned_data[entry].strip()
         stashed = self.cleaned_data['stashed']
         stash_id = self.cleaned_data['stash_id']
         if (len(cleanEntry) > 0):
            newGratitudeEntry = Gratitude()
            newGratitudeEntry.user_id = self.user.id
            newGratitudeEntry.text = cleanEntry
            newGratitudeEntry.stashed = stashed
            newGratitudeEntry.stash_id = stash_id
            newGratitudeEntry.save()

class SignupFormOnePage(SignupFormOnlyEmail):
   first_name = forms.CharField(max_length=30, widget=forms.TextInput(attrs={'placeholder':_('First name')}))
   last_name = forms.CharField(max_length=30, widget=forms.TextInput(attrs={'placeholder':_('Last name')}))
   email = forms.EmailField(widget=forms.TextInput(attrs={'class':'required', 'placeholder':_('Email'), 'maxlength':75}), label=_('Email'))
   password1 = forms.CharField(widget=forms.PasswordInput(attrs={'class':'required', 'placeholder':_('Password')}, render_value=False), label=_('Create password'))
   accept_tos = forms.BooleanField(required=True, initial=False, label=_('I accept the Art of Gratitude Terms of Service'))

   def __init__(self, *args, **kwargs):
      super(SignupFormOnlyEmail, self).__init__(*args, **kwargs)
      del self.fields['username']
      del self.fields['password2']
      self.helper = FormHelper()
      self.helper.form_id = 'signup-form'
      self.helper.form_class = 'form-horizontal pull-left'
      self.helper.layout = Layout(
         Fieldset(
             '',
             'first_name',
             'last_name',
             'email',
             'password1',
             'accept_tos'
         ),
         FormActions(
             Submit('gratitude_signup', 'Be Grateful', ),
             css_class='pagination-centered',
         )
      )

   def clean_accept_tos(self):
      if self.cleaned_data['accept_tos'] is not True:
         raise forms.ValidationError(_('You must accept the Terms of Service to sign up.'))

   def save(self):
      """ Saves the user details then calls the base class."""
      user =  super(SignupFormOnePage, self).save()
      userDetail = UserDetail()
      userDetail.user = user
      userDetail.save()
      return user

   def clean(self):
      "Overrride base class to prevent checking that password1 and password2 match, because we don't use password2"
      return self.cleaned_data

class RememberMeAuthenticationForm(AuthenticationForm):
    remember_me = forms.BooleanField(widget=forms.HiddenInput(), initial=True, required=True)
<commit_msg>fixed bug where stash_id was getting set to 'False'<commit_after>from django import forms
from django.contrib.auth.models import User
from django.contrib.localflavor.us.forms import USPhoneNumberField
from django.utils.translation import ugettext_lazy as _
from django.forms.util import ErrorList
from django.utils.safestring import mark_safe
from django.utils.html import conditional_escape
from django.utils.encoding import force_unicode

from crispy_forms.helper import FormHelper
from crispy_forms.layout import Layout, Fieldset, ButtonHolder, Submit
from crispy_forms.bootstrap import FormActions

from userena.forms import SignupFormOnlyEmail, AuthenticationForm

from gratitude.gratitude.models import UserDetail, Gratitude

PROFILE_PLACEHOLDER = 'I am grateful for...'

class AlertErrorList(ErrorList):
    def as_ul(self):
        if not self: return u''
        return mark_safe(u'
%s
'
                % ''.join([u'
%s
' % conditional_escape(force_unicode(e)) for e in self]))

def get_datetime_field():
   return StrippingDateTimeField(required=False, widget=forms.TextInput(attrs={'class':'jquery-datetime'}))

def get_time_field():
   return StrippingTimeField(required=False, widget=forms.TimeInput(format="%H:%M", attrs={'class':'jquery-time'}))

class StrippingTimeField(forms.TimeField):
   def to_python(self, value):
      return super(StrippingTimeField, self).to_python(value.strip())

class StrippingDateTimeField(forms.DateTimeField):
   def to_python(self, value):
      return super(StrippingDateTimeField, self).to_python(value.strip())

class EmailForm(forms.Form):
   message = forms.CharField(max_length=160,widget=forms.Textarea)
   email_address = forms.CharField()

class MessagingForm(forms.Form):
   user_id = forms.IntegerField(widget=forms.HiddenInput, required=True)
   no_messages = forms.BooleanField(required=False, label=_('Do not send me any more emails'), help_text='If you want to stop getting email messages from us, check this box.')

   def save(self):
      userDetail = UserDetail.objects.filter(user_id__exact=self.cleaned_data['user_id'])[:1].get()
      userDetail.no_messages = self.cleaned_data['no_messages']
      userDetail.save()

class ProfileForm(forms.Form):
   entry0 = forms.CharField(required=False,max_length=5000,widget=forms.TextInput(attrs={'placeholder':PROFILE_PLACEHOLDER, 'autofocus':'autofocus', 'tabindex': 1}))
   entry1 = forms.CharField(required=False,max_length=5000,widget=forms.TextInput(attrs={'placeholder':PROFILE_PLACEHOLDER, 'tabindex': 2}))
   entry2 = forms.CharField(required=False,max_length=5000,widget=forms.TextInput(attrs={'placeholder':PROFILE_PLACEHOLDER, 'tabindex': 3}))
   stashed = forms.BooleanField(required=False,initial=False,widget=forms.HiddenInput())
   stash_id = forms.BooleanField(required=False,initial="",widget=forms.HiddenInput())

   def __init__(self, *args, **kwargs):
      self.user = kwargs.pop("user")
      super(ProfileForm, self).__init__(*args, **kwargs)

   def save(self):
      for entry in ['entry0', 'entry1', 'entry2']:
         cleanEntry = self.cleaned_data[entry].strip()
         stashed = self.cleaned_data['stashed']
         stash_id = self.cleaned_data['stash_id']
         if (len(cleanEntry) > 0):
            newGratitudeEntry = Gratitude()
            newGratitudeEntry.user_id = self.user.id
            newGratitudeEntry.text = cleanEntry
            newGratitudeEntry.stashed = stashed
            newGratitudeEntry.stash_id = stash_id
            newGratitudeEntry.save()

   def clean_stash_id(self):
      data = self.cleaned_data['stash_id']
      if type(data) != type(''):
         data = ''
      return data

class SignupFormOnePage(SignupFormOnlyEmail):
   first_name = forms.CharField(max_length=30, widget=forms.TextInput(attrs={'placeholder':_('First name')}))
   last_name = forms.CharField(max_length=30, widget=forms.TextInput(attrs={'placeholder':_('Last name')}))
   email = forms.EmailField(widget=forms.TextInput(attrs={'class':'required', 'placeholder':_('Email'), 'maxlength':75}), label=_('Email'))
   password1 = forms.CharField(widget=forms.PasswordInput(attrs={'class':'required', 'placeholder':_('Password')}, render_value=False), label=_('Create password'))
   accept_tos = forms.BooleanField(required=True, initial=False, label=_('I accept the Art of Gratitude Terms of Service'))

   def __init__(self, *args, **kwargs):
      super(SignupFormOnlyEmail, self).__init__(*args, **kwargs)
      del self.fields['username']
      del self.fields['password2']
      self.helper = FormHelper()
      self.helper.form_id = 'signup-form'
      self.helper.form_class = 'form-horizontal pull-left'
      self.helper.layout = Layout(
         Fieldset(
             '',
             'first_name',
             'last_name',
             'email',
             'password1',
             'accept_tos'
         ),
         FormActions(
             Submit('gratitude_signup', 'Be Grateful', ),
             css_class='pagination-centered',
         )
      )

   def clean_accept_tos(self):
      if self.cleaned_data['accept_tos'] is not True:
         raise forms.ValidationError(_('You must accept the Terms of Service to sign up.'))

   def save(self):
      """ Saves the user details then calls the base class."""
      user =  super(SignupFormOnePage, self).save()
      userDetail = UserDetail()
      userDetail.user = user
      userDetail.save()
      return user

   def clean(self):
      "Overrride base class to prevent checking that password1 and password2 match, because we don't use password2"
      return self.cleaned_data

class RememberMeAuthenticationForm(AuthenticationForm):
    remember_me = forms.BooleanField(widget=forms.HiddenInput(), initial=True, required=True)
<|endoftext|>


<commit_before>import requests
import lxml.html
import logging
import pprint

logging.basicConfig(level=logging.DEBUG)

class Client:

    def __init__(self):
        self.session = requests.Session()
        self.state = dict.fromkeys((
            '__EVENTVALIDATION',
            '__VIEWSTATE',
            '__EVENTTARGET',
            '__EVENTARGUMENT',
            ), '')
        # req = requests.Request(method="GET", url='http://nyc.legistar.com/')
        # resp = self.session.send(req.prepare())
        # self.write_resp(req, resp)
        # self.update_state(resp)
        req = requests.Request(method="GET", url='http://nyc.legistar.com/Legislation.aspx')
        resp = self.session.send(req.prepare())
        self.write_resp(req, resp)
        self.update_state(resp)

    def hydrate_request(self, data):
        req = requests.Request(**data)
        return req

    def update_state(self, resp):
        '''Get the weird ASPX client state nonsense from the response
        and update the Client's state so it can be sent with future requests.
        '''
        doc = lxml.html.fromstring(resp.text)
        form = dict(doc.forms[0].fields)
        for key in set(self.state.keys()) & set(form.keys()):
            self.state[key] = form.get(key)

    def update_request(self, request):
        if request.method != 'POST':
            return
        update_keys = set(self.state.keys()) & set(request.data.keys())
        update_keys |= set(['__VIEWSTATE', '__EVENTVALIDATION'])
        for key in update_keys:
            if self.state[key]:
                request.data[key] = self.state[key]

    def write_resp(self, req, resp):
        with open('resp.html', 'wb') as f:
            f.write(resp.content)
            _data = dict(req.data)
            _data.pop('__VIEWSTATE', None)
            pprint.pprint(_data)
            import pdb; pdb.set_trace()

    def send(self, req):
        if isinstance(req, dict):
            req = self.hydrate_request(req)
        self.update_request(req)
        resp = self.session.send(req.prepare())
        client.update_state(resp)
        self.write_resp(req, resp)

# This request switches to the advanced search page.
req1 = {
    'method': 'POST',
    'url': 'http://nyc.legistar.com/Legislation.aspx',

    'cookies': {
        ' Setting-61-ASP.calendar_aspx.gridCalendar.SortExpression': 'MeetingStartDate DESC',
        ' Setting-61-ASP.legislation_aspx.gridMain.SortExpression': 'MatterID DESC',
        ' Setting-61-Calendar Body': 'All',
        ' Setting-61-Calendar Options': 'info|',
        ' Setting-61-Calendar Year': 'This Month',
        ' Setting-61-Legislation Type': 'All',
        ' Setting-61-Legislation Year': 'This Year',
        # ' __atuvc': '13%7C29',
        # ' __utma': '196938163.1626544919.1405688217.1405698268.1405700259.5',
        # ' __utmb': '196938163.2.10.1405700259',
        # ' __utmc': '196938163',
        # ' __utmz': '196938163.1405688217.1.1.utmcsr=(direct)|utmccn=(direct)|utmcmd=(none)',
        'Setting-61-Legislation Options': 'ID|Text|'
        },

    'data': {
        '__EVENTTARGET': 'ctl00$ContentPlaceHolder1$btnSwitch',
        'ctl00$ContentPlaceHolder1$chkID': 'on',
        'ctl00$ContentPlaceHolder1$chkText': 'on',
        'ctl00$ContentPlaceHolder1$lstTypeBasic': 'All Types',
        'ctl00$ContentPlaceHolder1$lstYears': 'This Year',
        'ctl00_RadScriptManager1_TSM': (
            ';;System.Web.Extensions, Version=4.0.0.0, Culture=neutral, '
            'PublicKeyToken=31bf3856ad364e35:en-US:fa6755fd-da1a-49d3-9eb4-1e473e780ecd:ea597d4b:b25378d2;'
            'Telerik.Web.UI, Version=2014.1.403.45, Culture=neutral, '
            'PublicKeyToken=121fae78165ba3d4:en-US:68d9452f-f268-45b2-8db7-8c3bbf305b8d:16e4e7cd:f76'
            '45509:24ee1bba:e330518b:88144a7a:1e771326:8e6f0d33:ed16cbdc:f46195d3:19620875:874f8ea2:'
            'cda80b3:383e4ce8:2003d0b8:aa288e2d:258f1c72:c128760b:c8618e41:1a73651d:333f8d94:58366029'),
        'ctl00_tabTop_ClientState': '{"selectedIndexes":["1"],"logEntries":[],"scrollState":{}}'
        },

    'headers': {
        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
        'Accept-Encoding': 'gzip, deflate',
        'Accept-Language': 'en-US,en;q=0.5',
        'Connection': 'keep-alive',
        'Content-Length': '37519',
        'Content-Type': 'application/x-www-form-urlencoded',
        'Cookie': 'Setting-61-Legislation Options=ID|Text|; Setting-61-Legislation Year=This Year; Setting-61-Legislation Type=All; Setting-61-ASP.legislation_aspx.gridMain.SortExpression=MatterID DESC; __utma=196938163.1626544919.1405688217.1405698268.1405700259.5; __utmc=196938163; __utmz=196938163.1405688217.1.1.utmcsr=(direct)|utmccn=(direct)|utmcmd=(none); __atuvc=13%7C29; Setting-61-Calendar Options=info|; Setting-61-Calendar Year=This Month; Setting-61-Calendar Body=All; Setting-61-ASP.calendar_aspx.gridCalendar.SortExpression=MeetingStartDate DESC; __utmb=196938163.2.10.1405700259',
        'Host': 'nyc.legistar.com',
        'Referer': 'http://nyc.legistar.com/Legislation.aspx',
        'User-Agent': 'Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:30.0) Gecko/20100101 Firefox/30.0'
        },

    }

req2 = {
    'method': 'POST',
    'url': 'http://nyc.legistar.com/Legislation.aspx',

    'cookies': {
        ' Setting-61-ASP.calendar_aspx.gridCalendar.SortExpression': 'MeetingStartDate DESC',
        ' Setting-61-ASP.legislation_aspx.gridMain.SortExpression': 'MatterID DESC',
        ' Setting-61-Calendar Body': 'All',
        ' Setting-61-Calendar Options': 'info|',
        ' Setting-61-Calendar Year': 'This Month',
        ' Setting-61-Legislation Type': 'All',
        ' Setting-61-Legislation Year': 'This Year',
        # ' __atuvc': '14%7C29',
        # ' __utma': '196938163.1626544919.1405688217.1405698268.1405700259.5',
        # ' __utmb': '196938163.3.10.1405700259',
        # ' __utmc': '196938163',
        # ' __utmz': '196938163.1405688217.1.1.utmcsr=(direct)|utmccn=(direct)|utmcmd=(none)',
        'Setting-61-Legislation Options': 'ID|Text|'
        },

    'data': {
        'ctl00$ContentPlaceHolder1$btnSearch2': 'Search Legislation',
        'ctl00$ContentPlaceHolder1$lstInControlOf': '-Select-',
        'ctl00$ContentPlaceHolder1$lstIndexedUnder': '-Select-',
        'ctl00$ContentPlaceHolder1$lstMax': 'All',
        'ctl00$ContentPlaceHolder1$lstSponsoredBy': '-Select-',
        'ctl00$ContentPlaceHolder1$lstStatus': '-Select-',
        'ctl00$ContentPlaceHolder1$lstType': '-Select-',
        'ctl00$ContentPlaceHolder1$lstYearsAdvanced': 'This Year',
        'ctl00$ContentPlaceHolder1$radFinalAction': '=',
        'ctl00$ContentPlaceHolder1$radOnAgenda': '=',
        'ctl00_ContentPlaceHolder1_lstMax_ClientState': '{"logEntries":[],"value":"1000000","text":"All","enabled":true,"checkedIndices":[],"checkedItemsTextOverflows":false}',
        'ctl00_ContentPlaceHolder1_txtFinalAction1_calendar_AD': '[[1980,1,1],[2099,12,30],[2014,7,18]]',
        'ctl00_ContentPlaceHolder1_txtFinalAction1_calendar_SD': '[]',
        'ctl00_ContentPlaceHolder1_txtFinalAction1_dateInput_ClientState': '{"enabled":true,"emptyMessage":"","validationText":"","valueAsString":"","minDateStr":"1980-01-01-00-00-00","maxDateStr":"2099-12-31-00-00-00","lastSetTextBoxValue":""}',
        'ctl00_ContentPlaceHolder1_txtFinalAction2_calendar_AD': '[[1980,1,1],[2099,12,30],[2014,7,18]]',
        'ctl00_ContentPlaceHolder1_txtFinalAction2_calendar_SD': '[]',
        'ctl00_ContentPlaceHolder1_txtFinalAction2_dateInput_ClientState': '{"enabled":true,"emptyMessage":"","validationText":"","valueAsString":"","minDateStr":"1980-01-01-00-00-00","maxDateStr":"2099-12-31-00-00-00","lastSetTextBoxValue":""}',
        'ctl00_ContentPlaceHolder1_txtOnAgenda1_calendar_AD': '[[1980,1,1],[2099,12,30],[2014,7,18]]',
        'ctl00_ContentPlaceHolder1_txtOnAgenda1_calendar_SD': '[]',
        'ctl00_ContentPlaceHolder1_txtOnAgenda1_dateInput_ClientState': '{"enabled":true,"emptyMessage":"","validationText":"","valueAsString":"","minDateStr":"1980-01-01-00-00-00","maxDateStr":"2099-12-31-00-00-00","lastSetTextBoxValue":""}',
        'ctl00_ContentPlaceHolder1_txtOnAgenda2_calendar_AD': '[[1980,1,1],[2099,12,30],[2014,7,18]]',
        'ctl00_ContentPlaceHolder1_txtOnAgenda2_calendar_SD': '[]',
        'ctl00_ContentPlaceHolder1_txtOnAgenda2_dateInput_ClientState': '{"enabled":true,"emptyMessage":"","validationText":"","valueAsString":"","minDateStr":"1980-01-01-00-00-00","maxDateStr":"2099-12-31-00-00-00","lastSetTextBoxValue":""}',
        'ctl00_RadScriptManager1_TSM': ';;System.Web.Extensions, Version=4.0.0.0, Culture=neutral, PublicKeyToken=31bf3856ad364e35:en-US:fa6755fd-da1a-49d3-9eb4-1e473e780ecd:ea597d4b:b25378d2;Telerik.Web.UI, Version=2014.1.403.45, Culture=neutral, PublicKeyToken=121fae78165ba3d4:en-US:68d9452f-f268-45b2-8db7-8c3bbf305b8d:16e4e7cd:f7645509:24ee1bba:e330518b:88144a7a:1e771326:8e6f0d33:ed16cbdc:f46195d3:19620875:874f8ea2:cda80b3:383e4ce8:2003d0b8:aa288e2d:258f1c72:8674cba1:7c926187:b7778d6c:c08e9f8a:a51ee93e:59462f1:c128760b:c8618e41:1a73651d:333f8d94:58366029',
        'ctl00_tabTop_ClientState': '{"selectedIndexes":["1"],"logEntries":[],"scrollState":{}}'
        },

    'headers': {
        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
        'Accept-Encoding': 'gzip, deflate',
        'Accept-Language': 'en-US,en;q=0.5',
        'Connection': 'keep-alive',
        'Content-Length': '834864',
        'Content-Type': 'application/x-www-form-urlencoded',
        'Cookie': 'Setting-61-Legislation Options=ID|Text|; Setting-61-Legislation Year=This Year; Setting-61-Legislation Type=All; Setting-61-ASP.legislation_aspx.gridMain.SortExpression=MatterID DESC; __utma=196938163.1626544919.1405688217.1405698268.1405700259.5; __utmc=196938163; __utmz=196938163.1405688217.1.1.utmcsr=(direct)|utmccn=(direct)|utmcmd=(none); __atuvc=14%7C29; Setting-61-Calendar Options=info|; Setting-61-Calendar Year=This Month; Setting-61-Calendar Body=All; Setting-61-ASP.calendar_aspx.gridCalendar.SortExpression=MeetingStartDate DESC; __utmb=196938163.3.10.1405700259',
        'Host': 'nyc.legistar.com',
        'Referer': 'http://nyc.legistar.com/Legislation.aspx',
        'User-Agent': 'Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:30.0) Gecko/20100101 Firefox/30.0'
        },
    }

req3 = {
    'method': 'POST',
    'url': 'http://nyc.legistar.com/Legislation.aspx',

    'cookies': {
        ' Setting-61-ASP.calendar_aspx.gridCalendar.SortExpression': 'MeetingStartDate DESC',
        ' Setting-61-ASP.legislation_aspx.gridMain.SortExpression': 'MatterID DESC',
        ' Setting-61-Calendar Body': 'All',
        ' Setting-61-Calendar Options': 'info|',
        ' Setting-61-Calendar Year': 'This Month',
        ' Setting-61-Legislation Type': '',
        ' Setting-61-Legislation Year': 'This Year',
        # ' __atuvc': '15%7C29',
        # ' __utma': '196938163.1626544919.1405688217.1405698268.1405700259.5',
        # ' __utmb': '196938163.4.10.1405700259',
        # ' __utmc': '196938163',
        # ' __utmz': '196938163.1405688217.1.1.utmcsr=(direct)|utmccn=(direct)|utmcmd=(none)',
        'Setting-61-Legislation Options': 'ID|Text|'
        },

    'data': {
        '__EVENTTARGET': 'ctl00$ContentPlaceHolder1$gridMain$ctl00$ctl02$ctl00$ctl04',
        'ctl00$ContentPlaceHolder1$lstInControlOf': '-Select-',
        'ctl00$ContentPlaceHolder1$lstIndexedUnder': '-Select-',
        'ctl00$ContentPlaceHolder1$lstMax': 'All',
        'ctl00$ContentPlaceHolder1$lstSponsoredBy': '-Select-',
        'ctl00$ContentPlaceHolder1$lstStatus': '-Select-',
        'ctl00$ContentPlaceHolder1$lstType': '-Select-',
        'ctl00$ContentPlaceHolder1$lstYearsAdvanced': 'This Year',
        'ctl00$ContentPlaceHolder1$radFinalAction': '=',
        'ctl00$ContentPlaceHolder1$radOnAgenda': '=',
        'ctl00_ContentPlaceHolder1_txtFinalAction1_calendar_AD': '[[1980,1,1],[2099,12,30],[2014,7,18]]',
        'ctl00_ContentPlaceHolder1_txtFinalAction1_calendar_SD': '[]',
        'ctl00_ContentPlaceHolder1_txtFinalAction1_dateInput_ClientState': '{"enabled":true,"emptyMessage":"","validationText":"","valueAsString":"","minDateStr":"1980-01-01-00-00-00","maxDateStr":"2099-12-31-00-00-00","lastSetTextBoxValue":""}',
        'ctl00_ContentPlaceHolder1_txtFinalAction2_calendar_AD': '[[1980,1,1],[2099,12,30],[2014,7,18]]',
        'ctl00_ContentPlaceHolder1_txtFinalAction2_calendar_SD': '[]',
        'ctl00_ContentPlaceHolder1_txtFinalAction2_dateInput_ClientState': '{"enabled":true,"emptyMessage":"","validationText":"","valueAsString":"","minDateStr":"1980-01-01-00-00-00","maxDateStr":"2099-12-31-00-00-00","lastSetTextBoxValue":""}',
        'ctl00_ContentPlaceHolder1_txtOnAgenda1_calendar_AD': '[[1980,1,1],[2099,12,30],[2014,7,18]]',
        'ctl00_ContentPlaceHolder1_txtOnAgenda1_calendar_SD': '[]',
        'ctl00_ContentPlaceHolder1_txtOnAgenda1_dateInput_ClientState': '{"enabled":true,"emptyMessage":"","validationText":"","valueAsString":"","minDateStr":"1980-01-01-00-00-00","maxDateStr":"2099-12-31-00-00-00","lastSetTextBoxValue":""}',
        'ctl00_ContentPlaceHolder1_txtOnAgenda2_calendar_AD': '[[1980,1,1],[2099,12,30],[2014,7,18]]',
        'ctl00_ContentPlaceHolder1_txtOnAgenda2_calendar_SD': '[]',
        'ctl00_ContentPlaceHolder1_txtOnAgenda2_dateInput_ClientState': '{"enabled":true,"emptyMessage":"","validationText":"","valueAsString":"","minDateStr":"1980-01-01-00-00-00","maxDateStr":"2099-12-31-00-00-00","lastSetTextBoxValue":""}',
        'ctl00_RadScriptManager1_TSM': ';;System.Web.Extensions, Version=4.0.0.0, Culture=neutral, PublicKeyToken=31bf3856ad364e35:en-US:fa6755fd-da1a-49d3-9eb4-1e473e780ecd:ea597d4b:b25378d2;Telerik.Web.UI, Version=2014.1.403.45, Culture=neutral, PublicKeyToken=121fae78165ba3d4:en-US:68d9452f-f268-45b2-8db7-8c3bbf305b8d:16e4e7cd:f7645509:24ee1bba:e330518b:88144a7a:1e771326:8e6f0d33:ed16cbdc:f46195d3:19620875:874f8ea2:cda80b3:383e4ce8:2003d0b8:aa288e2d:258f1c72:8674cba1:7c926187:b7778d6c:c08e9f8a:a51ee93e:59462f1:c128760b:c8618e41:1a73651d:333f8d94:58366029',
        'ctl00_tabTop_ClientState': '{"selectedIndexes":["1"],"logEntries":[],"scrollState":{}}'
        },

    'headers': {
        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
        'Accept-Encoding': 'gzip, deflate',
        'Accept-Language': 'en-US,en;q=0.5',
        'Connection': 'keep-alive',
        'Content-Length': '953155',
        'Content-Type': 'application/x-www-form-urlencoded',
        'Cookie': 'Setting-61-Legislation Options=ID|Text|; Setting-61-Legislation Year=This Year; Setting-61-Legislation Type=; Setting-61-ASP.legislation_aspx.gridMain.SortExpression=MatterID DESC; __utma=196938163.1626544919.1405688217.1405698268.1405700259.5; __utmc=196938163; __utmz=196938163.1405688217.1.1.utmcsr=(direct)|utmccn=(direct)|utmcmd=(none); __atuvc=15%7C29; Setting-61-Calendar Options=info|; Setting-61-Calendar Year=This Month; Setting-61-Calendar Body=All; Setting-61-ASP.calendar_aspx.gridCalendar.SortExpression=MeetingStartDate DESC; __utmb=196938163.4.10.1405700259',
        'Host': 'nyc.legistar.com',
        'Referer': 'http://nyc.legistar.com/Legislation.aspx',
        'User-Agent': 'Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:30.0) Gecko/20100101 Firefox/30.0'
        },
    }

def incr_page(req):
    event_target = req['data']['__EVENTTARGET']
    event_target = event_target[:-1] + str(int(event_target[-1]) + 2)
    req['data']['__EVENTTARGET'] = event_target
    return req

client = Client()
client.send(req1)
client.send(req2)
client.send(req3)
client.send(incr_page(req3))
import pdb; pdb.set_trace()<commit_msg>old pagination hackery checked in?<commit_after><|endoftext|>


<commit_before>from contest.contest_info import get_running_contests
from problem.models import Submission, SubmissionDetail
from users.models import User
from utils.user_info import validate_user

def regroup_submission(submissions):
    submission_groups = []
    for submission in submissions:
        submission_groups.append({
            'grouper': submission,
            'list': SubmissionDetail.objects.filter(
                sid=submission.id
            ).order_by('tid')
        })

    return submission_groups

def get_visible_submission(user):
    """Get all submissions that can be viewed by the given user."""

    user = validate_user(user)
    submissions = Submission.objects.all()

    # Admin can view all submissions
    if user.has_admin_auth():
        return submissions

    # No one can view admins' submissions
    submissions = submissions.exclude(
        user__in=User.objects.filter(user_level=User.ADMIN)
    )

    # During the contest, only owner/coowner can view contestants' detail
    contests = get_running_contests()

    return submissions
<commit_msg>exclude contest<commit_after>from datetime import datetime
from django.db.models import Q

from contest.contest_info import get_running_contests, get_freeze_time_datetime
from contest.models import Contest
from problem.models import Problem, Submission, SubmissionDetail
from users.models import User
from utils.user_info import validate_user, has_contest_ownership

def regroup_submission(submissions):
    submission_groups = []
    for submission in submissions:
        submission_groups.append({
            'grouper': submission,
            'list': SubmissionDetail.objects.filter(
                sid=submission.id
            ).order_by('tid')
        })

    return submission_groups

def get_visible_submission(user):
    """Get all submissions that can be viewed by the given user."""
    user = validate_user(user)
    submissions = Submission.objects.all()

    # Admin can view all submissions
    if user.has_admin_auth():
        return submissions

    # No one can view admins' submissions
    submissions = submissions.exclude(
        user__in=User.objects.filter(user_level=User.ADMIN)
    )

    # Invisible problems' submission can't be seen
    invisible_problem = Problem.objects.filter(
        visible=False
    ).exclude(
        owner=user
    )

    # Not ended contest has something to judge
    contests = Contest.objects.filter(
        is_homework=False,
        end_time__gte=datetime.now()
    )
    for contest in contests:
        # Those who don't have contest ownership has some limitations
        if not has_contest_ownership(user, contest):
            # 1. Can't view contest owner/coowners' submission before the end
            submissions = submissions.exclude(
                Q(user__in=contest.coowner.all()) | Q(user=contest.owner),
                problem__in=contest.problem.all(),
                submit_time__gte=contest.creation_time
            )
            # 2. Can't view other's submission after contest freeze
            submissions = submissions.exclude(
                problem__in=contest.problem.all(),
                submit_time__gte=get_freeze_time_datetime(contest)
            )

    return submissions
<|endoftext|>