<commit_before># Authors: Denis Engemann <d.engemann@fz-juelich.de>
#
# License: BSD (3-clause)

import pandas as pd
import copy
import numpy as np
from .event import Discrete
from .viz import plot_epochs
from .utils import string_types

class Epochs(object):
    """ Create epoched data

    Parameters
    ----------
    raw : instance of Raw | list
        The raw instance to create epochs from. Can also be a list of raw
        instances to use.
    events : ndarray (n_epochs) | list
        The events to construct epochs around. Can also be a list of
        arrays.
    event_id : int | dict
        The event ID to use. Can be a dict to supply multiple event types
        by name.
    tmin : float
        The time window before a particular event in seconds.
    tmax : float
        The time window after a particular event in seconds.

    Returns
    -------
    epochs : instance of Epochs
        The epoched dataset.
    """
    def __init__(self, raw, events, event_id, tmin, tmax):
        self.event_id = event_id
        self.tmin = tmin
        self.tmax = tmax
        self._current = 0
        event_keys = None
        if isinstance(event_id, dict):
            my_event_id = event_id.values()
            event_keys = dict((v, k) for k, v in event_id.items())
        elif np.isscalar(event_id):
            my_event_id = [event_id]
        if not isinstance(raw, list):
            raw = [raw]
        if not isinstance(events, list):
            events = [events]
        if len(raw) != len(events):
            raise ValueError('raw and events must match')
        assert len(raw) > 0
        # figure out parameters to use
        idx_offsets = raw[0].time_as_index([self.tmin, self.tmax])
        n_samples = idx_offsets[1] - idx_offsets[0]
        self._n_times = n_samples
        self.times = np.linspace(self.tmin, self.tmax, self._n_times)
        self.info = dict(sfreq=raw[0].info['sfreq'],
                         data_cols=raw[0].info['data_cols'])
        for r in raw[1:]:
            if r.info['sfreq'] != raw[0].info['sfreq']:
                raise RuntimeError('incompatible raw files')
        # process each raw file
        outs = list()

        for rr, ee in zip(raw, events):
            out = self._process_raw_events(rr, ee, my_event_id,
                                           event_keys, idx_offsets)
            outs.append(out)

        _samples, _discretes, _events = zip(*outs)
        _events = np.concatenate(_events)
        self._n_epochs = len(_events)
        # ignore index to allow for sorting + keep unique values

        _flatten = lambda x: [i for i in x for i in i]
        _samples = _flatten(_samples)
        if len(raw) > 1:  # update sample counts for index
            this_n_epochs = 0
            for ss in _samples:
                if this_n_epochs:
                    ss['epoch_idx'] += this_n_epochs
                this_n_epochs += len(ss.epoch_idx.unique())

        _data = pd.concat(_samples, ignore_index=True)
        # important for multiple conditions
        _data = _data.sort(['epoch_idx', 'time'])
        self._data = _data
        self._data['times'] = np.tile(self.times, self._n_epochs)
        self._data.set_index(['epoch_idx', 'times'], drop=True,
                             inplace=True, verify_integrity=True)
        assert self._n_epochs == self._data.index.values.max()[0] + 1
        self.info['discretes'] = _flatten(_discretes)
        self.events = _events

    def _process_raw_events(self, raw, events, my_event_id, event_keys,
                            idx_offsets):
        data, times = raw[:]
        discrete_inds = [[] for _ in range(3)]
        sample_inds = dict((k, []) for k in my_event_id)
        saccade_inds, fixation_inds, blink_inds = discrete_inds
        keep_idx = []
        # prevent the evil
        events = events[events[:, 0].argsort()]
        for ii, (event, this_id) in enumerate(events):
            if this_id not in my_event_id:
                continue
            this_time = times[event]
            this_tmin, this_tmax = this_time + self.tmin, this_time + self.tmax
            inds_min, inds_max = raw.time_as_index(this_time)[0] + idx_offsets
            if max([inds_min, inds_max]) >= len(raw.samples):
                continue
            if min([inds_min, inds_max]) < 0:
                continue
            inds = np.arange(inds_min, inds_max)

            sample_inds[this_id].append([inds, ii])
            for kind, parsed in zip(raw.info['event_types'], discrete_inds):
                df = raw.discrete.get(kind, kind)
                assert(set([a <= b for a, b in
                       df[['stime', 'etime']].values]) == set([True]))
                event_in_window = np.where((df['stime'] >= this_tmin) &
                                           (df['etime'] <= this_tmax))
                parsed.append([event_in_window[0], ii, this_id, this_time])
            keep_idx.append(ii)
        events = events[keep_idx]

        discretes = []
        for kind, parsed in zip(['saccades', 'fixations', 'blinks'],
                                discrete_inds):
            this_in = raw.discrete.get(kind, None)
            if this_in is not None:
                this_discrete = Discrete()
                for inds, epochs_idx, this_id, this_time in parsed:
                    this_id = (this_id if event_keys is None else
                               event_keys[this_id])
                    if inds.any().any():
                        df = this_in.ix[inds]
                        df['event_id'] = this_id
                        # don't use -= here b/c pandas complains
                        df['stime'] -= this_time
                        df['etime'] -= this_time
                        this_discrete.append(df)
                    else:
                        this_discrete.append([])
                this_name = kind + '_'
                setattr(self, this_name, this_discrete)  # XXX FIX
                discretes += [this_name]

        _samples = []
        c = np.concatenate
        track_inds = []
        for this_id, values in sample_inds.items():
            ind, _ = zip(*values)
            ind = [i[:self._n_times] for i in ind]
            df = raw.samples.ix[c(ind)]
            this_id = this_id if event_keys is None else event_keys[this_id]
            df['event_id'] = this_id
            count = c([np.repeat(vv, self._n_times) for _, vv in values])
            df['epoch_idx'] = count
            _samples.append(df)
            track_inds.extend([len(i) for i in ind])

        assert set(track_inds) == set([self._n_times])
        return _samples, discretes, events

    def __repr__(self):
        s = '<Epochs | {0} events | tmin: {1} tmax: {2}>'
        return s.format(len(self.events), self.tmin, self.tmax)

    def __iter__(self):
        """To make iteration over epochs easy.
        """
        self._current = 0
        return self

    def next(self, return_event_id=False):
        """To make iteration over epochs easy.
        """
        if self._current >= self._n_epochs:
            raise StopIteration
        epoch = self.data_frame.ix[self._current]
        epoch = epoch[self.info['data_cols']].values.T
        self._current += 1
        if not return_event_id:
            return epoch
        else:
            return epoch, self.events[self._current - 1][-1]

    def __next__(self, *args, **kwargs):
        return self.next(*args, **kwargs)

    @property
    def data(self):
        out = self._data[self.info['data_cols']].values
        out = out.reshape(len(self.events),
                          len(self.times),
                          len(self.info['data_cols']))
        return np.transpose(out, [0, 2, 1])

    @property
    def len(self):
        return self._n_epochs

    @property
    def data_frame(self):
        return self._data

    @property
    def ch_names(self):
        return [k for k in self.data_frame.columns]

    def __getitem__(self, idx):
        out = self.copy()
        if isinstance(idx, string_types):
            if idx not in self.event_id:
                raise ValueError('ID not found')
            idx = self.event_id[idx]
            idx = np.where(self.events[:, -1] == idx)[0]
        elif (isinstance(idx, list) and isinstance(idx[0],
              string_types)):
            idx_list = []
            for ii in idx:
                ii = self.event_id[ii]
                idx_list.append(np.where(self.events[:, -1] == ii)[0])
            idx = np.concatenate(idx_list)
        elif np.isscalar(idx):
            idx = [idx]
        elif isinstance(idx, slice):
            idx = np.arange(*idx.indices(idx.stop))
        # XXX inquire whether Index.map works across pandas versions (fast)
        idx = np.sort(idx)
        midx = [i for i in out._data.index if i[0] in idx]  # ... slow
        out._data = out._data.ix[midx]
        out.events = out.events[idx]
        out._n_epochs = len(idx)
        for discrete in self.info['discretes']:
            disc = vars(self)[discrete]
            setattr(out, discrete, Discrete(disc[k] for k in idx))
        return out

    def copy(self):
        """Return a copy of Epochs.
        """
        return copy.deepcopy(self)

    def plot(self, epoch_idx=None, picks=None, n_chunks=20,
             title_str='#%003i', show=True, draw_discrete=None,
             discrete_colors=None, block=False):
        """ Visualize single trials using Trellis plot.

        Parameters
        ----------
        epoch_idx : array-like | int | None
            The epochs to visualize. If None, the first 20 epochs are shown.
            Defaults to None.
        n_chunks : int
            The number of chunks to use for display.
        picks : array-like | None
            Channels to be included. If None only good data channels are used.
            Defaults to None
        lines : array-like | list of tuple
            Events to draw as vertical lines
        title_str : None | str
            The string formatting to use for axes titles. If None, no titles
            will be shown. Defaults expand to ``#001, #002, ...``
        show : bool
            Whether to show the figure or not.
        draw_discrete : {saccades, blinks, fixations} | list-like | None |
            The events to draw as vertical lines.
        discrete_colors: : list-like | None
            list of str or color objects with length of discrete events drawn.
        block : bool
            Whether to halt program execution until the figure is closed.
            Useful for rejecting bad trials on the fly by clicking on a
            sub plot.

        Returns
        -------
        fig : Instance of matplotlib.figure.Figure
            The figure.
        """
        return plot_epochs(epochs=self, epoch_idx=epoch_idx, picks=picks,
                           n_chunks=n_chunks, title_str=title_str,
                           show=show, draw_discrete=draw_discrete,
                           discrete_colors=discrete_colors,
                           block=block)
<commit_msg>more flakes<commit_after># Authors: Denis Engemann <d.engemann@fz-juelich.de>
#
# License: BSD (3-clause)

import pandas as pd
import copy
import numpy as np
from .event import Discrete
from .viz import plot_epochs
from .utils import string_types

class Epochs(object):
    """ Create epoched data

    Parameters
    ----------
    raw : instance of Raw | list
        The raw instance to create epochs from. Can also be a list of raw
        instances to use.
    events : ndarray (n_epochs) | list
        The events to construct epochs around. Can also be a list of
        arrays.
    event_id : int | dict
        The event ID to use. Can be a dict to supply multiple event types
        by name.
    tmin : float
        The time window before a particular event in seconds.
    tmax : float
        The time window after a particular event in seconds.

    Returns
    -------
    epochs : instance of Epochs
        The epoched dataset.
    """
    def __init__(self, raw, events, event_id, tmin, tmax):
        self.event_id = event_id
        self.tmin = tmin
        self.tmax = tmax
        self._current = 0
        event_keys = None
        if isinstance(event_id, dict):
            my_event_id = event_id.values()
            event_keys = dict((v, k) for k, v in event_id.items())
        elif np.isscalar(event_id):
            my_event_id = [event_id]
        if not isinstance(raw, list):
            raw = [raw]
        if not isinstance(events, list):
            events = [events]
        if len(raw) != len(events):
            raise ValueError('raw and events must match')
        assert len(raw) > 0
        # figure out parameters to use
        idx_offsets = raw[0].time_as_index([self.tmin, self.tmax])
        n_samples = idx_offsets[1] - idx_offsets[0]
        self._n_times = n_samples
        self.times = np.linspace(self.tmin, self.tmax, self._n_times)
        self.info = dict(sfreq=raw[0].info['sfreq'],
                         data_cols=raw[0].info['data_cols'])
        for r in raw[1:]:
            if r.info['sfreq'] != raw[0].info['sfreq']:
                raise RuntimeError('incompatible raw files')
        # process each raw file
        outs = list()

        for rr, ee in zip(raw, events):
            out = self._process_raw_events(rr, ee, my_event_id,
                                           event_keys, idx_offsets)
            outs.append(out)

        _samples, _discretes, _events = zip(*outs)
        _events = np.concatenate(_events)
        self._n_epochs = len(_events)
        # ignore index to allow for sorting + keep unique values

        _flatten = lambda x: [i for i in x for i in i]
        _samples = _flatten(_samples)
        if len(raw) > 1:  # update sample counts for index
            this_n_epochs = 0
            for ss in _samples:
                if this_n_epochs:
                    ss['epoch_idx'] += this_n_epochs
                this_n_epochs += len(ss.epoch_idx.unique())

        _data = pd.concat(_samples, ignore_index=True)
        # important for multiple conditions
        _data = _data.sort(['epoch_idx', 'time'])
        self._data = _data
        self._data['times'] = np.tile(self.times, self._n_epochs)
        self._data.set_index(['epoch_idx', 'times'], drop=True,
                             inplace=True, verify_integrity=True)
        assert self._n_epochs == self._data.index.values.max()[0] + 1
        self.info['discretes'] = _flatten(_discretes)
        self.events = _events

    def _process_raw_events(self, raw, events, my_event_id, event_keys,
                            idx_offsets):
        data, times = raw[:]
        discrete_inds = [[] for _ in range(3)]
        sample_inds = dict((k, []) for k in my_event_id)
        saccade_inds, fixation_inds, blink_inds = discrete_inds
        keep_idx = []
        # prevent the evil
        events = events[events[:, 0].argsort()]
        for ii, (event, this_id) in enumerate(events):
            if this_id not in my_event_id:
                continue
            this_time = times[event]
            this_tmin, this_tmax = this_time + self.tmin, this_time + self.tmax
            inds_min, inds_max = raw.time_as_index(this_time)[0] + idx_offsets
            if max([inds_min, inds_max]) >= len(raw.samples):
                continue
            if min([inds_min, inds_max]) < 0:
                continue
            inds = np.arange(inds_min, inds_max)

            sample_inds[this_id].append([inds, ii])
            for kind, parsed in zip(raw.info['event_types'], discrete_inds):
                df = raw.discrete.get(kind, kind)
                assert(set([a <= b for a, b in
                       df[['stime', 'etime']].values]) == set([True]))
                event_in_window = np.where((df['stime'] >= this_tmin) &
                                           (df['etime'] <= this_tmax))
                parsed.append([event_in_window[0], ii, this_id, this_time])
            keep_idx.append(ii)
        events = events[keep_idx]

        discretes = []
        for kind, parsed in zip(['saccades', 'fixations', 'blinks'],
                                discrete_inds):
            this_in = raw.discrete.get(kind, None)
            if this_in is not None:
                this_discrete = Discrete()
                for inds, epochs_idx, this_id, this_time in parsed:
                    this_id = (this_id if event_keys is None else
                               event_keys[this_id])
                    if inds.any().any():
                        df = this_in.ix[inds]
                        df['event_id'] = this_id
                        # don't use -= here b/c pandas complains
                        df['stime'] -= this_time
                        df['etime'] -= this_time
                        this_discrete.append(df)
                    else:
                        this_discrete.append([])
                this_name = kind + '_'
                setattr(self, this_name, this_discrete)  # XXX FIX
                discretes += [this_name]

        _samples = []
        c = np.concatenate
        track_inds = []
        for this_id, values in sample_inds.items():
            ind, _ = zip(*values)
            ind = [i[:self._n_times] for i in ind]
            df = raw.samples.ix[c(ind)]
            this_id = this_id if event_keys is None else event_keys[this_id]
            df['event_id'] = this_id
            count = c([np.repeat(vv, self._n_times) for _, vv in values])
            df['epoch_idx'] = count
            _samples.append(df)
            track_inds.extend([len(i) for i in ind])

        assert set(track_inds) == set([self._n_times])
        return _samples, discretes, events

    def __repr__(self):
        s = '<Epochs | {0} events | tmin: {1} tmax: {2}>'
        return s.format(len(self.events), self.tmin, self.tmax)

    def __iter__(self):
        """To make iteration over epochs easy.
        """
        self._current = 0
        return self

    def next(self, return_event_id=False):
        """To make iteration over epochs easy.
        """
        if self._current >= self._n_epochs:
            raise StopIteration
        epoch = self.data_frame.ix[self._current]
        epoch = epoch[self.info['data_cols']].values.T
        self._current += 1
        if not return_event_id:
            return epoch
        else:
            return epoch, self.events[self._current - 1][-1]

    def __next__(self, *args, **kwargs):
        return self.next(*args, **kwargs)

    @property
    def data(self):
        out = self._data[self.info['data_cols']].values
        out = out.reshape(len(self.events),
                          len(self.times),
                          len(self.info['data_cols']))
        return np.transpose(out, [0, 2, 1])

    @property
    def len(self):
        return self._n_epochs

    @property
    def data_frame(self):
        return self._data

    @property
    def ch_names(self):
        return [k for k in self.data_frame.columns]

    def __getitem__(self, idx):
        out = self.copy()
        if isinstance(idx, string_types):
            if idx not in self.event_id:
                raise ValueError('ID not found')
            idx = self.event_id[idx]
            idx = np.where(self.events[:, -1] == idx)[0]
        elif (isinstance(idx, list) and isinstance(idx[0],
              string_types)):
            idx_list = []
            for ii in idx:
                ii = self.event_id[ii]
                idx_list.append(np.where(self.events[:, -1] == ii)[0])
            idx = np.concatenate(idx_list)
        elif np.isscalar(idx):
            idx = [idx]
        elif isinstance(idx, slice):
            idx = np.arange(*idx.indices(idx.stop))
        # XXX inquire whether Index.map works across pandas versions (fast)
        idx = np.sort(idx)
        midx = [i for i in out._data.index if i[0] in idx]  # ... slow
        out._data = out._data.ix[midx]
        out.events = out.events[idx]
        out._n_epochs = len(idx)
        for discrete in self.info['discretes']:
            disc = vars(self)[discrete]
            setattr(out, discrete, Discrete(disc[k] for k in idx))
        return out

    def copy(self):
        """Return a copy of Epochs.
        """
        return copy.deepcopy(self)

    def plot(self, epoch_idx=None, picks=None, n_chunks=20,
             title_str='#%003i', show=True, draw_discrete=None,
             discrete_colors=None, block=False):
        """ Visualize single trials using Trellis plot.

        Parameters
        ----------
        epoch_idx : array-like | int | None
            The epochs to visualize. If None, the first 20 epochs are shown.
            Defaults to None.
        n_chunks : int
            The number of chunks to use for display.
        picks : array-like | None
            Channels to be included. If None only good data channels are used.
            Defaults to None
        lines : array-like | list of tuple
            Events to draw as vertical lines
        title_str : None | str
            The string formatting to use for axes titles. If None, no titles
            will be shown. Defaults expand to ``#001, #002, ...``
        show : bool
            Whether to show the figure or not.
        draw_discrete : {saccades, blinks, fixations} | list-like | None |
            The events to draw as vertical lines.
        discrete_colors: : list-like | None
            list of str or color objects with length of discrete events drawn.
        block : bool
            Whether to halt program execution until the figure is closed.
            Useful for rejecting bad trials on the fly by clicking on a
            sub plot.

        Returns
        -------
        fig : Instance of matplotlib.figure.Figure
            The figure.
        """
        return plot_epochs(epochs=self, epoch_idx=epoch_idx, picks=picks,
                           n_chunks=n_chunks, title_str=title_str,
                           show=show, draw_discrete=draw_discrete,
                           discrete_colors=discrete_colors,
                           block=block)
<|endoftext|>


<commit_before># Authors: Denis Engemann <d.engemann@fz-juelich.de>
#
# License: BSD (3-clause)

import pandas as pd
import copy
import numpy as np
from .event import Discrete
from .viz import plot_epochs
from .utils import string_types

class Epochs(object):
    """ Create epoched data

    Parameters
    ----------
    raw : instance of Raw | list
        The raw instance to create epochs from. Can also be a list of raw
        instances to use.
    events : ndarray (n_epochs) | list
        The events to construct epochs around. Can also be a list of
        arrays.
    event_id : int | dict
        The event ID to use. Can be a dict to supply multiple event types
        by name.
    tmin : float
        The time window before a particular event in seconds.
    tmax : float
        The time window after a particular event in seconds.

    Returns
    -------
    epochs : instance of Epochs
        The epoched dataset.
    """
    def __init__(self, raw, events, event_id, tmin, tmax):
        self.event_id = event_id
        self.tmin = tmin
        self.tmax = tmax
        self._current = 0
        event_keys = None
        if isinstance(event_id, dict):
            my_event_id = event_id.values()
            event_keys = dict((v, k) for k, v in event_id.items())
        elif np.isscalar(event_id):
            my_event_id = [event_id]
        if not isinstance(raw, list):
            raw = [raw]
        if not isinstance(events, list):
            events = [events]
        if len(raw) != len(events):
            raise ValueError('raw and events must match')
        assert len(raw) > 0
        # figure out parameters to use
        idx_offsets = raw[0].time_as_index([self.tmin, self.tmax])
        self.info = dict(sfreq=raw[0].info['sfreq'],
                         data_cols=raw[0].info['data_cols'])
        for r in raw[1:]:
            if r.info['sfreq'] != raw[0].info['sfreq']:
                raise RuntimeError('incompatible raw files')
        n_samples = idx_offsets[1] - idx_offsets[0]
        self._n_times = n_samples
        self.times = np.linspace(self.tmin, self.tmax, self._n_times)
        # process each raw file
        outs = list()
        for rr, ee in zip(raw, events):
            out = self._process_raw_events(rr, ee, my_event_id,
                                           event_keys, idx_offsets)
            outs.append(out)
        _samples = outs[0][0]
        _discretes = outs[0][2]
        _events = np.concatenate([out[3] for out in outs])
        for out in outs[1:]:
            _samples.extend(out[0])
            _discretes.extend(out[2])  # XXX NEED TO ADJUST TIME INDICES?
        self._n_epochs = len(_events)
        # ignore index to allow for sorting + keep unique values
        _data = pd.concat(_samples, ignore_index=True)
        # important for multiple conditions
        _data = _data.sort(['epoch_idx', 'time'])
        self._data = _data
        self._data['times'] = np.tile(self.times, self._n_epochs)
        self._data.set_index(['epoch_idx', 'times'], drop=True,
                             inplace=True, verify_integrity=True)
        assert self._n_epochs == self._data.index.values.max()[0]
        self.info['discretes'] = _discretes
        self.events = _events

    def _process_raw_events(self, raw, events, my_event_id, event_keys,
                            idx_offsets):
        data, times = raw[:]
        discrete_inds = [[] for _ in range(3)]
        sample_inds = dict((k, []) for k in my_event_id)
        saccade_inds, fixation_inds, blink_inds = discrete_inds
        keep_idx = []
        # prevent the evil
        events = events[events[:, 0].argsort()]
        for ii, (event, this_id) in enumerate(events):
            if this_id not in my_event_id:
                continue
            this_time = times[event]
            this_tmin, this_tmax = this_time + self.tmin, this_time + self.tmax
            inds_min, inds_max = raw.time_as_index(this_time)[0] + idx_offsets
            if max([inds_min, inds_max]) >= len(raw.samples):
                continue
            if min([inds_min, inds_max]) < 0:
                continue
            inds = np.arange(inds_min, inds_max)

            sample_inds[this_id].append([inds, ii])
            for kind, parsed in zip(raw.info['event_types'], discrete_inds):
                df = raw.discrete.get(kind, kind)
                assert(set([a <= b for a, b in
                       df[['stime', 'etime']].values]) == set([True]))
                event_in_window = np.where((df['stime'] >= this_tmin) &
                                           (df['etime'] <= this_tmax))
                parsed.append([event_in_window[0], ii, this_id, this_time])
            keep_idx.append(ii)
        events = events[keep_idx]

        discretes = []
        for kind, parsed in zip(['saccades', 'fixations', 'blinks'],
                                discrete_inds):
            this_in = raw.discrete.get(kind, None)
            if this_in is not None:
                this_discrete = Discrete()
                for inds, epochs_idx, this_id, this_time in parsed:
                    this_id = (this_id if event_keys is None else
                               event_keys[this_id])
                    if inds.any().any():
                        df = this_in.ix[inds]
                        df['event_id'] = this_id
                        # don't use -= here b/c pandas complains
                        df['stime'] -= this_time
                        df['etime'] -= this_time
                        this_discrete.append(df)
                    else:
                        this_discrete.append([])
                this_name = kind + '_'
                setattr(self, this_name, this_discrete)  # XXX FIX
                discretes += [this_name]

        _samples = []
        c = np.concatenate
        track_inds = []
        for this_id, values in sample_inds.items():
            ind, _ = zip(*values)
            ind = [i[:self._n_times] for i in ind]
            df = raw.samples.ix[c(ind)]
            this_id = this_id if event_keys is None else event_keys[this_id]
            df['event_id'] = this_id
            count = c([np.repeat(vv, self._n_times) for _, vv in values])
            df['epoch_idx'] = count
            _samples.append(df)
            track_inds.extend([len(i) for i in ind])

        assert set(track_inds) == set([self._n_times])
        return _samples, discretes, events

    def __repr__(self):
        s = '<Epochs | {0} events | tmin: {1} tmax: {2}>'
        return s.format(len(self.events), self.tmin, self.tmax)

    def __iter__(self):
        """To make iteration over epochs easy.
        """
        self._current = 0
        return self

    def next(self, return_event_id=False):
        """To make iteration over epochs easy.
        """
        if self._current >= self._n_epochs:
            raise StopIteration
        epoch = self.data_frame.ix[self._current]
        epoch = epoch[self.info['data_cols']].values.T
        self._current += 1
        if not return_event_id:
            return epoch
        else:
            return epoch, self.events[self._current - 1][-1]

    def __next__(self, *args, **kwargs):
        return self.next(*args, **kwargs)

    @property
    def data(self):
        out = self._data[self.info['data_cols']].values
        out = out.reshape(len(self.events),
                          len(self.times),
                          len(self.info['data_cols']))
        return np.transpose(out, [0, 2, 1])

    @property
    def len(self):
        return self._n_epochs

    @property
    def data_frame(self):
        return self._data

    @property
    def ch_names(self):
        return [k for k in self.data_frame.columns]

    def __getitem__(self, idx):
        out = self.copy()
        if isinstance(idx, string_types):
            if idx not in self.event_id:
                raise ValueError('ID not found')
            idx = self.event_id[idx]
            idx = np.where(self.events[:, -1] == idx)[0]
        elif (isinstance(idx, list) and isinstance(idx[0],
              string_types)):
            idx_list = []
            for ii in idx:
                ii = self.event_id[ii]
                idx_list.append(np.where(self.events[:, -1] == ii)[0])
            idx = np.concatenate(idx_list)
        elif np.isscalar(idx):
            idx = [idx]
        elif isinstance(idx, slice):
            idx = np.arange(*idx.indices(idx.stop))
        # XXX inquire whether Index.map works across pandas versions (fast)
        idx = np.sort(idx)
        midx = [i for i in out._data.index if i[0] in idx]  # ... slow
        out._data = out._data.ix[midx]
        out.events = out.events[idx]
        out._n_epochs = len(idx)
        for discrete in self.info['discretes']:
            disc = vars(self)[discrete]
            setattr(out, discrete, Discrete(disc[k] for k in idx))
        return out

    def copy(self):
        """Return a copy of Epochs.
        """
        return copy.deepcopy(self)

    def plot(self, epoch_idx=None, picks=None, n_chunks=20,
             title_str='#%003i', show=True, draw_discrete=None,
             discrete_colors=None, block=False):
        """ Visualize single trials using Trellis plot.

        Parameters
        ----------
        epoch_idx : array-like | int | None
            The epochs to visualize. If None, the first 20 epochs are shown.
            Defaults to None.
        n_chunks : int
            The number of chunks to use for display.
        picks : array-like | None
            Channels to be included. If None only good data channels are used.
            Defaults to None
        lines : array-like | list of tuple
            Events to draw as vertical lines
        title_str : None | str
            The string formatting to use for axes titles. If None, no titles
            will be shown. Defaults expand to ``#001, #002, ...``
        show : bool
            Whether to show the figure or not.
        draw_discrete : {saccades, blinks, fixations} | list-like | None |
            The events to draw as vertical lines.
        discrete_colors: : list-like | None
            list of str or color objects with length of discrete events drawn.
        block : bool
            Whether to halt program execution until the figure is closed.
            Useful for rejecting bad trials on the fly by clicking on a
            sub plot.

        Returns
        -------
        fig : Instance of matplotlib.figure.Figure
            The figure.
        """
        return plot_epochs(epochs=self, epoch_idx=epoch_idx, picks=picks,
                           n_chunks=n_chunks, title_str=title_str,
                           show=show, draw_discrete=draw_discrete,
                           discrete_colors=discrete_colors,
                           block=block)
<commit_msg>FIX various stuff<commit_after># Authors: Denis Engemann <d.engemann@fz-juelich.de>
#
# License: BSD (3-clause)

import pandas as pd
import copy
import numpy as np
from .event import Discrete
from .viz import plot_epochs
from .utils import string_types

class Epochs(object):
    """ Create epoched data

    Parameters
    ----------
    raw : instance of Raw | list
        The raw instance to create epochs from. Can also be a list of raw
        instances to use.
    events : ndarray (n_epochs) | list
        The events to construct epochs around. Can also be a list of
        arrays.
    event_id : int | dict
        The event ID to use. Can be a dict to supply multiple event types
        by name.
    tmin : float
        The time window before a particular event in seconds.
    tmax : float
        The time window after a particular event in seconds.

    Returns
    -------
    epochs : instance of Epochs
        The epoched dataset.
    """
    def __init__(self, raw, events, event_id, tmin, tmax):
        self.event_id = event_id
        self.tmin = tmin
        self.tmax = tmax
        self._current = 0
        event_keys = None
        if isinstance(event_id, dict):
            my_event_id = event_id.values()
            event_keys = dict((v, k) for k, v in event_id.items())
        elif np.isscalar(event_id):
            my_event_id = [event_id]
        if not isinstance(raw, list):
            raw = [raw]
        if not isinstance(events, list):
            events = [events]
        if len(raw) != len(events):
            raise ValueError('raw and events must match')
        assert len(raw) > 0
        # figure out parameters to use
        idx_offsets = raw[0].time_as_index([self.tmin, self.tmax])
        self.info = dict(sfreq=raw[0].info['sfreq'],
                         data_cols=raw[0].info['data_cols'])
        for r in raw[1:]:
            if r.info['sfreq'] != raw[0].info['sfreq']:
                raise RuntimeError('incompatible raw files')
        n_samples = idx_offsets[1] - idx_offsets[0]
        self._n_times = n_samples
        self.times = np.linspace(self.tmin, self.tmax, self._n_times)
        # process each raw file
        outs = list()
        for rr, ee in zip(raw, events):
            out = self._process_raw_events(rr, ee, my_event_id,
                                           event_keys, idx_offsets)
            outs.append(out)

        _samples, _discretes, _events = zip(*outs)
        _events = np.concatenate(_events)

        # import pdb;pdb.set_trace()
        self._n_epochs = len(_events)
        # ignore index to allow for sorting + keep unique values
        _data = pd.concat(_samples[0], ignore_index=True)
        # important for multiple conditions
        _data = _data.sort(['epoch_idx', 'time'])
        self._data = _data
        self._data['times'] = np.tile(self.times, self._n_epochs)
        self._data.set_index(['epoch_idx', 'times'], drop=True,
                             inplace=True, verify_integrity=True)
        assert self._n_epochs == self._data.index.values.max()[0] + 1
        self.info['discretes'] = _discretes[0]
        self.events = _events

    def _process_raw_events(self, raw, events, my_event_id, event_keys,
                            idx_offsets):
        data, times = raw[:]
        discrete_inds = [[] for _ in range(3)]
        sample_inds = dict((k, []) for k in my_event_id)
        saccade_inds, fixation_inds, blink_inds = discrete_inds
        keep_idx = []
        # prevent the evil
        events = events[events[:, 0].argsort()]
        for ii, (event, this_id) in enumerate(events):
            if this_id not in my_event_id:
                continue
            this_time = times[event]
            this_tmin, this_tmax = this_time + self.tmin, this_time + self.tmax
            inds_min, inds_max = raw.time_as_index(this_time)[0] + idx_offsets
            if max([inds_min, inds_max]) >= len(raw.samples):
                continue
            if min([inds_min, inds_max]) < 0:
                continue
            inds = np.arange(inds_min, inds_max)

            sample_inds[this_id].append([inds, ii])
            for kind, parsed in zip(raw.info['event_types'], discrete_inds):
                df = raw.discrete.get(kind, kind)
                assert(set([a <= b for a, b in
                       df[['stime', 'etime']].values]) == set([True]))
                event_in_window = np.where((df['stime'] >= this_tmin) &
                                           (df['etime'] <= this_tmax))
                parsed.append([event_in_window[0], ii, this_id, this_time])
            keep_idx.append(ii)
        events = events[keep_idx]

        discretes = []
        for kind, parsed in zip(['saccades', 'fixations', 'blinks'],
                                discrete_inds):
            this_in = raw.discrete.get(kind, None)
            if this_in is not None:
                this_discrete = Discrete()
                for inds, epochs_idx, this_id, this_time in parsed:
                    this_id = (this_id if event_keys is None else
                               event_keys[this_id])
                    if inds.any().any():
                        df = this_in.ix[inds]
                        df['event_id'] = this_id
                        # don't use -= here b/c pandas complains
                        df['stime'] -= this_time
                        df['etime'] -= this_time
                        this_discrete.append(df)
                    else:
                        this_discrete.append([])
                this_name = kind + '_'
                setattr(self, this_name, this_discrete)  # XXX FIX
                discretes += [this_name]

        _samples = []
        c = np.concatenate
        track_inds = []
        for this_id, values in sample_inds.items():
            ind, _ = zip(*values)
            ind = [i[:self._n_times] for i in ind]
            df = raw.samples.ix[c(ind)]
            this_id = this_id if event_keys is None else event_keys[this_id]
            df['event_id'] = this_id
            count = c([np.repeat(vv, self._n_times) for _, vv in values])
            df['epoch_idx'] = count
            _samples.append(df)
            track_inds.extend([len(i) for i in ind])

        assert set(track_inds) == set([self._n_times])
        return _samples, discretes, events

    def __repr__(self):
        s = '<Epochs | {0} events | tmin: {1} tmax: {2}>'
        return s.format(len(self.events), self.tmin, self.tmax)

    def __iter__(self):
        """To make iteration over epochs easy.
        """
        self._current = 0
        return self

    def next(self, return_event_id=False):
        """To make iteration over epochs easy.
        """
        if self._current >= self._n_epochs:
            raise StopIteration
        epoch = self.data_frame.ix[self._current]
        epoch = epoch[self.info['data_cols']].values.T
        self._current += 1
        if not return_event_id:
            return epoch
        else:
            return epoch, self.events[self._current - 1][-1]

    def __next__(self, *args, **kwargs):
        return self.next(*args, **kwargs)

    @property
    def data(self):
        out = self._data[self.info['data_cols']].values
        out = out.reshape(len(self.events),
                          len(self.times),
                          len(self.info['data_cols']))
        return np.transpose(out, [0, 2, 1])

    @property
    def len(self):
        return self._n_epochs

    @property
    def data_frame(self):
        return self._data

    @property
    def ch_names(self):
        return [k for k in self.data_frame.columns]

    def __getitem__(self, idx):
        out = self.copy()
        if isinstance(idx, string_types):
            if idx not in self.event_id:
                raise ValueError('ID not found')
            idx = self.event_id[idx]
            idx = np.where(self.events[:, -1] == idx)[0]
        elif (isinstance(idx, list) and isinstance(idx[0],
              string_types)):
            idx_list = []
            for ii in idx:
                ii = self.event_id[ii]
                idx_list.append(np.where(self.events[:, -1] == ii)[0])
            idx = np.concatenate(idx_list)
        elif np.isscalar(idx):
            idx = [idx]
        elif isinstance(idx, slice):
            idx = np.arange(*idx.indices(idx.stop))
        # XXX inquire whether Index.map works across pandas versions (fast)
        idx = np.sort(idx)
        midx = [i for i in out._data.index if i[0] in idx]  # ... slow
        out._data = out._data.ix[midx]
        out.events = out.events[idx]
        out._n_epochs = len(idx)
        for discrete in self.info['discretes']:
            disc = vars(self)[discrete]
            setattr(out, discrete, Discrete(disc[k] for k in idx))
        return out

    def copy(self):
        """Return a copy of Epochs.
        """
        return copy.deepcopy(self)

    def plot(self, epoch_idx=None, picks=None, n_chunks=20,
             title_str='#%003i', show=True, draw_discrete=None,
             discrete_colors=None, block=False):
        """ Visualize single trials using Trellis plot.

        Parameters
        ----------
        epoch_idx : array-like | int | None
            The epochs to visualize. If None, the first 20 epochs are shown.
            Defaults to None.
        n_chunks : int
            The number of chunks to use for display.
        picks : array-like | None
            Channels to be included. If None only good data channels are used.
            Defaults to None
        lines : array-like | list of tuple
            Events to draw as vertical lines
        title_str : None | str
            The string formatting to use for axes titles. If None, no titles
            will be shown. Defaults expand to ``#001, #002, ...``
        show : bool
            Whether to show the figure or not.
        draw_discrete : {saccades, blinks, fixations} | list-like | None |
            The events to draw as vertical lines.
        discrete_colors: : list-like | None
            list of str or color objects with length of discrete events drawn.
        block : bool
            Whether to halt program execution until the figure is closed.
            Useful for rejecting bad trials on the fly by clicking on a
            sub plot.

        Returns
        -------
        fig : Instance of matplotlib.figure.Figure
            The figure.
        """
        return plot_epochs(epochs=self, epoch_idx=epoch_idx, picks=picks,
                           n_chunks=n_chunks, title_str=title_str,
                           show=show, draw_discrete=draw_discrete,
                           discrete_colors=discrete_colors,
                           block=block)
<|endoftext|>


<commit_before>#!/usr/bin/env python

'''
    nrepl.bencode
    -------------

    This module provides BEncode-protocol support.

    :copyright: (c) 2013 by Chas Emerick.
    :license: MIT, see LICENSE for more details.
'''

try:
    from cStringIO import StringIO
except ImportError:
    from io import StringIO

import sys

# Some code so we can use different features without worrying about versions.
PY2 = sys.version_info[0] == 2
if not PY2:
    text_type = str
    string_types = (str, bytes)
    unichr = chr
else:
    text_type = unicode
    string_types = (str, unicode)
    unichr = unichr

def _read_byte(s):
    return s.read(1)

def _read_int(s, terminator=None, init_data=None):
    int_chrs = init_data or []
    while True:
        c = _read_byte(s)
        if not c.isdigit() or c == terminator or not c:
            break
        else:
            int_chrs.append(c)
    return int(''.join(int_chrs))

def _read_bytes(s, n):
    data = StringIO()
    cnt = 0
    while cnt < n:
        m = s.read(n - cnt)
        if not m:
            raise Exception("Invalid bytestring, unexpected end of input.")
        data.write(m)
        cnt += len(m)
    data.flush()
    # Taking into account that Python3 can't decode strings
    try:
        ret = data.getvalue().decode("UTF-8")
    except AttributeError:
        ret = data.getvalue()
    return ret

def _read_delimiter(s):
    d = _read_byte(s)
    if d.isdigit():
        d = _read_int(s, ":", [d])
    return d

def _read_list(s):
    data = []
    while True:
        datum = _read_datum(s)
        if not datum:
            break
        data.append(datum)
    return data

def _read_map(s):
    i = iter(_read_list(s))
    return dict(zip(i, i))

_read_fns = {"i": _read_int,
             "l": _read_list,
             "d": _read_map,
             "e": lambda _: None,
             # EOF
             None: lambda _: None}

def _read_datum(s):
    delim = _read_delimiter(s)
    if delim:
        return _read_fns.get(delim, lambda s: _read_bytes(s, delim))(s)

def encode_bytes(x):
    return (str(len(x)).encode('utf-8'), b':', x)

def encode_string(x):
    try:
        s = x.encode('utf-8')
    except UnicodeDecodeError:
        return encode_bytes(x)

    return (str(len(s)).encode('utf-8'), b':', s)

def _write_datum(x, out):
    if isinstance(x, string_types):
        #x = x.encode("utf-8")
        # TODO revisit encodings, this is surely not right. Python
        # (2.x, anyway) conflates bytes and strings, but 3.x does not...
        l = len(x.replace("\n", "nn"))
        out.write(str(l))
        out.write(":")
        out.write(x)
    elif isinstance(x, int):
        out.write("i")
        out.write(str(x))
        out.write("e")
    elif isinstance(x, (list, tuple)):
        out.write("l")
        for v in x:
            _write_datum(v, out)
        out.write("e")
    elif isinstance(x, dict):
        out.write("d")
        for k, v in x.items():
            _write_datum(k, out)
            _write_datum(v, out)
        out.write("e")
    out.flush()

def encode(v):
    "bencodes the given value, may be a string, integer, list, or dict."
    s = StringIO()
    _write_datum(v, s)
    return s.getvalue()

def decode_file(file):
    while True:
        x = _read_datum(file)
        if not x:
            break
        yield x

def decode(string):
    "Generator that yields decoded values from the input string."
    return decode_file(StringIO(string))

class BencodeIO(object):
    def __init__(self, file, on_close=None):
        self._file = file
        self._on_close = on_close

    def read(self):
        return _read_datum(self._file)

    def __iter__(self):
        return self

    def next(self):
        v = self.read()
        if not v:
            raise StopIteration
        return v

    def __next__(self):
        # In Python3, __next__ it is an own special class.
        v = self.read()
        if not v:
            raise StopIteration
        return v

    def write(self, v):
        return _write_datum(v, self._file)

    def flush(self):
        if self._file.flush:
            self._file.flush()

    def close(self):
        # Run the on_close handler if one exists, which can do something
        # useful like cleanly close a socket. (Note that .close() on a
        # socket.makefile('rw') does some kind of unclean close.)
        if self._on_close is not None:
            self._on_close()
        else:
            self._file.close()<commit_msg>fix for linux bencode newline encodings<commit_after>#!/usr/bin/env python

'''
    nrepl.bencode
    -------------

    This module provides BEncode-protocol support.

    :copyright: (c) 2013 by Chas Emerick.
    :license: MIT, see LICENSE for more details.
'''

try:
    from cStringIO import StringIO
except ImportError:
    from io import StringIO

import sys, os

# Some code so we can use different features without worrying about versions.
PY2 = sys.version_info[0] == 2
if not PY2:
    text_type = str
    string_types = (str, bytes)
    unichr = chr
else:
    text_type = unicode
    string_types = (str, unicode)
    unichr = unichr

def _read_byte(s):
    return s.read(1)

def _read_int(s, terminator=None, init_data=None):
    int_chrs = init_data or []
    while True:
        c = _read_byte(s)
        if not c.isdigit() or c == terminator or not c:
            break
        else:
            int_chrs.append(c)
    return int(''.join(int_chrs))

def _read_bytes(s, n):
    data = StringIO()
    cnt = 0
    while cnt < n:
        m = s.read(n - cnt)
        if not m:
            raise Exception("Invalid bytestring, unexpected end of input.")
        data.write(m)
        cnt += len(m)
    data.flush()
    # Taking into account that Python3 can't decode strings
    try:
        ret = data.getvalue().decode("UTF-8")
    except AttributeError:
        ret = data.getvalue()
    return ret

def _read_delimiter(s):
    d = _read_byte(s)
    if d.isdigit():
        d = _read_int(s, ":", [d])
    return d

def _read_list(s):
    data = []
    while True:
        datum = _read_datum(s)
        if not datum:
            break
        data.append(datum)
    return data

def _read_map(s):
    i = iter(_read_list(s))
    return dict(zip(i, i))

_read_fns = {"i": _read_int,
             "l": _read_list,
             "d": _read_map,
             "e": lambda _: None,
             # EOF
             None: lambda _: None}

def _read_datum(s):
    delim = _read_delimiter(s)
    if delim:
        return _read_fns.get(delim, lambda s: _read_bytes(s, delim))(s)

def encode_bytes(x):
    return (str(len(x)).encode('utf-8'), b':', x)

def encode_string(x):
    try:
        s = x.encode('utf-8')
    except UnicodeDecodeError:
        return encode_bytes(x)

    return (str(len(s)).encode('utf-8'), b':', s)

def _write_datum(x, out):
    if isinstance(x, string_types):
        #x = x.encode("utf-8")
        # TODO revisit encodings, this is surely not right. Python
        # (2.x, anyway) conflates bytes and strings, but 3.x does not...
        l = len(x.replace("\n", "nn"))
        if os.name is "posix":
            l = len(x)
        out.write(str(l))
        out.write(":")
        out.write(x)
    elif isinstance(x, int):
        out.write("i")
        out.write(str(x))
        out.write("e")
    elif isinstance(x, (list, tuple)):
        out.write("l")
        for v in x:
            _write_datum(v, out)
        out.write("e")
    elif isinstance(x, dict):
        out.write("d")
        for k, v in x.items():
            _write_datum(k, out)
            _write_datum(v, out)
        out.write("e")
    out.flush()

def encode(v):
    "bencodes the given value, may be a string, integer, list, or dict."
    s = StringIO()
    _write_datum(v, s)
    return s.getvalue()

def decode_file(file):
    while True:
        x = _read_datum(file)
        if not x:
            break
        yield x

def decode(string):
    "Generator that yields decoded values from the input string."
    return decode_file(StringIO(string))

class BencodeIO(object):
    def __init__(self, file, on_close=None):
        self._file = file
        self._on_close = on_close

    def read(self):
        return _read_datum(self._file)

    def __iter__(self):
        return self

    def next(self):
        v = self.read()
        if not v:
            raise StopIteration
        return v

    def __next__(self):
        # In Python3, __next__ it is an own special class.
        v = self.read()
        if not v:
            raise StopIteration
        return v

    def write(self, v):
        return _write_datum(v, self._file)

    def flush(self):
        if self._file.flush:
            self._file.flush()

    def close(self):
        # Run the on_close handler if one exists, which can do something
        # useful like cleanly close a socket. (Note that .close() on a
        # socket.makefile('rw') does some kind of unclean close.)
        if self._on_close is not None:
            self._on_close()
        else:
            self._file.close()<|endoftext|>


<commit_before>class BigJSON:

    _WHITESPACE = '\t\n '
    _READBUF_CHUNK_SIZE = 1024*1024

    def __init__(self, file):
        self.file = file

        # This buffer is for reading and peeking
        self.readbuf = ''

        self.root_node = Node(self)

    def __getitem__(self, index):
        return self.root_node[index]

    def _skip_whitespace(self):
        while True:
            # Read whitespace from current readbuffer
            whitespace_in_readbuf = 0
            while whitespace_in_readbuf < len(self.readbuf) and self.readbuf[whitespace_in_readbuf] in BigJSON._WHITESPACE:
                whitespace_in_readbuf += 1
            self.readbuf = self.readbuf[whitespace_in_readbuf:]

            # If readbuffer is not empty, then read more
            if len(self.readbuf) > 0:
                return

            # Readbuffer is empty, so load more
            self._fill_readbuf(BigJSON._READBUF_CHUNK_SIZE)

    def _get(self):
        self._fill_readbuf(1)
        if len(self.readbuf) == 0:
            raise Exception('Unexpected end of file when getting next byte!')
        result = self.readbuf[0]
        self.readbuf = self.readbuf[1:]
        return result

    def _read(self, amount):
        self._fill_readbuf(amount)
        if len(self.readbuf) < amount:
            raise Exception('Unexpected end of file when getting next byte!')
        result = self.readbuf[:amount]
        self.readbuf = self.readbuf[amount:]
        return result

    def _peek(self):
        self._fill_readbuf(1)
        if len(self.readbuf) == 0:
            return None
        return self.readbuf[0]

    def _is_next(self, s):
        self._fill_readbuf(len(s))
        return self.readbuf.startswith(s)

    def _skip_if_next(self, s):
        """ If next bytes are same as in 's', then skip them and return True.
        """
        if self._is_next(s):
            self.readbuf = self.readbuf[len(s):]
            return True
        return False

    def _fill_readbuf(self, desired_size):
        read_amount = desired_size - len(self.readbuf)
        if read_amount <= 0:
            return
        self.readbuf += self.file.read(read_amount)

    def _tell_read_pos(self):
        return self.file.tell() - len(self.readbuf)

    def _seek(self, pos):
        self.readbuf = u''
        self.file.seek(pos)

class Node:

    TYPE_NULL = 0
    TYPE_FALSE = 1
    TYPE_TRUE = 2
    TYPE_NUMBER = 3
    TYPE_STRING = 4
    TYPE_ARRAY = 5
    TYPE_OBJECT = 6

    def __init__(self, bj):
        self.bj = bj

        self.bj._skip_whitespace()

        # Read part of file to analyze type of it
        if self.bj._skip_if_next('null'):
            self.type = Node.TYPE_NULL

        elif self.bj._skip_if_next('false'):
            self.type = Node.TYPE_FALSE

        elif self.bj._skip_if_next('true'):
            self.type = Node.TYPE_TRUE

        elif self.bj._peek() in '-0123456789':
            self.type = Node.TYPE_NUMBER
            num = self.bj._get()
            # Check sign
            if num == '-':
                num += self.bj._get()
            # Read integer part
            if num[-1] != '0':
                while self.bj._peek() in '0123456789':
                    num += self.bj._get()
            # Read possible decimal part and convert to float
            if self.bj._peek() == '.':
                self.bj._get()
                num += '.' + self.bj._get()
                if num[-1] not in '0123456790':
                    raise Exception('Expected digit after dot!')
                while self.bj._peek() in '0123456789':
                    num += self.bj._get()
                self.number = float(num)
            else:
                self.number = int(num)
            # Read possible exponent
            if self.bj._peek() in 'eE':
                self.bj._get()
                exp = self.bj._get()
                exp_neg = False
                if exp == '-':
                    exp_neg = True
                    exp = self.bj._get()
                elif exp == '+':
                    exp = self.bj._get()
                while self.bj._peek() in '0123456789':
                    exp += self.bj._get()
                exp = int(exp)
                exp = 10 ** exp
                if exp_neg:
                    self.number = float(self.number) / exp
                else:
                    self.number *= exp

        elif self.bj._skip_if_next('"'):
            self.type = Node.TYPE_STRING

            self.string = u''

            while True:
                c = self.bj._get()

                if c == u'"':
                    break

                if c == u'\\':
                    c = self.bj._get()
                    if c == u'"':
                        self.string += u'"'
                    elif c == u'\\':
                        self.string += u'\\'
                    elif c == u'/':
                        self.string += u'/'
                    elif c == u'b':
                        self.string += u'\b'
                    elif c == u'f':
                        self.string += u'\f'
                    elif c == u'n':
                        self.string += u'\n'
                    elif c == u'r':
                        self.string += u'\r'
                    elif c == u't':
                        self.string += u'\t'
                    elif c == u'u':
                        unicode_bytes = self.bj._read(4)
                        self.string += ('\\u' + unicode_bytes).decode('unicode_escape')
                    else:
                        raise Exception(u'Unexpected {} in backslash encoding!'.format(c))

                else:
                    self.string += c

        elif self.bj._skip_if_next('['):
            self.type = Node.TYPE_ARRAY
            self.fully_read = False
            # TODO: Replace this single lookup position with some kind of lookup table!
            self.children_begin = self.bj._tell_read_pos()

        elif self.bj._skip_if_next('{'):
            self.type = Node.TYPE_OBJECT
            self.fully_read = False
            raise Exception('Not implemented yet!')

        else:
            raise Exception('Unexpected bytes!')

    def get_type(self):
        return ['null', 'false', 'true', 'number', 'string', 'array', 'object'][self.type]

    def __getitem__(self, index):
        if self.type == Node.TYPE_ARRAY:

            # TODO: Support negative indexes!

            # Rewind to requested child
            self.bj._seek(self.children_begin)
            self.bj._skip_whitespace()
            while True:
                if self.bj._is_next(']'):
                    raise IndexError('Out of range!')

                child = Node(self.bj)
                child._read_fully()

                if index == 0:
                    return child._get_value()

                # Skip comma and whitespace around it
                self.bj._skip_whitespace()
                if not self.bj._skip_if_next(','):
                    raise IndexError('Out of range!')
                self.bj._skip_whitespace()

                index -= 1

        elif self.type == Node.TYPE_OBJECT:

            raise Exception('NOT IMPLEMENTED YET!')

        elif self.type == Node.TYPE_STRING:

            return self.string[index]

        else:

            raise Exception('Index operator does not work with type "{}"!'.format(self.get_type))

    def _get_value(self):
        """ If possible, returns basic python object, but
        with complex Nodes, returns the Node itself.
        """
        if self.type == Node.TYPE_NULL:
            return None
        if self.type == Node.TYPE_FALSE:
            return False
        if self.type == Node.TYPE_TRUE:
            return True
        if self.type == Node.TYPE_NUMBER:
            return self.number
        if self.type == Node.TYPE_STRING:
            return self.string
        return self

    def _read_fully(self):

        # If already fully read
        if self.type in [Node.TYPE_NULL, Node.TYPE_FALSE, Node.TYPE_TRUE, Node.TYPE_NUMBER, Node.TYPE_STRING]:
            return

        if self.type == Node.TYPE_ARRAY:
            raise Exception('Not implemented yet!')

        else:
            raise Exception('Not implemented yet!')
<commit_msg>Reading object.<commit_after>class BigJSON:

    _WHITESPACE = '\t\n '
    _READBUF_CHUNK_SIZE = 1024*1024

    def __init__(self, file):
        self.file = file

        # This buffer is for reading and peeking
        self.readbuf = ''

        self.root_node = Node(self)

    def __getitem__(self, index):
        return self.root_node[index]

    def _skip_whitespace(self):
        while True:
            # Read whitespace from current readbuffer
            whitespace_in_readbuf = 0
            while whitespace_in_readbuf < len(self.readbuf) and self.readbuf[whitespace_in_readbuf] in BigJSON._WHITESPACE:
                whitespace_in_readbuf += 1
            self.readbuf = self.readbuf[whitespace_in_readbuf:]

            # If readbuffer is not empty, then read more
            if len(self.readbuf) > 0:
                return

            # Readbuffer is empty, so load more
            self._fill_readbuf(BigJSON._READBUF_CHUNK_SIZE)

    def _get(self):
        self._fill_readbuf(1)
        if len(self.readbuf) == 0:
            raise Exception('Unexpected end of file when getting next byte!')
        result = self.readbuf[0]
        self.readbuf = self.readbuf[1:]
        return result

    def _read(self, amount):
        self._fill_readbuf(amount)
        if len(self.readbuf) < amount:
            raise Exception('Unexpected end of file when getting next byte!')
        result = self.readbuf[:amount]
        self.readbuf = self.readbuf[amount:]
        return result

    def _peek(self):
        self._fill_readbuf(1)
        if len(self.readbuf) == 0:
            return None
        return self.readbuf[0]

    def _is_next(self, s):
        self._fill_readbuf(len(s))
        return self.readbuf.startswith(s)

    def _skip_if_next(self, s):
        """ If next bytes are same as in 's', then skip them and return True.
        """
        if self._is_next(s):
            self.readbuf = self.readbuf[len(s):]
            return True
        return False

    def _fill_readbuf(self, desired_size):
        read_amount = desired_size - len(self.readbuf)
        if read_amount <= 0:
            return
        self.readbuf += self.file.read(read_amount)

    def _tell_read_pos(self):
        return self.file.tell() - len(self.readbuf)

    def _seek(self, pos):
        self.readbuf = u''
        self.file.seek(pos)

class Node:

    TYPE_NULL = 0
    TYPE_FALSE = 1
    TYPE_TRUE = 2
    TYPE_NUMBER = 3
    TYPE_STRING = 4
    TYPE_ARRAY = 5
    TYPE_OBJECT = 6

    def __init__(self, bj):
        self.bj = bj

        self.bj._skip_whitespace()

        # Read part of file to analyze type of it
        if self.bj._skip_if_next('null'):
            self.type = Node.TYPE_NULL

        elif self.bj._skip_if_next('false'):
            self.type = Node.TYPE_FALSE

        elif self.bj._skip_if_next('true'):
            self.type = Node.TYPE_TRUE

        elif self.bj._peek() in '-0123456789':
            self.type = Node.TYPE_NUMBER
            num = self.bj._get()
            # Check sign
            if num == '-':
                num += self.bj._get()
            # Read integer part
            if num[-1] != '0':
                while self.bj._peek() in '0123456789':
                    num += self.bj._get()
            # Read possible decimal part and convert to float
            if self.bj._peek() == '.':
                self.bj._get()
                num += '.' + self.bj._get()
                if num[-1] not in '0123456790':
                    raise Exception('Expected digit after dot!')
                while self.bj._peek() in '0123456789':
                    num += self.bj._get()
                self.number = float(num)
            else:
                self.number = int(num)
            # Read possible exponent
            if self.bj._peek() in 'eE':
                self.bj._get()
                exp = self.bj._get()
                exp_neg = False
                if exp == '-':
                    exp_neg = True
                    exp = self.bj._get()
                elif exp == '+':
                    exp = self.bj._get()
                while self.bj._peek() in '0123456789':
                    exp += self.bj._get()
                exp = int(exp)
                exp = 10 ** exp
                if exp_neg:
                    self.number = float(self.number) / exp
                else:
                    self.number *= exp

        elif self.bj._skip_if_next('"'):
            self.type = Node.TYPE_STRING

            self.string = u''

            while True:
                c = self.bj._get()

                if c == u'"':
                    break

                if c == u'\\':
                    c = self.bj._get()
                    if c == u'"':
                        self.string += u'"'
                    elif c == u'\\':
                        self.string += u'\\'
                    elif c == u'/':
                        self.string += u'/'
                    elif c == u'b':
                        self.string += u'\b'
                    elif c == u'f':
                        self.string += u'\f'
                    elif c == u'n':
                        self.string += u'\n'
                    elif c == u'r':
                        self.string += u'\r'
                    elif c == u't':
                        self.string += u'\t'
                    elif c == u'u':
                        unicode_bytes = self.bj._read(4)
                        self.string += ('\\u' + unicode_bytes).decode('unicode_escape')
                    else:
                        raise Exception(u'Unexpected {} in backslash encoding!'.format(c))

                else:
                    self.string += c

        elif self.bj._skip_if_next('['):
            self.type = Node.TYPE_ARRAY
            self.fully_read = False
            # TODO: Replace this single lookup position with some kind of lookup table!
            self.children_begin = self.bj._tell_read_pos()

        elif self.bj._skip_if_next('{'):
            self.type = Node.TYPE_OBJECT
            self.fully_read = False
            # TODO: Replace this single lookup position with some kind of lookup table!
            self.children_begin = self.bj._tell_read_pos()

        else:
            raise Exception('Unexpected bytes!')

    def get_type(self):
        return ['null', 'false', 'true', 'number', 'string', 'array', 'object'][self.type]

    def __getitem__(self, index):
        if self.type == Node.TYPE_ARRAY:

            # TODO: Support negative indexes!

            # Rewind to requested child
            self.bj._seek(self.children_begin)
            self.bj._skip_whitespace()
            while True:
                if self.bj._is_next(']'):
                    raise IndexError('Out of range!')

                child = Node(self.bj)
                child._read_fully()

                if index == 0:
                    return child._get_value()

                # Skip comma and whitespace around it
                self.bj._skip_whitespace()
                if not self.bj._skip_if_next(','):
                    raise IndexError('Out of range!')
                self.bj._skip_whitespace()

                index -= 1

        elif self.type == Node.TYPE_OBJECT:

            if not isinstance(index, basestring):
                raise TypeError(u'Key must be string!')

            # Rewind to requested child
            self.bj._seek(self.children_begin)
            self.bj._skip_whitespace()

            if self.bj._is_next('}'):
                raise KeyError('Key not found!')

            while True:

                # Read key
                key = Node(self.bj)
                if key.type != Node.TYPE_STRING:
                    raise Exception('Invalid key type in JSON!')

                # Read colon
                self.bj._skip_whitespace()
                if not self.bj._skip_if_next(':'):
                    raise Exception('Missing ":"!')
                self.bj._skip_whitespace()

                # Read child and return it if the key matches.
                child = Node(self.bj)
                child._read_fully()
                if key._get_value() == index:
                    return child._get_value()

                # Read comma
                self.bj._skip_whitespace()
                if self.bj._skip_if_next(','):
                    self.bj._skip_whitespace()
                elif self.bj._is_next('}'):
                    raise KeyError('Key not found!')
                else:
                    raise Exception('Expected "," or "}"!')

        elif self.type == Node.TYPE_STRING:

            return self.string[index]

        else:

            raise Exception('Index operator does not work with type "{}"!'.format(self.get_type))

    def _get_value(self):
        """ If possible, returns basic python object, but
        with complex Nodes, returns the Node itself.
        """
        if self.type == Node.TYPE_NULL:
            return None
        if self.type == Node.TYPE_FALSE:
            return False
        if self.type == Node.TYPE_TRUE:
            return True
        if self.type == Node.TYPE_NUMBER:
            return self.number
        if self.type == Node.TYPE_STRING:
            return self.string
        return self

    def _read_fully(self):

        # If already fully read
        if self.type in [Node.TYPE_NULL, Node.TYPE_FALSE, Node.TYPE_TRUE, Node.TYPE_NUMBER, Node.TYPE_STRING]:
            return

        if self.type == Node.TYPE_ARRAY:
            raise Exception('Not implemented yet!')

        else:
            raise Exception('Not implemented yet!')
<|endoftext|>


<commit_before># ----------------------------------------------------------------------------
# pymunk
# Copyright (c) 2007-2012 Victor Blomqvist
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# ----------------------------------------------------------------------------

"""This submodule contains helper functions to help with quick prototyping
using pymunk together with pyglet.

Intended to help with debugging and prototyping, not for actual production use
in a full application. The methods contained in this module is opinionated
about your coordinate system and not very optimized (they use batched
drawing, but there is probably room for optimizations still).
"""

__version__ = "$Id$"
__docformat__ = "reStructuredText"

__all__ = ["draw_space", "draw_shape", "draw_circle", "draw_poly"
         , "draw_segment"]

import math

import pyglet
from pyglet.gl import *

import pymunk
from pymunk.vec2d import Vec2d

def draw_space(space, batch = None):
    """Draw the contents of a pymunk.Space object

    This method currently supports drawing of
        * pymunk.Segment
        * pymunk.Circle
        * pymunk.Poly

    You can control the color of a shape by setting shape.color to the color
    you want it drawn in.

    >>> my_shape.color = (255, 0, 0) # will draw my_shape in red

    If you do not want a shape to be drawn, set shape.ignore_draw to True.

    >>> my_shape.ignore_draw = True

    See pyglet_util.demo.py for a full example

    :Parameters:
            space : pymunk.Space
                The contents of this Space will be drawn on the surface.
            batch : pyglet.graphics.Batch
                Use this batch to draw on. You can pass None to have it
                create a new batch. If you pass in your own batch you need
                to call draw on it yourself.
    """
    new_batch = False
    if batch == None:
        new_batch = True
        batch = pyglet.graphics.Batch()
    for s in space.shapes:
        if not (hasattr(s, "ignore_draw") and s.ignore_draw):
            draw_shape(s, batch)
    if new_batch:
        batch.draw()

def draw_shape(shape, batch = None):
    """Draw a pymunk.Shape object

    See the documentation of draw_space for full details

    :Parameters:
            surface : pygame.Surface
                Surface that the space will be drawn on
            shape : pymunk.Shape
                The Shape object to draw
    """
    if isinstance(shape, pymunk.Circle):
        draw_circle(shape, batch)
    elif isinstance(shape, pymunk.Segment):
        draw_segment(shape, batch)
    elif  isinstance(shape, pymunk.Poly):
        draw_poly(shape, batch)

def draw_circle(circle, batch = None):
    """Draw a pymunk.Circle object

    See help of draw_space for full details

    :Parameters:
            shape : pymunk.Circle
                The circle shape to draw
    """
    circle_center = circle.body.position + circle.offset.rotated(circle.body.angle)

    r = 0
    if hasattr(circle, "color"):
        color = circle.color
    elif circle.body.is_static:
        color = (200, 200, 200)
        r = 1
    else:
        color = (255, 0, 0)

    #http://slabode.exofire.net/circle_draw.shtml
    num_segments = int(4 * math.sqrt(circle.radius))
    theta = 2 * math.pi / num_segments
    c = math.cos(theta)
    s = math.sin(theta)

    x = circle.radius # we start at angle 0
    y = 0

    ps = []

    for i in range(num_segments):
        ps += [Vec2d(circle_center.x + x, circle_center.y + y)]
        t = x
        x = c * x - s * y
        y = s * t + c * y

    if circle.body.is_static:
        mode = pyglet.gl.GL_LINES
        ps = [p for p in ps+ps[:1] for _ in (0, 1)]
    else:
        mode = pyglet.gl.GL_TRIANGLE_STRIP
        ps2 = [ps[0]]
        for i in range(1, len(ps)+1/2):
            ps2.append(ps[i])
            ps2.append(ps[-i])
        ps = ps2
    vs = []
    for p in [ps[0]] + ps + [ps[-1]]:
            vs += [p.x, p.y]

    c = circle_center + Vec2d(circle.radius, 0).rotated(circle.body.angle)
    cvs = [circle_center.x, circle_center.y, c.x, c.y]

    bg = pyglet.graphics.OrderedGroup(0)

    l = len(vs)/2
    if batch == None:
        pyglet.graphics.draw(l, mode,
                            ('v2f', vs),
                            ('c3B', color*l))
        pyglet.graphics.draw(2, pyglet.gl.GL_LINES,
                            ('v2f', cvs),
                            ('c3B', (0,0,255)*2))
    else:
        batch.add(len(vs)/2, mode, bg,
                 ('v2f', vs),
                 ('c3B', color*l))
        batch.add(2, pyglet.gl.GL_LINES, None,
                 ('v2f', cvs),
                 ('c3B', (0,0,255)*2))
    return

def draw_poly(poly, batch = None):
    """Draw a pymunk.Poly object

    See help of draw_space for full details

    :Parameters:
            shape : pymunk.Poly
                The poly shape to draw
    """
    ps = poly.get_points()

    if hasattr(poly, "color"):
        color = poly.color
    elif poly.body.is_static:
        color = (200, 200, 200)
    else:
        color = (0, 255, 0)

    if poly.body.is_static:
        mode = pyglet.gl.GL_LINES
        ps = [p for p in ps+ps[:1] for _ in (0, 1)]
    else:
        mode = pyglet.gl.GL_TRIANGLE_STRIP
        ps = [ps[1],ps[2], ps[0]] + ps[3:]

    vs = []
    for p in [ps[0]] + ps + [ps[-1]]:
            vs += [p.x, p.y]

    l = len(vs)/2
    if batch == None:
        pyglet.graphics.draw(l, mode,
                            ('v2f', vs),
                            ('c3B', color*l))
    else:
        batch.add(l, mode, None,
                 ('v2f', vs),
                 ('c3B', color*l))

def draw_segment(segment, batch = None):
    """Draw a pymunk.Segment object

    See help of draw_space for full details

    :Parameters:
            shape : pymunk.Segment
                The segment shape to draw
    """
    body = segment.body
    pv1 = body.position + segment.a.rotated(body.angle)
    pv2 = body.position + segment.b.rotated(body.angle)

    d = pv2-pv1
    a = math.atan2(d.x, d.y)
    dx = segment.radius * math.cos(a)
    dy = segment.radius * math.sin(a)

    p1 = pv1 + Vec2d(dx,dy)
    p2 = pv1 - Vec2d(dx,dy)
    p3 = pv2 + Vec2d(dx,dy)
    p4 = pv2 - Vec2d(dx,dy)

    vs = [i for xy in [p1,p2,p3]+[p2,p3,p4] for i in xy]

    if hasattr(segment, "color"):
        color = segment.color
    elif segment.body.is_static:
        color = (200, 200, 200)
    else:
        color = (0, 0, 255)

    l = len(vs)/2
    if batch == None:
        pyglet.graphics.draw(l, pyglet.gl.GL_TRIANGLES,
                            ('v2f', vs),
                            ('c3B', color * l))
    else:
        batch.add(l,pyglet.gl.GL_TRIANGLES, None,
                 ('v2f', vs),
                 ('c3B', color * l))
<commit_msg>fixed segment drawing bug in pyglet util<commit_after># ----------------------------------------------------------------------------
# pymunk
# Copyright (c) 2007-2012 Victor Blomqvist
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# ----------------------------------------------------------------------------

"""This submodule contains helper functions to help with quick prototyping
using pymunk together with pyglet.

Intended to help with debugging and prototyping, not for actual production use
in a full application. The methods contained in this module is opinionated
about your coordinate system and not very optimized (they use batched
drawing, but there is probably room for optimizations still).
"""

__version__ = "$Id$"
__docformat__ = "reStructuredText"

__all__ = ["draw_space", "draw_shape", "draw_circle", "draw_poly"
         , "draw_segment"]

import math

import pyglet
from pyglet.gl import *

import pymunk
from pymunk.vec2d import Vec2d

def draw_space(space, batch = None):
    """Draw the contents of a pymunk.Space object

    This method currently supports drawing of
        * pymunk.Segment
        * pymunk.Circle
        * pymunk.Poly

    You can control the color of a shape by setting shape.color to the color
    you want it drawn in.

    >>> my_shape.color = (255, 0, 0) # will draw my_shape in red

    If you do not want a shape to be drawn, set shape.ignore_draw to True.

    >>> my_shape.ignore_draw = True

    See pyglet_util.demo.py for a full example

    :Parameters:
            space : pymunk.Space
                The contents of this Space will be drawn on the surface.
            batch : pyglet.graphics.Batch
                Use this batch to draw on. You can pass None to have it
                create a new batch. If you pass in your own batch you need
                to call draw on it yourself.
    """
    new_batch = False
    if batch == None:
        new_batch = True
        batch = pyglet.graphics.Batch()
    for s in space.shapes:
        if not (hasattr(s, "ignore_draw") and s.ignore_draw):
            draw_shape(s, batch)
    if new_batch:
        batch.draw()

def draw_shape(shape, batch = None):
    """Draw a pymunk.Shape object

    See the documentation of draw_space for full details

    :Parameters:
            surface : pygame.Surface
                Surface that the space will be drawn on
            shape : pymunk.Shape
                The Shape object to draw
    """
    if isinstance(shape, pymunk.Circle):
        draw_circle(shape, batch)
    elif isinstance(shape, pymunk.Segment):
        draw_segment(shape, batch)
    elif  isinstance(shape, pymunk.Poly):
        draw_poly(shape, batch)

def draw_circle(circle, batch = None):
    """Draw a pymunk.Circle object

    See help of draw_space for full details

    :Parameters:
            shape : pymunk.Circle
                The circle shape to draw
    """
    circle_center = circle.body.position + circle.offset.rotated(circle.body.angle)

    r = 0
    if hasattr(circle, "color"):
        color = circle.color
    elif circle.body.is_static:
        color = (200, 200, 200)
        r = 1
    else:
        color = (255, 0, 0)

    #http://slabode.exofire.net/circle_draw.shtml
    num_segments = int(4 * math.sqrt(circle.radius))
    theta = 2 * math.pi / num_segments
    c = math.cos(theta)
    s = math.sin(theta)

    x = circle.radius # we start at angle 0
    y = 0

    ps = []

    for i in range(num_segments):
        ps += [Vec2d(circle_center.x + x, circle_center.y + y)]
        t = x
        x = c * x - s * y
        y = s * t + c * y

    if circle.body.is_static:
        mode = pyglet.gl.GL_LINES
        ps = [p for p in ps+ps[:1] for _ in (0, 1)]
    else:
        mode = pyglet.gl.GL_TRIANGLE_STRIP
        ps2 = [ps[0]]
        for i in range(1, len(ps)+1/2):
            ps2.append(ps[i])
            ps2.append(ps[-i])
        ps = ps2
    vs = []
    for p in [ps[0]] + ps + [ps[-1]]:
            vs += [p.x, p.y]

    c = circle_center + Vec2d(circle.radius, 0).rotated(circle.body.angle)
    cvs = [circle_center.x, circle_center.y, c.x, c.y]

    bg = pyglet.graphics.OrderedGroup(0)

    l = len(vs)/2
    if batch == None:
        pyglet.graphics.draw(l, mode,
                            ('v2f', vs),
                            ('c3B', color*l))
        pyglet.graphics.draw(2, pyglet.gl.GL_LINES,
                            ('v2f', cvs),
                            ('c3B', (0,0,255)*2))
    else:
        batch.add(len(vs)/2, mode, bg,
                 ('v2f', vs),
                 ('c3B', color*l))
        batch.add(2, pyglet.gl.GL_LINES, None,
                 ('v2f', cvs),
                 ('c3B', (0,0,255)*2))
    return

def draw_poly(poly, batch = None):
    """Draw a pymunk.Poly object

    See help of draw_space for full details

    :Parameters:
            shape : pymunk.Poly
                The poly shape to draw
    """
    ps = poly.get_points()

    if hasattr(poly, "color"):
        color = poly.color
    elif poly.body.is_static:
        color = (200, 200, 200)
    else:
        color = (0, 255, 0)

    if poly.body.is_static:
        mode = pyglet.gl.GL_LINES
        ps = [p for p in ps+ps[:1] for _ in (0, 1)]
    else:
        mode = pyglet.gl.GL_TRIANGLE_STRIP
        ps = [ps[1],ps[2], ps[0]] + ps[3:]

    vs = []
    for p in [ps[0]] + ps + [ps[-1]]:
            vs += [p.x, p.y]

    l = len(vs)/2
    if batch == None:
        pyglet.graphics.draw(l, mode,
                            ('v2f', vs),
                            ('c3B', color*l))
    else:
        batch.add(l, mode, None,
                 ('v2f', vs),
                 ('c3B', color*l))

def draw_segment(segment, batch = None):
    """Draw a pymunk.Segment object

    See help of draw_space for full details

    :Parameters:
            shape : pymunk.Segment
                The segment shape to draw
    """
    body = segment.body
    pv1 = body.position + segment.a.rotated(body.angle)
    pv2 = body.position + segment.b.rotated(body.angle)

    d = pv2 - pv1
    a = -math.atan2(d.x, d.y)
    dx = segment.radius * math.cos(a)
    dy = segment.radius * math.sin(a)

    p1 = pv1 + Vec2d(dx,dy)
    p2 = pv1 - Vec2d(dx,dy)
    p3 = pv2 + Vec2d(dx,dy)
    p4 = pv2 - Vec2d(dx,dy)

    vs = [i for xy in [p1,p2,p3]+[p2,p3,p4] for i in xy]

    if hasattr(segment, "color"):
        color = segment.color
    elif segment.body.is_static:
        color = (200, 200, 200)
    else:
        color = (0, 0, 255)

    l = len(vs)/2
    if batch == None:
        pyglet.graphics.draw(l, pyglet.gl.GL_TRIANGLES,
                            ('v2f', vs),
                            ('c3B', color * l))
    else:
        batch.add(l,pyglet.gl.GL_TRIANGLES, None,
                 ('v2f', vs),
                 ('c3B', color * l))
<|endoftext|>


<commit_before>"""Run PyNetgear from the command-line."""
import sys
from pynetgear import Netgear

def main():
    """Scan for devices and print results."""
    netgear = Netgear(*sys.argv[1:])

    devices = netgear.get_attached_devices()

    if devices is None:
        print("Error communicating with the Netgear router")

    else:
        for i in devices:
            print(i)

if __name__ == '__main__':
    main()
<commit_msg>print get_traffic_meter in main<commit_after>"""Run PyNetgear from the command-line."""
import sys
from pynetgear import Netgear

def main():
    """Scan for devices and print results."""
    netgear = Netgear(*sys.argv[1:])

    print(netgear.get_traffic_meter())

    devices = netgear.get_attached_devices()

    if devices is None:
        print("Error communicating with the Netgear router")

    else:
        for i in devices:
            print(i)

if __name__ == '__main__':
    main()
<|endoftext|>