<commit_before>"""This module provides utility classes and functions to load spike sorting
data sets."""

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import os
import os.path
import shutil
import re
from collections import Counter

import numpy as np
import pandas as pd
import tables as tb

from loader import (Loader, default_group_info, reorder, renumber_clusters,
    default_cluster_info)
from klustersloader import (find_filenames, save_clusters, convert_to_clu,
    find_filename, find_filename_or_new)
from tools import (load_text, normalize,
    load_binary, load_pickle, save_text, get_array,
    first_row, load_binary_memmap)
from selection import (select, select_pairs, get_spikes_in_clusters,
    get_some_spikes_in_clusters, get_some_spikes, get_indices, pandaize)
from kwiklib.utils.logger import (debug, info, warn, exception, FileLogger,
    register, unregister)
from kwiklib.utils.colors import COLORS_COUNT, generate_colors
from kwiklib.dataio.kwik import add_cluster
from kwiklib.dataio.klusterskwik import klusters_to_kwik
from .experiment import Experiment

def add_missing_clusters(exp):

    shanks = sorted(exp.channel_groups.keys())

    for shank in shanks:
        cg = exp.channel_groups[shank]
        clusters = cg.clusters.main.keys()
        clusters_unique = np.unique(cg.spikes.clusters.main[:])
        # Find missing clusters in the kwik file.
        missing = sorted(set(clusters_unique)-set(clusters))

        # Add all missing clusters with a default color and "Unsorted" cluster group (group #3).
        for idx in missing:
            info("Adding missing cluster %d in shank %d." % (idx, shank))
            add_cluster(exp._files, channel_group_id='%d' % shank,
                        id=str(idx),
                        clustering='main',
                        cluster_group=3)

# -----------------------------------------------------------------------------
# HDF5 Loader
# -----------------------------------------------------------------------------
class KwikLoader(Loader):
    # TODO: change the clustering ('main' by default)

    def __init__(self, parent=None, filename=None, userpref=None):
        self.experiment = None
        super(KwikLoader, self).__init__(parent=parent, filename=filename, userpref=userpref)

    # Read functions.
    # ---------------
    def _report_progress_open(self, spike, nspikes, shank, nshanks):
        i = shank * 100 + float(spike)/nspikes*100
        n = nshanks * 100
        self.report_progress(i, n)

    def _consistency_check(self):
        exp = self.experiment
        chgrp = self.shank

        cg = exp.channel_groups[chgrp]
        clusters = cg.clusters.main.keys()
        clusters_unique = np.unique(cg.spikes.clusters.main[:])

        # Find missing clusters in the kwik file.
        missing = sorted(set(clusters_unique)-set(clusters))

        # Add all missing clusters with a default color and "Unsorted" cluster group (group #3).
        for idx in missing:
            warn("Consistency check: adding cluster %d in the kwik file" % idx)
            add_cluster(exp._files, channel_group_id='%d' % chgrp,
                        id=idx,
                        clustering='main',
                        cluster_group=3)

    def open(self, filename=None):
        """Open everything."""
        if filename is None:
            filename = self.filename
        else:
            self.filename = filename
        dir, basename = os.path.split(filename)

        # Converting to kwik if needed
        # kwik = find_filename(basename, 'kwik', dir=dir)
        # xml = find_filename(basename, 'xml', dir=dir)
        # self.filename_clu = find_filename(basename, 'clu', dir=dir)
        self._filenames = find_filenames(filename)
        kwik = find_filename(basename, 'kwik', dir=dir)
        xml = self._filenames['xml']
        clu = self._filenames['clu']

        self.log_filename = find_filename_or_new(filename, 'kvlog', dir=dir)

        # Backup the .clu file.
        clu_original = find_filename_or_new(filename, 'clu_original')
        if os.path.exists(clu) and not os.path.exists(clu_original):
            shutil.copyfile(clu, clu_original)

        if not kwik:
            assert xml, ValueError("I need a valid .kwik file")
            return

        self.experiment = Experiment(basename, dir=dir, mode='a')

        # CONSISTENCY CHECK
        # add missing clusters
        add_missing_clusters(self.experiment)

        # TODO
        # self.initialize_logfile()
        # Load the similarity measure chosen by the user in the preferences
        # file: 'gaussian' or 'kl'.
        # Refresh the preferences file when a new file is opened.
        # USERPREF.refresh()
        self.similarity_measure = self.userpref['similarity_measure'] or 'gaussian'
        debug("Similarity measure: {0:s}.".format(self.similarity_measure))
        info("Opening {0:s}.".format(self.experiment.name))
        self.shanks = sorted(self.experiment.channel_groups.keys())

        self.freq = self.experiment.application_data.spikedetekt.sample_rate

        self.fetdim = self.experiment.application_data.spikedetekt.n_features_per_channel
        self.nsamples = self.experiment.application_data.spikedetekt.extract_s_before + self.experiment.application_data.spikedetekt.extract_s_after

        self.set_shank(self.shanks[0])

    # Shank functions.
    # ----------------
    def get_shanks(self):
        """Return the list of shanks available in the file."""
        return self.shanks

    def set_shank(self, shank):
        """Change the current shank and read the corresponding tables."""
        if not shank in self.shanks:
            warn("Shank {0:d} is not in the list of shanks: {1:s}".format(
                shank, str(self.shanks)))
            return
        self.shank = shank

        # CONSISTENCY CHECK
        # self._consistency_check()

        self.nchannels = len(self.experiment.channel_groups[self.shank].channels)

        clusters = self.experiment.channel_groups[self.shank].spikes.clusters.main[:]
        self.clusters = pd.Series(clusters, dtype=np.int32)
        self.nspikes = len(self.clusters)

        self.features = self.experiment.channel_groups[self.shank].spikes.features
        self.masks = self.experiment.channel_groups[self.shank].spikes.masks
        self.waveforms = self.experiment.channel_groups[self.shank].spikes.waveforms_filtered

        if self.features is not None:
            nfet = self.features.shape[1]
            self.nextrafet = (nfet - self.nchannels * self.fetdim)
        else:
            self.nextrafet = 0

        # Load concatenated time samples: those are the time samples +
        # the start time of the corresponding recordings.
        spiketimes = self.experiment.channel_groups[self.shank].spikes.concatenated_time_samples[:] * (1. / self.freq)
        self.spiketimes = pd.Series(spiketimes, dtype=np.float64)
        self.duration = spiketimes[-1]

        self._update_data()

        self.read_clusters()

    def copy_clustering(self, clustering_from='original',
                        clustering_to='main'):
        clusters = self.experiment.channel_groups[self.shank].spikes.clusters
        clusters.copy(clustering_from, clustering_to)

    # Read contents.
    # ---------------------
    def get_probe_geometry(self):
        return np.array([c.position
            for c in self.experiment.channel_groups[self.shank].channels])

    def read_clusters(self):
        # Read the cluster info.
        clusters = self.experiment.channel_groups[self.shank].clusters.main.keys()
        cluster_groups = [c.cluster_group or 0 for c in self.experiment.channel_groups[self.shank].clusters.main.values()]

        # cluster_colors = [c.application_data.klustaviewa.color
        #     if c.application_data.klustaviewa.color is not None
        #     else 1
        #     for c in self.experiment.channel_groups[self.shank].clusters.main.values()]

        groups = self.experiment.channel_groups[self.shank].cluster_groups.main.keys()
        group_names = [g.name or 'Group' for g in self.experiment.channel_groups[self.shank].cluster_groups.main.values()]
        # group_colors = [g.application_data.klustaviewa.color or 1 for g in self.experiment.channel_groups[self.shank].cluster_groups.main.values()]

        # Create the cluster_info DataFrame.
        self.cluster_info = pd.DataFrame(dict(
            # color=cluster_colors,
            group=cluster_groups,
            ), index=clusters)
        # self.cluster_colors = self.cluster_info['color'].astype(np.int32)
        self.cluster_groups = self.cluster_info['group'].astype(np.int32)

        # Create the group_info DataFrame.
        self.group_info = pd.DataFrame(dict(
            # color=group_colors,
            name=group_names,
            ), index=groups)
        # self.group_colors = self.group_info['color'].astype(np.int32)
        self.group_names = self.group_info['name']

    # Writing capabilities.
    # ---------------------
    def set_cluster(self, spikes, cluster):
        if not hasattr(spikes, '__len__'):
            spikes = [spikes]

        self.experiment.channel_groups[self.shank].spikes.clusters.main[spikes] = cluster
        clusters = self.experiment.channel_groups[self.shank].spikes.clusters.main[:]
        self.clusters = pd.Series(clusters, dtype=np.int32)

        self._update_data()

    def set_cluster_groups(self, clusters, group):
        # self.cluster_groups.ix[clusters] = group
        if not hasattr(clusters, '__len__'):
            clusters = [clusters]

        clusters_gr = self.experiment.channel_groups[self.shank].clusters.main
        for cl in clusters:
            clusters_gr[cl].cluster_group = group

        self.read_clusters()

    def set_cluster_colors(self, clusters, color):
        # self.cluster_colors.ix[clusters] = color
        if not hasattr(clusters, '__len__'):
            clusters = [clusters]
        clusters_gr = self.experiment.channel_groups[self.shank].clusters.main
        for cl in clusters:
            clusters_gr[cl].application_data.klustaviewa.color = color

        self.read_clusters()

    def set_group_names(self, groups, name):
        # self.group_names.ix[groups] = name
        if not hasattr(groups, '__len__'):
            groups = [groups]
        groups_gr = self.experiment.channel_groups[self.shank].cluster_groups.main
        for gr in groups:
            groups_gr[gr].name = name

        self.read_clusters()

    def set_group_colors(self, groups, color):
        # self.group_colors.ix[groups] = color
        if not hasattr(groups, '__len__'):
            groups = [groups]

        groups_gr = self.experiment.channel_groups[self.shank].cluster_groups.main
        # for gr in groups:
        #     groups_gr[gr].application_data.klustaviewa.color = color

        self.read_clusters()

    # Add.
    def add_cluster(self, cluster, group, color):
        # if cluster not in self.cluster_groups.index:
            # self.cluster_groups = self.cluster_groups.append(
                # pd.Series([group], index=[cluster])).sort_index()
        # if cluster not in self.cluster_colors.index:
            # self.cluster_colors = self.cluster_colors.append(
                # pd.Series([color], index=[cluster])).sort_index()

        self.experiment.channel_groups[self.shank].clusters.main.add_cluster(
            id=cluster,
            # color=color,
            cluster_group=group)

        self.read_clusters()

    def add_clusters(self, clusters, groups):
        # if cluster not in self.cluster_groups.index:
            # self.cluster_groups = self.cluster_groups.append(
                # pd.Series([group], index=[cluster])).sort_index()
        # if cluster not in self.cluster_colors.index:
            # self.cluster_colors = self.cluster_colors.append(
                # pd.Series([color], index=[cluster])).sort_index()
        for cluster, group in zip(clusters, groups):
            self.experiment.channel_groups[self.shank].clusters.main.add_cluster(
                id=cluster, cluster_group=group)
        self.read_clusters()

    def add_group(self, group, name):
        # if group not in self.group_colors.index:
            # self.group_colors = self.group_colors.append(
                # pd.Series([color], index=[group])).sort_index()
        # if group not in self.group_names.index:
            # self.group_names = self.group_names.append(
                # pd.Series([name], index=[group])).sort_index()

        groups = self.experiment.channel_groups[self.shank].cluster_groups.main
        groups.add_group(id=group, name=name,)

        self.read_clusters()

    # Remove.
    def remove_cluster(self, cluster):
        if np.any(np.in1d(cluster, self.clusters)):
            raise ValueError(("Cluster {0:d} is not empty and cannot "
            "be removed.").format(cluster))

        self.experiment.channel_groups[self.shank].clusters.main.remove_cluster(
            id=cluster,)

        self.read_clusters()

    def remove_group(self, group):
        if np.any(np.in1d(group, self.cluster_groups)):
            raise ValueError(("Group {0:d} is not empty and cannot "
            "be removed.").format(group))

        self.experiment.channel_groups[self.shank].cluster_groups.main.remove_group(
            id=group,)

        self.read_clusters()

    # Access to the data: spikes
    # --------------------------
    def select(self, spikes=None, clusters=None):
        if clusters is not None:
            if not hasattr(clusters, '__len__'):
                clusters = [clusters]
            spikes = get_spikes_in_clusters(clusters, self.clusters)
        self.spikes_selected = spikes
        self.clusters_selected = clusters

    # Log file.
    # ---------
    def initialize_logfile(self):
        self.logfile = FileLogger(self.filename_log, name='datafile',
            level=self.userpref['loglevel_file'])
        # Register log file.
        register(self.logfile)

    # Save.
    # -----
    def save(self, renumber=False):
        self.report_progress_save(1, 4)

        if renumber:
            self.renumber()
            self.clusters = self.clusters_renumbered
            self.cluster_info = self.cluster_info_renumbered
            self._update_data()

        # Save the clusters in the .clu file.
        clu = self._filenames['clu']
        clu_split = clu.split('.')
        clu_split[-1] = str(self.shank)
        clu = '.'.join(clu_split)
        save_clusters(clu,
            convert_to_clu(self.clusters, self.cluster_info['group']))

        self.report_progress_save(2, 4)

        # self.close()
        self.report_progress_save(3, 4)

        # self.open()
        self.report_progress_save(4, 4)

    # Close functions.
    # ----------------
    def close(self):
        """Close the kwik HDF5 file."""
        # if hasattr(self, 'kwik') and self.kwik.isopen:
            # self.kwik.flush()
            # self.kwik.close()
        if self.experiment is not None:
            self.experiment.close()
            self.experiment = None
        if hasattr(self, 'logfile'):
            unregister(self.logfile)

<commit_msg>Add shank option when opening a dataset<commit_after>"""This module provides utility classes and functions to load spike sorting
data sets."""

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import os
import os.path
import shutil
import re
from collections import Counter

import numpy as np
import pandas as pd
import tables as tb

from loader import (Loader, default_group_info, reorder, renumber_clusters,
    default_cluster_info)
from klustersloader import (find_filenames, save_clusters, convert_to_clu,
    find_filename, find_filename_or_new)
from tools import (load_text, normalize,
    load_binary, load_pickle, save_text, get_array,
    first_row, load_binary_memmap)
from selection import (select, select_pairs, get_spikes_in_clusters,
    get_some_spikes_in_clusters, get_some_spikes, get_indices, pandaize)
from kwiklib.utils.logger import (debug, info, warn, exception, FileLogger,
    register, unregister)
from kwiklib.utils.colors import COLORS_COUNT, generate_colors
from kwiklib.dataio.kwik import add_cluster
from kwiklib.dataio.klusterskwik import klusters_to_kwik
from .experiment import Experiment

def add_missing_clusters(exp):

    shanks = sorted(exp.channel_groups.keys())

    for shank in shanks:
        cg = exp.channel_groups[shank]
        clusters = cg.clusters.main.keys()
        clusters_unique = np.unique(cg.spikes.clusters.main[:])
        # Find missing clusters in the kwik file.
        missing = sorted(set(clusters_unique)-set(clusters))

        # Add all missing clusters with a default color and "Unsorted" cluster group (group #3).
        for idx in missing:
            info("Adding missing cluster %d in shank %d." % (idx, shank))
            add_cluster(exp._files, channel_group_id='%d' % shank,
                        id=str(idx),
                        clustering='main',
                        cluster_group=3)

# -----------------------------------------------------------------------------
# HDF5 Loader
# -----------------------------------------------------------------------------
class KwikLoader(Loader):
    # TODO: change the clustering ('main' by default)

    def __init__(self, parent=None, filename=None, userpref=None):
        self.experiment = None
        super(KwikLoader, self).__init__(parent=parent, filename=filename, userpref=userpref)

    # Read functions.
    # ---------------
    def _report_progress_open(self, spike, nspikes, shank, nshanks):
        i = shank * 100 + float(spike)/nspikes*100
        n = nshanks * 100
        self.report_progress(i, n)

    def _consistency_check(self):
        exp = self.experiment
        chgrp = self.shank

        cg = exp.channel_groups[chgrp]
        clusters = cg.clusters.main.keys()
        clusters_unique = np.unique(cg.spikes.clusters.main[:])

        # Find missing clusters in the kwik file.
        missing = sorted(set(clusters_unique)-set(clusters))

        # Add all missing clusters with a default color and "Unsorted" cluster group (group #3).
        for idx in missing:
            warn("Consistency check: adding cluster %d in the kwik file" % idx)
            add_cluster(exp._files, channel_group_id='%d' % chgrp,
                        id=idx,
                        clustering='main',
                        cluster_group=3)

    def open(self, filename=None, shank=None):
        """Open everything."""
        if filename is None:
            filename = self.filename
        else:
            self.filename = filename
        dir, basename = os.path.split(filename)

        # Converting to kwik if needed
        # kwik = find_filename(basename, 'kwik', dir=dir)
        # xml = find_filename(basename, 'xml', dir=dir)
        # self.filename_clu = find_filename(basename, 'clu', dir=dir)
        self._filenames = find_filenames(filename)
        kwik = find_filename(basename, 'kwik', dir=dir)
        xml = self._filenames['xml']
        clu = self._filenames['clu']

        self.log_filename = find_filename_or_new(filename, 'kvlog', dir=dir)

        # Backup the .clu file.
        clu_original = find_filename_or_new(filename, 'clu_original')
        if os.path.exists(clu) and not os.path.exists(clu_original):
            shutil.copyfile(clu, clu_original)

        if not kwik:
            assert xml, ValueError("I need a valid .kwik file")
            return

        self.experiment = Experiment(basename, dir=dir, mode='a')

        # CONSISTENCY CHECK
        # add missing clusters
        add_missing_clusters(self.experiment)

        # TODO
        # self.initialize_logfile()
        # Load the similarity measure chosen by the user in the preferences
        # file: 'gaussian' or 'kl'.
        # Refresh the preferences file when a new file is opened.
        # USERPREF.refresh()
        self.similarity_measure = self.userpref['similarity_measure'] or 'gaussian'
        debug("Similarity measure: {0:s}.".format(self.similarity_measure))
        info("Opening {0:s}.".format(self.experiment.name))
        self.shanks = sorted(self.experiment.channel_groups.keys())

        self.freq = self.experiment.application_data.spikedetekt.sample_rate

        self.fetdim = self.experiment.application_data.spikedetekt.n_features_per_channel
        self.nsamples = self.experiment.application_data.spikedetekt.extract_s_before + self.experiment.application_data.spikedetekt.extract_s_after

        self.set_shank(shank or self.shanks[0])

    # Shank functions.
    # ----------------
    def get_shanks(self):
        """Return the list of shanks available in the file."""
        return self.shanks

    def set_shank(self, shank):
        """Change the current shank and read the corresponding tables."""
        if not shank in self.shanks:
            warn("Shank {0:d} is not in the list of shanks: {1:s}".format(
                shank, str(self.shanks)))
            return
        self.shank = shank

        # CONSISTENCY CHECK
        # self._consistency_check()

        self.nchannels = len(self.experiment.channel_groups[self.shank].channels)

        clusters = self.experiment.channel_groups[self.shank].spikes.clusters.main[:]
        self.clusters = pd.Series(clusters, dtype=np.int32)
        self.nspikes = len(self.clusters)

        self.features = self.experiment.channel_groups[self.shank].spikes.features
        self.masks = self.experiment.channel_groups[self.shank].spikes.masks
        self.waveforms = self.experiment.channel_groups[self.shank].spikes.waveforms_filtered

        if self.features is not None:
            nfet = self.features.shape[1]
            self.nextrafet = (nfet - self.nchannels * self.fetdim)
        else:
            self.nextrafet = 0

        # Load concatenated time samples: those are the time samples +
        # the start time of the corresponding recordings.
        spiketimes = self.experiment.channel_groups[self.shank].spikes.concatenated_time_samples[:] * (1. / self.freq)
        self.spiketimes = pd.Series(spiketimes, dtype=np.float64)
        self.duration = spiketimes[-1]

        self._update_data()

        self.read_clusters()

    def copy_clustering(self, clustering_from='original',
                        clustering_to='main'):
        clusters = self.experiment.channel_groups[self.shank].spikes.clusters
        clusters.copy(clustering_from, clustering_to)

    # Read contents.
    # ---------------------
    def get_probe_geometry(self):
        return np.array([c.position
            for c in self.experiment.channel_groups[self.shank].channels])

    def read_clusters(self):
        # Read the cluster info.
        clusters = self.experiment.channel_groups[self.shank].clusters.main.keys()
        cluster_groups = [c.cluster_group or 0 for c in self.experiment.channel_groups[self.shank].clusters.main.values()]

        # cluster_colors = [c.application_data.klustaviewa.color
        #     if c.application_data.klustaviewa.color is not None
        #     else 1
        #     for c in self.experiment.channel_groups[self.shank].clusters.main.values()]

        groups = self.experiment.channel_groups[self.shank].cluster_groups.main.keys()
        group_names = [g.name or 'Group' for g in self.experiment.channel_groups[self.shank].cluster_groups.main.values()]
        # group_colors = [g.application_data.klustaviewa.color or 1 for g in self.experiment.channel_groups[self.shank].cluster_groups.main.values()]

        # Create the cluster_info DataFrame.
        self.cluster_info = pd.DataFrame(dict(
            # color=cluster_colors,
            group=cluster_groups,
            ), index=clusters)
        # self.cluster_colors = self.cluster_info['color'].astype(np.int32)
        self.cluster_groups = self.cluster_info['group'].astype(np.int32)

        # Create the group_info DataFrame.
        self.group_info = pd.DataFrame(dict(
            # color=group_colors,
            name=group_names,
            ), index=groups)
        # self.group_colors = self.group_info['color'].astype(np.int32)
        self.group_names = self.group_info['name']

    # Writing capabilities.
    # ---------------------
    def set_cluster(self, spikes, cluster):
        if not hasattr(spikes, '__len__'):
            spikes = [spikes]

        self.experiment.channel_groups[self.shank].spikes.clusters.main[spikes] = cluster
        clusters = self.experiment.channel_groups[self.shank].spikes.clusters.main[:]
        self.clusters = pd.Series(clusters, dtype=np.int32)

        self._update_data()

    def set_cluster_groups(self, clusters, group):
        # self.cluster_groups.ix[clusters] = group
        if not hasattr(clusters, '__len__'):
            clusters = [clusters]

        clusters_gr = self.experiment.channel_groups[self.shank].clusters.main
        for cl in clusters:
            clusters_gr[cl].cluster_group = group

        self.read_clusters()

    def set_cluster_colors(self, clusters, color):
        # self.cluster_colors.ix[clusters] = color
        if not hasattr(clusters, '__len__'):
            clusters = [clusters]
        clusters_gr = self.experiment.channel_groups[self.shank].clusters.main
        for cl in clusters:
            clusters_gr[cl].application_data.klustaviewa.color = color

        self.read_clusters()

    def set_group_names(self, groups, name):
        # self.group_names.ix[groups] = name
        if not hasattr(groups, '__len__'):
            groups = [groups]
        groups_gr = self.experiment.channel_groups[self.shank].cluster_groups.main
        for gr in groups:
            groups_gr[gr].name = name

        self.read_clusters()

    def set_group_colors(self, groups, color):
        # self.group_colors.ix[groups] = color
        if not hasattr(groups, '__len__'):
            groups = [groups]

        groups_gr = self.experiment.channel_groups[self.shank].cluster_groups.main
        # for gr in groups:
        #     groups_gr[gr].application_data.klustaviewa.color = color

        self.read_clusters()

    # Add.
    def add_cluster(self, cluster, group, color):
        # if cluster not in self.cluster_groups.index:
            # self.cluster_groups = self.cluster_groups.append(
                # pd.Series([group], index=[cluster])).sort_index()
        # if cluster not in self.cluster_colors.index:
            # self.cluster_colors = self.cluster_colors.append(
                # pd.Series([color], index=[cluster])).sort_index()

        self.experiment.channel_groups[self.shank].clusters.main.add_cluster(
            id=cluster,
            # color=color,
            cluster_group=group)

        self.read_clusters()

    def add_clusters(self, clusters, groups):
        # if cluster not in self.cluster_groups.index:
            # self.cluster_groups = self.cluster_groups.append(
                # pd.Series([group], index=[cluster])).sort_index()
        # if cluster not in self.cluster_colors.index:
            # self.cluster_colors = self.cluster_colors.append(
                # pd.Series([color], index=[cluster])).sort_index()
        for cluster, group in zip(clusters, groups):
            self.experiment.channel_groups[self.shank].clusters.main.add_cluster(
                id=cluster, cluster_group=group)
        self.read_clusters()

    def add_group(self, group, name):
        # if group not in self.group_colors.index:
            # self.group_colors = self.group_colors.append(
                # pd.Series([color], index=[group])).sort_index()
        # if group not in self.group_names.index:
            # self.group_names = self.group_names.append(
                # pd.Series([name], index=[group])).sort_index()

        groups = self.experiment.channel_groups[self.shank].cluster_groups.main
        groups.add_group(id=group, name=name,)

        self.read_clusters()

    # Remove.
    def remove_cluster(self, cluster):
        if np.any(np.in1d(cluster, self.clusters)):
            raise ValueError(("Cluster {0:d} is not empty and cannot "
            "be removed.").format(cluster))

        self.experiment.channel_groups[self.shank].clusters.main.remove_cluster(
            id=cluster,)

        self.read_clusters()

    def remove_group(self, group):
        if np.any(np.in1d(group, self.cluster_groups)):
            raise ValueError(("Group {0:d} is not empty and cannot "
            "be removed.").format(group))

        self.experiment.channel_groups[self.shank].cluster_groups.main.remove_group(
            id=group,)

        self.read_clusters()

    # Access to the data: spikes
    # --------------------------
    def select(self, spikes=None, clusters=None):
        if clusters is not None:
            if not hasattr(clusters, '__len__'):
                clusters = [clusters]
            spikes = get_spikes_in_clusters(clusters, self.clusters)
        self.spikes_selected = spikes
        self.clusters_selected = clusters

    # Log file.
    # ---------
    def initialize_logfile(self):
        self.logfile = FileLogger(self.filename_log, name='datafile',
            level=self.userpref['loglevel_file'])
        # Register log file.
        register(self.logfile)

    # Save.
    # -----
    def save(self, renumber=False):
        self.report_progress_save(1, 4)

        if renumber:
            self.renumber()
            self.clusters = self.clusters_renumbered
            self.cluster_info = self.cluster_info_renumbered
            self._update_data()

        # Save the clusters in the .clu file.
        clu = self._filenames['clu']
        clu_split = clu.split('.')
        clu_split[-1] = str(self.shank)
        clu = '.'.join(clu_split)
        save_clusters(clu,
            convert_to_clu(self.clusters, self.cluster_info['group']))

        self.report_progress_save(2, 4)

        # self.close()
        self.report_progress_save(3, 4)

        # self.open()
        self.report_progress_save(4, 4)

    # Close functions.
    # ----------------
    def close(self):
        """Close the kwik HDF5 file."""
        # if hasattr(self, 'kwik') and self.kwik.isopen:
            # self.kwik.flush()
            # self.kwik.close()
        if self.experiment is not None:
            self.experiment.close()
            self.experiment = None
        if hasattr(self, 'logfile'):
            unregister(self.logfile)

<|endoftext|>


<commit_before>import asyncore
import json
import socket

class TestClient(asyncore.dispatcher_with_send):
	def __init__(self, host):
		asyncore.dispatcher_with_send.__init__(self)
		self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
		self.connect( (host, 8900) )
		self.data =  json.dumps({ "action": "register", "name" : "xmpp"})

	def writable(self):
		return len(self.data) > 0

	def handle_close(self):
		self.close()

	def handle_read(self):
		msg = self.recv(1024)
		print msg

	def handle_write(self):
		self.send(self.data)
		self.data = ""

test = TestClient("localhost")
asyncore.loop(1)<commit_msg>cleaning<commit_after><|endoftext|>


<commit_before># encoding:UTF-8
import logging

from chariot.influx import influx
from datetime import datetime
from django.core.urlresolvers import reverse
from django.http import JsonResponse, HttpResponse, StreamingHttpResponse
from django.shortcuts import get_object_or_404
import json
from django.views.generic import CreateView
from rest_framework import status, serializers, authentication
from rest_framework.authtoken.models import Token
from rest_framework.generics import ListAPIView
from rest_framework.response import Response
from rest_framework.views import APIView

from chariot import utils
from chariot.mixins import LoginRequiredMixin, BackButtonMixin
from chariot.utils import LOCALE_DATE_FMT
from deployments.models import Deployment, DeploymentSensor, DeploymentThermostatSetting
from graphs.views import generate_data, last
from hubs.forms import HubCreateForm
from hubs.models import Hub
from sensors.models import Sensor
from sensors.serializers import SensorIDSerializer

logger = logging.getLogger(__name__)

class DeploymentSensorIDSerializer(serializers.ModelSerializer):
    sensor = SensorIDSerializer()

    class Meta:
        model = DeploymentSensor
        fields = ('sensor', 'cost', 'location', 'nearest_thermostat', 'room_height', 'room_area')

class DeploymentThermostatSerializer(serializers.ModelSerializer):
    time = serializers.TimeField(format='%H:%M')

    class Meta:
        model = DeploymentThermostatSetting
        fields = ('setting', 'time', 'days')

class DeploymentSerializer(serializers.ModelSerializer):
    hub = serializers.PrimaryKeyRelatedField(read_only=True)
    sensors = DeploymentSensorIDSerializer(many=True)
    thermostats = DeploymentThermostatSerializer(many=True)
    boiler_type = serializers.CharField(source='get_boiler_type_display')

    class Meta:
        model = Deployment
        fields = ('id', 'hub', 'sensors', 'thermostats',
                  'boiler_manufacturer', 'boiler_model', 'boiler_output', 'boiler_efficiency',
                  'boiler_type', 'building_area', 'building_height')

class DeploymentListView(ListAPIView):
    model = Deployment
    serializer_class = DeploymentSerializer
    authentication_classes = (authentication.TokenAuthentication,)

    def get_queryset(self):
        return Deployment.objects.filter(end_date__isnull=True)

class DeploymentView(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def get_object(self, pk):
        try:
            return Deployment.objects.get(pk=pk)
        except Deployment.DoesNotExist:
            raise HttpResponse(status=404)

    def get(self, request, pk, format=None):
        deployment = self.get_object(pk)
        serializer = DeploymentSerializer(deployment)
        return Response(serializer.data)

def get_token(request, id):
    try:
        hub = Hub.objects.get(id=id)
        token = Token.objects.latest('created')
        return JsonResponse({'token': token.key, 'deployment': hub.deployment.pk})
    except Hub.DoesNotExist:
        return HttpResponse(status=404)
    except Token.DoesNotExist:
        return HttpResponse(status=404)

class HubPingView(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def put(self, request, id, format=None):
        mac_address = utils.decode_mac_address(id)

        hub = get_object_or_404(Hub, id=mac_address)

        # May raise a permission denied
        self.check_object_permissions(self.request, hub)

        hub.ping()
        hub.save()

        return Response("success", status=status.HTTP_200_OK)

class SensorListView(ListAPIView):
    model = Sensor
    serializer_class = SensorIDSerializer
    authentication_classes = (authentication.TokenAuthentication,)

    def get_queryset(self):
        mac_address = utils.decode_mac_address(self.kwargs.get('mac_address'))

        hub = get_object_or_404(Hub, id=mac_address)

        return hub.sensors

class SensorReading(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def post(self, request):
        hub_mac_address = utils.decode_mac_address(request.data['hub'])
        hub = Hub.objects.get(id=hub_mac_address)

        reading = {
            "measurement": request.data['channel'],
            "tags": {
                "sensor": request.data['sensor'],
                "deployment": hub.deployment.pk
            },
            "fields": {
                "value": float(request.data['value'])
            }
        }
        logger.info(reading)
        if 'timestamp' in request.data:
            reading['time'] = request.data['timestamp']
        if influx.write_points([reading]):
            return Response(reading, status=status.HTTP_201_CREATED)
        return Response("failed", status=status.HTTP_400_BAD_REQUEST)

class DeploymentPrediction(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def put(self, request, pk):
        prediction_json = json.loads(request.body)
        for prediction in prediction_json:
            try:
                deployment_id = prediction['Deployment_Id']
                deployment = Deployment.objects.get(pk=deployment_id)
                deployment.prediction = json.dumps(prediction)
                deployment.save()
            except Deployment.DoesNotExist:
                return HttpResponse(status=404)

        return Response("success", status=status.HTTP_200_OK)

class HubCreateView(LoginRequiredMixin, BackButtonMixin, CreateView):
    form_class = HubCreateForm
    model = Hub
    template_name = 'sensors/device_create.html'

    def get_success_url(self):
        return reverse('devices')

    def get_back_url(self):
        return reverse('devices')

class LatestDataView(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def get(self, request, pk, format=None):
        try:
            deployment = Deployment.objects.get(pk=pk)
        except Deployment.DoesNotExist:
            return HttpResponse(status=404)
        result = []

        for sensor in deployment.sensors.all():
            sensor_obj = {'id': sensor.sensor.id, 'channels': []}
            for channel in sensor.sensor.channels.all():
                value = last(deployment, sensor.sensor, channel)
                if value.has_data():
                    data = value.data
                    if data:
                        channel_obj = {'id': channel.id,
                                       'value': data[0]['last'],
                                       'time': data[0]['time']}
                        sensor_obj['channels'].append(channel_obj)
            if sensor_obj['channels']:
                result.append(sensor_obj)

        return JsonResponse(result, safe=False)

class DataView(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def get(self, request, pk, format=None):
        aggregate = None
        sensors = None
        channels = None
        start = None
        end = None

        if 'sensors' in request.GET:
            sensors = request.GET['sensors'].split(",")

        if 'channels' in request.GET:
            if request.GET['channels'].lower() == 'all':
                channels = 'all'
            else:
                channels = request.GET['channels'].split(",")

        if 'start' in request.GET:
            start = datetime.strptime(request.GET['start'], LOCALE_DATE_FMT)

        if 'end' in request.GET:
            end = datetime.strptime(request.GET['end'], LOCALE_DATE_FMT)

        simplify = 'simplify' not in request.GET or request.GET['simplify'].lower() != 'false'
        if 'aggregate' in request.GET:
            aggregate = request.GET['aggregate']

        return StreamingHttpResponse(
            generate_data(pk, sensors, channels, simplify, start, end, aggregate),
            content_type="application/json")
<commit_msg>Assume data is per-deployment<commit_after># encoding:UTF-8
import logging

from chariot.influx import influx
from datetime import datetime
from django.core.urlresolvers import reverse
from django.http import JsonResponse, HttpResponse, StreamingHttpResponse
from django.shortcuts import get_object_or_404
import json
from django.views.generic import CreateView
from rest_framework import status, serializers, authentication
from rest_framework.authtoken.models import Token
from rest_framework.generics import ListAPIView
from rest_framework.response import Response
from rest_framework.views import APIView

from chariot import utils
from chariot.mixins import LoginRequiredMixin, BackButtonMixin
from chariot.utils import LOCALE_DATE_FMT
from deployments.models import Deployment, DeploymentSensor, DeploymentThermostatSetting
from graphs.views import generate_data, last
from hubs.forms import HubCreateForm
from hubs.models import Hub
from sensors.models import Sensor
from sensors.serializers import SensorIDSerializer

logger = logging.getLogger(__name__)

class DeploymentSensorIDSerializer(serializers.ModelSerializer):
    sensor = SensorIDSerializer()

    class Meta:
        model = DeploymentSensor
        fields = ('sensor', 'cost', 'location', 'nearest_thermostat', 'room_height', 'room_area')

class DeploymentThermostatSerializer(serializers.ModelSerializer):
    time = serializers.TimeField(format='%H:%M')

    class Meta:
        model = DeploymentThermostatSetting
        fields = ('setting', 'time', 'days')

class DeploymentSerializer(serializers.ModelSerializer):
    hub = serializers.PrimaryKeyRelatedField(read_only=True)
    sensors = DeploymentSensorIDSerializer(many=True)
    thermostats = DeploymentThermostatSerializer(many=True)
    boiler_type = serializers.CharField(source='get_boiler_type_display')

    class Meta:
        model = Deployment
        fields = ('id', 'hub', 'sensors', 'thermostats',
                  'boiler_manufacturer', 'boiler_model', 'boiler_output', 'boiler_efficiency',
                  'boiler_type', 'building_area', 'building_height')

class DeploymentListView(ListAPIView):
    model = Deployment
    serializer_class = DeploymentSerializer
    authentication_classes = (authentication.TokenAuthentication,)

    def get_queryset(self):
        return Deployment.objects.filter(end_date__isnull=True)

class DeploymentView(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def get_object(self, pk):
        try:
            return Deployment.objects.get(pk=pk)
        except Deployment.DoesNotExist:
            raise HttpResponse(status=404)

    def get(self, request, pk, format=None):
        deployment = self.get_object(pk)
        serializer = DeploymentSerializer(deployment)
        return Response(serializer.data)

def get_token(request, id):
    try:
        hub = Hub.objects.get(id=id)
        token = Token.objects.latest('created')
        return JsonResponse({'token': token.key, 'deployment': hub.deployment.pk})
    except Hub.DoesNotExist:
        return HttpResponse(status=404)
    except Token.DoesNotExist:
        return HttpResponse(status=404)

class HubPingView(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def put(self, request, id, format=None):
        mac_address = utils.decode_mac_address(id)

        hub = get_object_or_404(Hub, id=mac_address)

        # May raise a permission denied
        self.check_object_permissions(self.request, hub)

        hub.ping()
        hub.save()

        return Response("success", status=status.HTTP_200_OK)

class SensorListView(ListAPIView):
    model = Sensor
    serializer_class = SensorIDSerializer
    authentication_classes = (authentication.TokenAuthentication,)

    def get_queryset(self):
        mac_address = utils.decode_mac_address(self.kwargs.get('mac_address'))

        hub = get_object_or_404(Hub, id=mac_address)

        return hub.sensors

class SensorReading(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def post(self, request):
        hub_mac_address = utils.decode_mac_address(request.data['hub'])
        hub = Hub.objects.get(id=hub_mac_address)

        reading = {
            "measurement": request.data['channel'],
            "tags": {
                "sensor": request.data['sensor'],
                "deployment": hub.deployment.pk
            },
            "fields": {
                "value": float(request.data['value'])
            }
        }
        logger.info(reading)
        if 'timestamp' in request.data:
            reading['time'] = request.data['timestamp']
        if influx.write_points([reading]):
            return Response(reading, status=status.HTTP_201_CREATED)
        return Response("failed", status=status.HTTP_400_BAD_REQUEST)

class DeploymentPrediction(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def put(self, request, pk):
        prediction = json.loads(request.body)
        try:
            deployment = Deployment.objects.get(pk=pk)
            deployment.prediction = prediction
            deployment.save()
        except Deployment.DoesNotExist:
            return HttpResponse(status=404)

        return Response("success", status=status.HTTP_200_OK)

class HubCreateView(LoginRequiredMixin, BackButtonMixin, CreateView):
    form_class = HubCreateForm
    model = Hub
    template_name = 'sensors/device_create.html'

    def get_success_url(self):
        return reverse('devices')

    def get_back_url(self):
        return reverse('devices')

class LatestDataView(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def get(self, request, pk, format=None):
        try:
            deployment = Deployment.objects.get(pk=pk)
        except Deployment.DoesNotExist:
            return HttpResponse(status=404)
        result = []

        for sensor in deployment.sensors.all():
            sensor_obj = {'id': sensor.sensor.id, 'channels': []}
            for channel in sensor.sensor.channels.all():
                value = last(deployment, sensor.sensor, channel)
                if value.has_data():
                    data = value.data
                    if data:
                        channel_obj = {'id': channel.id,
                                       'value': data[0]['last'],
                                       'time': data[0]['time']}
                        sensor_obj['channels'].append(channel_obj)
            if sensor_obj['channels']:
                result.append(sensor_obj)

        return JsonResponse(result, safe=False)

class DataView(APIView):
    authentication_classes = (authentication.TokenAuthentication,)

    def get(self, request, pk, format=None):
        aggregate = None
        sensors = None
        channels = None
        start = None
        end = None

        if 'sensors' in request.GET:
            sensors = request.GET['sensors'].split(",")

        if 'channels' in request.GET:
            if request.GET['channels'].lower() == 'all':
                channels = 'all'
            else:
                channels = request.GET['channels'].split(",")

        if 'start' in request.GET:
            start = datetime.strptime(request.GET['start'], LOCALE_DATE_FMT)

        if 'end' in request.GET:
            end = datetime.strptime(request.GET['end'], LOCALE_DATE_FMT)

        simplify = 'simplify' not in request.GET or request.GET['simplify'].lower() != 'false'
        if 'aggregate' in request.GET:
            aggregate = request.GET['aggregate']

        return StreamingHttpResponse(
            generate_data(pk, sensors, channels, simplify, start, end, aggregate),
            content_type="application/json")
<|endoftext|>


<commit_before># Copyright 2019 KMEE INFORMATICA LTDA
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl).

{
    'name': 'NFS-e',
    'summary': """
        NFS-e""",
    'version': '12.0.1.1.0',
    'license': 'AGPL-3',
    'author': 'KMEE, Odoo Community Association (OCA)',
    'maintainers': ['gabrielcardoso21', 'mileo', 'luismalta'],
    'website': 'https://github.com/OCA/l10n-brazil',
    'external_dependencies': {
        'python': [
            'erpbrasil.edoc',
            'erpbrasil.assinatura',
            'erpbrasil.transmissao',
            'erpbrasil.base',
        ],
    },
    'depends': [
        'l10n_br_fiscal',
    ],
    'data': [
        'security/ir.model.access.csv',

        'views/document_view.xml',
        'views/product_template_view.xml',
        'views/document_line_view.xml',
        'views/res_company_view.xml',

        'wizards/wizard_document_status.xml',

        'report/danfse.xml',
    ],
    'demo': [
        'demo/product_demo.xml',
        'demo/fiscal_document_demo.xml',
    ],
}
<commit_msg>l10n_br_nfse 12.0.1.2.0<commit_after># Copyright 2019 KMEE INFORMATICA LTDA
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl).

{
    'name': 'NFS-e',
    'summary': """
        NFS-e""",
    'version': '12.0.1.2.0',
    'license': 'AGPL-3',
    'author': 'KMEE, Odoo Community Association (OCA)',
    'maintainers': ['gabrielcardoso21', 'mileo', 'luismalta'],
    'website': 'https://github.com/OCA/l10n-brazil',
    'external_dependencies': {
        'python': [
            'erpbrasil.edoc',
            'erpbrasil.assinatura',
            'erpbrasil.transmissao',
            'erpbrasil.base',
        ],
    },
    'depends': [
        'l10n_br_fiscal',
    ],
    'data': [
        'security/ir.model.access.csv',

        'views/document_view.xml',
        'views/product_template_view.xml',
        'views/document_line_view.xml',
        'views/res_company_view.xml',

        'wizards/wizard_document_status.xml',

        'report/danfse.xml',
    ],
    'demo': [
        'demo/product_demo.xml',
        'demo/fiscal_document_demo.xml',
    ],
}
<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
###############################################################################
#                                                                             #
# Copyright (C) 2009  Renato Lima - Akretion, Gabriel C. Stabel               #
# Copyright (C) 2012  Raphaël Valyi - Akretion                                #
#                                                                             #
#This program is free software: you can redistribute it and/or modify         #
#it under the terms of the GNU Affero General Public License as published by  #
#the Free Software Foundation, either version 3 of the License, or            #
#(at your option) any later version.                                          #
#                                                                             #
#This program is distributed in the hope that it will be useful,              #
#but WITHOUT ANY WARRANTY; without even the implied warranty of               #
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the                #
#GNU Affero General Public License for more details.                          #
#                                                                             #
#You should have received a copy of the GNU Affero General Public License     #
#along with this program.  If not, see <http://www.gnu.org/licenses/>.        #
###############################################################################

from osv import osv, fields
from tools.translate import _
import decimal_precision as dp

class purchase_order(osv.Model):
    _inherit = 'purchase.order'

    def _amount_all(self, cr, uid, ids, field_name, arg, context=None):
        res = {}
        cur_obj = self.pool.get('res.currency')
        for order in self.browse(cr, uid, ids, context=context):
            res[order.id] = {
                'amount_untaxed': 0.0,
                'amount_tax': 0.0,
                'amount_total': 0.0,
            }
            val = val1 = 0.0
            cur = order.pricelist_id.currency_id
            for line in order.order_line:
                val1 += line.price_subtotal

                for c in self.pool.get('account.tax').compute_all(
                    cr, uid, line.taxes_id, line.price_unit, line.product_qty,
                    line.product_id.id, order.partner_id,
                    fiscal_position=line.fiscal_position)['taxes']:

                    tax_brw = self.pool.get('account.tax').browse(
                        cr, uid, c['id'])
                    if not tax_brw.tax_code_id.tax_discount:
                        val += c.get('amount', 0.0)
            res[order.id]['amount_tax'] = cur_obj.round(cr, uid, cur, val)
            res[order.id]['amount_untaxed'] = cur_obj.round(cr, uid, cur, val1)
            res[order.id]['amount_total'] = res[order.id]['amount_untaxed'] + res[order.id]['amount_tax']
        return res

    def _get_order(self, cr, uid, ids, context=None):
        result = {}
        for line in self.pool.get('purchase.order.line').browse(
            cr, uid, ids, context=context):
            result[line.order_id.id] = True
        return result.keys()

    def _default_fiscal_category(self, cr, uid, context=None):
        user = self.pool.get('res.users').browse(cr, uid, uid, context=context)
        return user.company_id.purchase_fiscal_category_id and \
        user.company_id.purchase_fiscal_category_id.id or False

    _columns = {
        'fiscal_category_id': fields.many2one(
            'l10n_br_account.fiscal.category', 'Categoria Fiscal',
            domain="[('type', '=', 'input'), \
            ('journal_type', '=', 'purchase')]"),
        'amount_untaxed': fields.function(
            _amount_all, method=True,
            digits_compute=dp.get_precision('Purchase Price'),
            string='Untaxed Amount',
            store={'purchase.order.line': (_get_order, None, 10)},
            multi="sums", help="The amount without tax"),
        'amount_tax': fields.function(
            _amount_all, method=True,
            digits_compute=dp.get_precision('Purchase Price'), string='Taxes',
            store={'purchase.order.line': (_get_order, None, 10)},
            multi="sums", help="The tax amount"),
        'amount_total': fields.function(
            _amount_all, method=True,
            digits_compute=dp.get_precision('Purchase Price'), string='Total',
            store={'purchase.order.line': (_get_order, None, 10)},
            multi="sums", help="The total amount")}

    _defaults = {
        'fiscal_category_id': _default_fiscal_category}

    def onchange_partner_id(self, cr, uid, ids, partner_id=False,
                            company_id=False, context=None,
                            fiscal_category_id=False, **kwargs):
        return super(purchase_order, self).onchange_partner_id(
            cr, uid, ids, partner_id, company_id, context=None,
            fiscal_category_id=fiscal_category_id)

    def onchange_dest_address_id(self, cr, uid, ids, partner_id,
                                 dest_address_id, company_id, context,
                                 fiscal_category_id, **kwargs):
        return super(purchase_order, self).onchange_dest_address_id(
            cr, uid, ids, partner_id, dest_address_id, company_id, context,
            fiscal_category_id=fiscal_category_id)

    def onchange_company_id(self, cr, uid, ids, partner_id,
                                 dest_address_id, company_id, context,
                                 fiscal_category_id, **kwargs):
        return super(purchase_order, self).onchange_company_id(
            cr, uid, ids, partner_id, dest_address_id, company_id, context,
            fiscal_category_id=fiscal_category_id)

    def onchange_fiscal_category_id(self, cr, uid, ids, partner_id=False,
                                    dest_address_id=False, company_id=False,
                                    context=None, fiscal_category_id=False,
                                    **kwargs):
        if not context:
            context = {}

        result = {'value': {'fiscal_position': False}}

        if not partner_id or not company_id:
            return result

        kwargs.update({
            'company_id': company_id,
            'partner_id': partner_id,
            'partner_invoice_id': partner_id,
            'partner_shipping_id': dest_address_id,
            'context': context
        })
        return self._fiscal_position_map(cr, uid, result, **kwargs)

    def _prepare_inv_line(self, cr, uid, account_id, order_line, context=None):

        result = super(purchase_order, self)._prepare_inv_line(
            cr, uid, account_id, order_line, context)

        order = order_line.order_id

        result['fiscal_category_id'] = order_line.fiscal_category_id and \
        order_line.fiscal_category_id.id or order.fiscal_category_id and \
        order.fiscal_category_id.id

        result['fiscal_position'] = order_line.fiscal_position and \
        order_line.fiscal_position.id or order.fiscal_position and \
        order.fiscal_position.id

        result['cfop_id'] = order.fiscal_position and \
        order.fiscal_position.cfop_id and order.fiscal_position.cfop_id.id or \
        order.fiscal_position and order.fiscal_position.cfop_id and \
        order.fiscal_position.cfop_id.id

        return result

    # TODO ask OpenERP SA for a _prepare_invoice method!
    def action_invoice_create(self, cr, uid, ids, *args):
        inv_id = super(purchase_order, self).action_invoice_create(cr, uid,
                                                                   ids, *args)
        for order in self.browse(cr, uid, ids):
            # REMARK: super method is ugly as it assumes only one invoice
            # for possibly several purchase orders.
            if inv_id:
                company_id = order.company_id
                if not company_id.document_serie_product_ids:
                    raise osv.except_osv(
                        _('No fiscal document serie found!'),
                        _("No fiscal document serie found for selected \
                        company %s") % (order.company_id.name))

                journal_id = order.fiscal_category_id and \
                order.fiscal_category_id.property_journal.id or False

                if not journal_id:
                    raise osv.except_osv(
                        _(u'Nenhuma Diário!'),
                        _(u"Categoria de operação fisca: '%s', não tem um \
                        diário contábil para a empresa %s") % (
                            order.fiscal_category_id.name,
                            order.company_id.name))

                comment = ''
                if order.fiscal_position.inv_copy_note:
                    comment = order.fiscal_position.note or ''
                if order.notes:
                    comment += ' - ' + order.notes

                self.pool.get('account.invoice').write(cr, uid, inv_id, {
                     'fiscal_category_id': order.fiscal_category_id and \
                     order.fiscal_category_id.id,
                     'fiscal_position': order.fiscal_position and \
                     order.fiscal_position.id,
                     'own_invoice': False,
                     'comment': comment,
                     'journal_id': journal_id})
        return inv_id

    def action_picking_create(self, cr, uid, ids, *args):
        picking_id = False
        for order in self.browse(cr, uid, ids):
            picking_id = super(purchase_order, self).action_picking_create(
                cr, uid, ids, *args)
            self.pool.get('stock.picking').write(
                cr, uid, picking_id,
                {'fiscal_category_id': order.fiscal_category_id.id,
                 'fiscal_position': order.fiscal_position.id})
        return picking_id

class purchase_order_line(osv.Model):
    _inherit = 'purchase.order.line'
    _columns = {
        'fiscal_category_id': fields.many2one(
            'l10n_br_account.fiscal.category', 'Categoria Fiscal',
            domain="[('type', '=', 'input'), \
            ('journal_type', '=', 'purchase')]"),
        'fiscal_position': fields.many2one(
            'account.fiscal.position', u'Posição Fiscal',
            domain="[('fiscal_category_id', '=', fiscal_category_id)]")
    }

    def _fiscal_position_map(self, cr, uid, result, **kwargs):

        kwargs['context'].update({'use_domain': ('use_purchase', '=', True)})
        fp_rule_obj = self.pool.get('account.fiscal.position.rule')
        result_rule = fp_rule_obj.apply_fiscal_mapping(cr, uid, result, kwargs)
        print result_rule
        if kwargs.get('product_id', False) and result_rule.get('fiscal_position', False):
            obj_fposition = self.pool.get('account.fiscal.position').browse(
                cr, uid, result_rule['fiscal_position'])
            obj_product = self.pool.get('product.product').browse(
                cr, uid, kwargs.get('product_id', False))
            kwargs['context'].update({'fiscal_type': obj_product.fiscal_type,
                            'type_tax_use': 'purchase'})
            taxes = obj_product.supplier_taxes_id or False
            tax_ids = self.pool.get('account.fiscal.position').map_tax(
                cr, uid, obj_fposition, taxes, kwargs.get('context'))

            result_rule['taxes_id'] = tax_ids

        return result_rule

    def product_id_change(self, cr, uid, ids, pricelist_id, product_id, qty,
                          uom_id, partner_id, date_order=False,
                          fiscal_position_id=False, date_planned=False,
                          name=False, price_unit=False, context=None,
                          company_id=False, parent_fiscal_position_id=False,
                          parent_fiscal_category_id=False, **kwargs):

        if context is None:
            context = {}

        result = super(purchase_order_line, self).product_id_change(
            cr, uid, ids, pricelist_id, product_id, qty, uom_id, partner_id,
            date_order, fiscal_position_id, date_planned, name, price_unit,
            context)

        if not product_id or not parent_fiscal_category_id:
            return result

        obj_fp_rule = self.pool.get('account.fiscal.position.rule')
        product_fiscal_category_id = obj_fp_rule.product_fiscal_category_map(
            cr, uid, product_id, parent_fiscal_category_id)

        if product_fiscal_category_id:
            parent_fiscal_category_id = product_fiscal_category_id

        result['value']['fiscal_category_id'] = parent_fiscal_category_id

        kwargs.update({
            'company_id': company_id,
            'product_id': product_id,
            'partner_id': partner_id,
            'partner_invoice_id': partner_id,
            'fiscal_category_id': parent_fiscal_category_id,
            'context': context,
        })
        return self._fiscal_position_map(cr, uid, result, **kwargs)

    def onchange_fiscal_category_id(self, cr, uid, ids, partner_id,
                                        dest_address_id=False,
                                        product_id=False,
                                        fiscal_category_id=False,
                                        company_id=False, context=None,
                                        **kwargs):
        result = {'value': {}}
        if not company_id or not partner_id:
            return result

        kwargs.update({
            'company_id': company_id,
            'product_id': product_id,
            'partner_id': partner_id,
            'partner_invoice_id': partner_id,
            'fiscal_category_id': fiscal_category_id,
            'context': context,
        })
        return self._fiscal_position_map(cr, uid, result, **kwargs)

    def onchange_fiscal_position(self, cr, uid, ids, partner_id,
                                 dest_address_id=False, product_id=False,
                                 fiscal_position=False,
                                 fiscal_category_id=False, company_id=False,
                                 context=None, **kwargs):
        result = {'value': {}}
        if not company_id or not partner_id or not fiscal_position:
            return result

        kwargs.update({
            'company_id': company_id,
            'product_id': product_id,
            'partner_id': partner_id,
            'partner_invoice_id': partner_id,
            'fiscal_category_id': fiscal_category_id,
            'context': context,
        })
        return self._fiscal_position_map(cr, uid, result, **kwargs)
<commit_msg>[MIGR] l10n_br_account: alterado o campo own_invoice para issuer como um campo selection indicando 0 emissão própria e 1 terceiros, para harmonizar com a tabela 5 do SPED Fiscal<commit_after># -*- encoding: utf-8 -*-
###############################################################################
#                                                                             #
# Copyright (C) 2009  Renato Lima - Akretion, Gabriel C. Stabel               #
# Copyright (C) 2012  Raphaël Valyi - Akretion                                #
#                                                                             #
#This program is free software: you can redistribute it and/or modify         #
#it under the terms of the GNU Affero General Public License as published by  #
#the Free Software Foundation, either version 3 of the License, or            #
#(at your option) any later version.                                          #
#                                                                             #
#This program is distributed in the hope that it will be useful,              #
#but WITHOUT ANY WARRANTY; without even the implied warranty of               #
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the                #
#GNU Affero General Public License for more details.                          #
#                                                                             #
#You should have received a copy of the GNU Affero General Public License     #
#along with this program.  If not, see <http://www.gnu.org/licenses/>.        #
###############################################################################

from osv import osv, fields
from tools.translate import _
import decimal_precision as dp

class purchase_order(osv.Model):
    _inherit = 'purchase.order'

    def _amount_all(self, cr, uid, ids, field_name, arg, context=None):
        res = {}
        cur_obj = self.pool.get('res.currency')
        for order in self.browse(cr, uid, ids, context=context):
            res[order.id] = {
                'amount_untaxed': 0.0,
                'amount_tax': 0.0,
                'amount_total': 0.0,
            }
            val = val1 = 0.0
            cur = order.pricelist_id.currency_id
            for line in order.order_line:
                val1 += line.price_subtotal

                for c in self.pool.get('account.tax').compute_all(
                    cr, uid, line.taxes_id, line.price_unit, line.product_qty,
                    line.product_id.id, order.partner_id,
                    fiscal_position=line.fiscal_position)['taxes']:

                    tax_brw = self.pool.get('account.tax').browse(
                        cr, uid, c['id'])
                    if not tax_brw.tax_code_id.tax_discount:
                        val += c.get('amount', 0.0)
            res[order.id]['amount_tax'] = cur_obj.round(cr, uid, cur, val)
            res[order.id]['amount_untaxed'] = cur_obj.round(cr, uid, cur, val1)
            res[order.id]['amount_total'] = res[order.id]['amount_untaxed'] + res[order.id]['amount_tax']
        return res

    def _get_order(self, cr, uid, ids, context=None):
        result = {}
        for line in self.pool.get('purchase.order.line').browse(
            cr, uid, ids, context=context):
            result[line.order_id.id] = True
        return result.keys()

    def _default_fiscal_category(self, cr, uid, context=None):
        user = self.pool.get('res.users').browse(cr, uid, uid, context=context)
        return user.company_id.purchase_fiscal_category_id and \
        user.company_id.purchase_fiscal_category_id.id or False

    _columns = {
        'fiscal_category_id': fields.many2one(
            'l10n_br_account.fiscal.category', 'Categoria Fiscal',
            domain="[('type', '=', 'input'), \
            ('journal_type', '=', 'purchase')]"),
        'amount_untaxed': fields.function(
            _amount_all, method=True,
            digits_compute=dp.get_precision('Purchase Price'),
            string='Untaxed Amount',
            store={'purchase.order.line': (_get_order, None, 10)},
            multi="sums", help="The amount without tax"),
        'amount_tax': fields.function(
            _amount_all, method=True,
            digits_compute=dp.get_precision('Purchase Price'), string='Taxes',
            store={'purchase.order.line': (_get_order, None, 10)},
            multi="sums", help="The tax amount"),
        'amount_total': fields.function(
            _amount_all, method=True,
            digits_compute=dp.get_precision('Purchase Price'), string='Total',
            store={'purchase.order.line': (_get_order, None, 10)},
            multi="sums", help="The total amount")}

    _defaults = {
        'fiscal_category_id': _default_fiscal_category}

    def onchange_partner_id(self, cr, uid, ids, partner_id=False,
                            company_id=False, context=None,
                            fiscal_category_id=False, **kwargs):
        return super(purchase_order, self).onchange_partner_id(
            cr, uid, ids, partner_id, company_id, context=None,
            fiscal_category_id=fiscal_category_id)

    def onchange_dest_address_id(self, cr, uid, ids, partner_id,
                                 dest_address_id, company_id, context,
                                 fiscal_category_id, **kwargs):
        return super(purchase_order, self).onchange_dest_address_id(
            cr, uid, ids, partner_id, dest_address_id, company_id, context,
            fiscal_category_id=fiscal_category_id)

    def onchange_company_id(self, cr, uid, ids, partner_id,
                                 dest_address_id, company_id, context,
                                 fiscal_category_id, **kwargs):
        return super(purchase_order, self).onchange_company_id(
            cr, uid, ids, partner_id, dest_address_id, company_id, context,
            fiscal_category_id=fiscal_category_id)

    def onchange_fiscal_category_id(self, cr, uid, ids, partner_id=False,
                                    dest_address_id=False, company_id=False,
                                    context=None, fiscal_category_id=False,
                                    **kwargs):
        if not context:
            context = {}

        result = {'value': {'fiscal_position': False}}

        if not partner_id or not company_id:
            return result

        kwargs.update({
            'company_id': company_id,
            'partner_id': partner_id,
            'partner_invoice_id': partner_id,
            'partner_shipping_id': dest_address_id,
            'context': context
        })
        return self._fiscal_position_map(cr, uid, result, **kwargs)

    def _prepare_inv_line(self, cr, uid, account_id, order_line, context=None):

        result = super(purchase_order, self)._prepare_inv_line(
            cr, uid, account_id, order_line, context)

        order = order_line.order_id

        result['fiscal_category_id'] = order_line.fiscal_category_id and \
        order_line.fiscal_category_id.id or order.fiscal_category_id and \
        order.fiscal_category_id.id

        result['fiscal_position'] = order_line.fiscal_position and \
        order_line.fiscal_position.id or order.fiscal_position and \
        order.fiscal_position.id

        result['cfop_id'] = order.fiscal_position and \
        order.fiscal_position.cfop_id and order.fiscal_position.cfop_id.id or \
        order.fiscal_position and order.fiscal_position.cfop_id and \
        order.fiscal_position.cfop_id.id

        return result

    # TODO ask OpenERP SA for a _prepare_invoice method!
    def action_invoice_create(self, cr, uid, ids, *args):
        inv_id = super(purchase_order, self).action_invoice_create(cr, uid,
                                                                   ids, *args)
        for order in self.browse(cr, uid, ids):
            # REMARK: super method is ugly as it assumes only one invoice
            # for possibly several purchase orders.
            if inv_id:
                company_id = order.company_id
                if not company_id.document_serie_product_ids:
                    raise osv.except_osv(
                        _('No fiscal document serie found!'),
                        _("No fiscal document serie found for selected \
                        company %s") % (order.company_id.name))

                journal_id = order.fiscal_category_id and \
                order.fiscal_category_id.property_journal.id or False

                if not journal_id:
                    raise osv.except_osv(
                        _(u'Nenhuma Diário!'),
                        _(u"Categoria de operação fisca: '%s', não tem um \
                        diário contábil para a empresa %s") % (
                            order.fiscal_category_id.name,
                            order.company_id.name))

                comment = ''
                if order.fiscal_position.inv_copy_note:
                    comment = order.fiscal_position.note or ''
                if order.notes:
                    comment += ' - ' + order.notes

                self.pool.get('account.invoice').write(cr, uid, inv_id, {
                     'fiscal_category_id': order.fiscal_category_id and
                     order.fiscal_category_id.id,
                     'fiscal_position': order.fiscal_position and
                     order.fiscal_position.id,
                     'issuer': '1',
                     'comment': comment,
                     'journal_id': journal_id})
        return inv_id

    def action_picking_create(self, cr, uid, ids, *args):
        picking_id = False
        for order in self.browse(cr, uid, ids):
            picking_id = super(purchase_order, self).action_picking_create(
                cr, uid, ids, *args)
            self.pool.get('stock.picking').write(
                cr, uid, picking_id,
                {'fiscal_category_id': order.fiscal_category_id.id,
                 'fiscal_position': order.fiscal_position.id})
        return picking_id

class purchase_order_line(osv.Model):
    _inherit = 'purchase.order.line'
    _columns = {
        'fiscal_category_id': fields.many2one(
            'l10n_br_account.fiscal.category', 'Categoria Fiscal',
            domain="[('type', '=', 'input'), \
            ('journal_type', '=', 'purchase')]"),
        'fiscal_position': fields.many2one(
            'account.fiscal.position', u'Posição Fiscal',
            domain="[('fiscal_category_id', '=', fiscal_category_id)]")
    }

    def _fiscal_position_map(self, cr, uid, result, **kwargs):

        kwargs['context'].update({'use_domain': ('use_purchase', '=', True)})
        fp_rule_obj = self.pool.get('account.fiscal.position.rule')
        result_rule = fp_rule_obj.apply_fiscal_mapping(cr, uid, result, kwargs)
        print result_rule
        if kwargs.get('product_id', False) and \
        result_rule.get('fiscal_position', False):
            obj_fposition = self.pool.get('account.fiscal.position').browse(
                cr, uid, result_rule['fiscal_position'])
            obj_product = self.pool.get('product.product').browse(
                cr, uid, kwargs.get('product_id', False))
            kwargs['context'].update({'fiscal_type': obj_product.fiscal_type,
                            'type_tax_use': 'purchase'})
            taxes = obj_product.supplier_taxes_id or False
            tax_ids = self.pool.get('account.fiscal.position').map_tax(
                cr, uid, obj_fposition, taxes, kwargs.get('context'))

            result_rule['taxes_id'] = tax_ids

        return result_rule

    def product_id_change(self, cr, uid, ids, pricelist_id, product_id, qty,
                          uom_id, partner_id, date_order=False,
                          fiscal_position_id=False, date_planned=False,
                          name=False, price_unit=False, context=None,
                          company_id=False, parent_fiscal_position_id=False,
                          parent_fiscal_category_id=False, **kwargs):

        if context is None:
            context = {}

        result = super(purchase_order_line, self).product_id_change(
            cr, uid, ids, pricelist_id, product_id, qty, uom_id, partner_id,
            date_order, fiscal_position_id, date_planned, name, price_unit,
            context)

        if not product_id or not parent_fiscal_category_id:
            return result

        obj_fp_rule = self.pool.get('account.fiscal.position.rule')
        product_fiscal_category_id = obj_fp_rule.product_fiscal_category_map(
            cr, uid, product_id, parent_fiscal_category_id)

        if product_fiscal_category_id:
            parent_fiscal_category_id = product_fiscal_category_id

        result['value']['fiscal_category_id'] = parent_fiscal_category_id

        kwargs.update({
            'company_id': company_id,
            'product_id': product_id,
            'partner_id': partner_id,
            'partner_invoice_id': partner_id,
            'fiscal_category_id': parent_fiscal_category_id,
            'context': context,
        })
        return self._fiscal_position_map(cr, uid, result, **kwargs)

    def onchange_fiscal_category_id(self, cr, uid, ids, partner_id,
                                        dest_address_id=False,
                                        product_id=False,
                                        fiscal_category_id=False,
                                        company_id=False, context=None,
                                        **kwargs):
        result = {'value': {}}
        if not company_id or not partner_id:
            return result

        kwargs.update({
            'company_id': company_id,
            'product_id': product_id,
            'partner_id': partner_id,
            'partner_invoice_id': partner_id,
            'fiscal_category_id': fiscal_category_id,
            'context': context,
        })
        return self._fiscal_position_map(cr, uid, result, **kwargs)

    def onchange_fiscal_position(self, cr, uid, ids, partner_id,
                                 dest_address_id=False, product_id=False,
                                 fiscal_position=False,
                                 fiscal_category_id=False, company_id=False,
                                 context=None, **kwargs):
        result = {'value': {}}
        if not company_id or not partner_id or not fiscal_position:
            return result

        kwargs.update({
            'company_id': company_id,
            'product_id': product_id,
            'partner_id': partner_id,
            'partner_invoice_id': partner_id,
            'fiscal_category_id': fiscal_category_id,
            'context': context,
        })
        return self._fiscal_position_map(cr, uid, result, **kwargs)
<|endoftext|>


<commit_before>import prairielearn as pl
import lxml.html
import chevron, os

def prepare(element_html, element_index, data):
    element = lxml.html.fragment_fromstring(element_html)
    required_attribs = ['script_name']
    optional_attribs = ['params_names', 'width', 'height']
    pl.check_attribs(element, required_attribs, optional_attribs)
    return data

def render(element_html, element_index, data):
    element = lxml.html.fragment_fromstring(element_html)
    script_name = pl.get_string_attrib(element, 'script-name', None)
    if script_name is None:
        raise Exception("no script-name attribute for pl_prairiedraw_figure")

    with open(os.path.join(data["options"]["question_path"], script_name)) as f:
        script = f.read()

    width = pl.get_string_attrib(element, 'width', '500')
    height = pl.get_string_attrib(element, 'height', '300')

    params_names = pl.get_string_attrib(element, 'params-names', None)
    if params_names is None:
        client_params = {}
    else:
        params_names = params_names.split(sep=',')
        client_params = {key: data['params'][key] for key in params_names}

    html_params = {
        'script': script,
        'width': width,
        'height': height,
        'client_params': client_params,
        'uuid': pl.get_uuid(),
    }

    with open('pl-prairiedraw-figure.mustache') as f:
        html = chevron.render(f, html_params).strip()

    return html
<commit_msg>Linter fixes<commit_after>import prairielearn as pl
import lxml.html
import chevron
import os

def prepare(element_html, element_index, data):
    element = lxml.html.fragment_fromstring(element_html)
    required_attribs = ['script_name']
    optional_attribs = ['params_names', 'width', 'height']
    pl.check_attribs(element, required_attribs, optional_attribs)
    return data

def render(element_html, element_index, data):
    element = lxml.html.fragment_fromstring(element_html)
    script_name = pl.get_string_attrib(element, 'script-name', None)
    if script_name is None:
        raise Exception('no script-name attribute for pl_prairiedraw_figure')

    with open(os.path.join(data['options']['question_path'], script_name)) as f:
        script = f.read()

    width = pl.get_string_attrib(element, 'width', '500')
    height = pl.get_string_attrib(element, 'height', '300')

    params_names = pl.get_string_attrib(element, 'params-names', None)
    if params_names is None:
        client_params = {}
    else:
        params_names = params_names.split(sep=',')
        client_params = {key: data['params'][key] for key in params_names}

    html_params = {
        'script': script,
        'width': width,
        'height': height,
        'client_params': client_params,
        'uuid': pl.get_uuid(),
    }

    with open('pl-prairiedraw-figure.mustache') as f:
        html = chevron.render(f, html_params).strip()

    return html
<|endoftext|>


<commit_before>import re
import os
import imaplib
import time
import urllib2
import email

THIS_DIR = os.path.dirname(os.path.abspath(__file__))
execfile(os.path.join(THIS_DIR, 'version.py'))

__version__ = VERSION

class ImapLibrary(object):

    ROBOT_LIBRARY_VERSION = VERSION
    ROBOT_LIBRARY_SCOPE = 'GLOBAL'

    port = 993

    def open_mailbox(self, server, user, password):
        """
        Open the mailbox on a mail server with a valid
        authentication.
        """
        self.imap = imaplib.IMAP4_SSL(server, self.port)
        self.imap.login(user, password)
        self.imap.select()
        self._init_walking_multipart()

    def wait_for_mail(self, fromEmail=None, toEmail=None, status=None,
                      timeout=60):
        """
        Wait for an incoming mail from a specific sender to
        a specific mail receiver. Check the mailbox every 10
        seconds for incoming mails until the timeout is exceeded.
        Returns the mail number of the latest email received.

        `timeout` sets the maximum waiting time until an error
        is raised.
        """
        timeout = int(timeout)
        while (timeout > 0):
            self.imap.recent()
            self.mails = self._check_emails(fromEmail, toEmail, status)
            if len(self.mails) > 0:
                return self.mails[-1]
            timeout -= 10
            if timeout > 0:
                time.sleep(10)
        raise AssertionError("No mail received within time")

    def get_links_from_email(self, mailNumber):
        '''
        Finds all links in an email body and returns them

        `mailNumber` is the index number of the mail to open
        '''
        body = self.get_email_body(mailNumber)
        return re.findall(r'href=[\'"]?([^\'" >]+)', body)

    def get_matches_from_email(self, mailNumber, regexp):
        """
        Finds all occurrences of a regular expression
        """
        body = self.get_email_body(mailNumber)
        return re.findall(regexp, body)

    def open_link_from_mail(self, mailNumber, linkNumber=0):
        """
        Find a link in an email body and open the link.
        Returns the link's html.

        `mailNumber` is the index number of the mail to open
        `linkNumber` declares which link shall be opened (link
        index in body text)
        """
        urls = self.get_links_from_email(mailNumber)

        if len(urls) > linkNumber:
            resp = urllib2.urlopen(urls[linkNumber])
            content_type = resp.headers.getheader('content-type')
            if content_type:
                enc = content_type.split('charset=')[-1]
                return unicode(resp.read(), enc)
            else:
                return resp.read()
        else:
            raise AssertionError("Link number %i not found!" % linkNumber)

    def close_mailbox(self):
        """
        Close the mailbox after finishing all mail activities of a user.
        """
        self.imap.close()

    def mark_as_read(self):
        """
        Mark all received mails as read
        """
        for mail in self.mails:
            self.imap.store(mail, '+FLAGS', '\SEEN')

    def get_email_body(self, mailNumber):
        """
        Returns an email body

        `mailNumber` is the index number of the mail to open
        """
        if self._is_walking_multipart(mailNumber):
            body = self.get_multipart_payload(decode=True)
        else:
            body = self.imap.fetch(mailNumber, '(BODY[TEXT])')[1][0][1].decode('quoted-printable')
        return body

    def walk_multipart_email(self, mailNumber):
        """
        Returns the number of parts of a multipart email. Content is stored internally
        to be used by other multipart keywords. Subsequent calls iterate over the
        elements, and the various Get Multipart keywords retrieve their contents.

        `mailNumber` is the index number of the mail to open
        """
        if not self._is_walking_multipart(mailNumber):
            data = self.imap.fetch(mailNumber, '(RFC822)')[1][0][1]
            msg = email.message_from_string(data.decode())
            self._start_walking_multipart(mailNumber, msg)

        try:
            self._part = next(self._mp_iter)
        except StopIteration:
            self._init_walking_multipart()
            return False

        # return number of parts
        return len(self._mp_msg.get_payload())

    def _is_walking_multipart(self, mailNumber):
        """
        Check if walking a multipart email is in progress
        """
        return self._mp_msg is not None and self._mailNumber == mailNumber

    def _start_walking_multipart(self, mailNumber, msg):
        self._mailNumber = mailNumber
        self._mp_msg = msg
        self._mp_iter = msg.walk()

    def _init_walking_multipart(self):
        self._mp_msg = None
        self._part = None
        self._mailNumber = None

    def get_multipart_content_type(self):
        """
        Return the content-type for the current part of a multipart email
        """
        return self._part.get_content_type()

    def get_multipart_payload(self, decode=False):
        """
        Return the payload for the current part of a multipart email

        decode is an optional flag that indicates whether to decoding
        """
        s = self._part.get_payload(decode=decode)
        charset = self._part.get_content_charset()
        if charset is not None:
            return s.decode(charset)
        return s

    def get_multipart_field_names(self):
        """
        Return the list of header field names for the current multipart email
        """
        return self._mp_msg.keys()

    def get_multipart_field(self, field):
        """
        Returns the content of a header field

        field is a string such as 'From', 'To', 'Subject', 'Date', etc.
        """
        return self._mp_msg[field]

    def _criteria(self, fromEmail, toEmail, status):
        crit = []
        if fromEmail:
            crit += ['FROM', fromEmail]
        if toEmail:
            crit += ['TO', toEmail]
        if status:
            crit += [status]
        if not crit:
            crit = ['UNSEEN']
        return crit

    def _check_emails(self, fromEmail, toEmail, status):
        crit = self._criteria(fromEmail, toEmail, status)
        type, msgnums = self.imap.search(None, *crit)
        return msgnums[0].split()
<commit_msg>Removed call to imap.recent() First email retrieved was ok, but calling imap.recent caused subsequent imap.fetch to fail within the same connection.<commit_after>import re
import os
import imaplib
import time
import urllib2
import email

THIS_DIR = os.path.dirname(os.path.abspath(__file__))
execfile(os.path.join(THIS_DIR, 'version.py'))

__version__ = VERSION

class ImapLibrary(object):

    ROBOT_LIBRARY_VERSION = VERSION
    ROBOT_LIBRARY_SCOPE = 'GLOBAL'

    port = 993

    def open_mailbox(self, server, user, password):
        """
        Open the mailbox on a mail server with a valid
        authentication.
        """
        self.imap = imaplib.IMAP4_SSL(server, self.port)
        self.imap.login(user, password)
        self.imap.select()
        self._init_walking_multipart()

    def wait_for_mail(self, fromEmail=None, toEmail=None, status=None,
                      timeout=60):
        """
        Wait for an incoming mail from a specific sender to
        a specific mail receiver. Check the mailbox every 10
        seconds for incoming mails until the timeout is exceeded.
        Returns the mail number of the latest email received.

        `timeout` sets the maximum waiting time until an error
        is raised.
        """
        timeout = int(timeout)
        while (timeout > 0):
            self.mails = self._check_emails(fromEmail, toEmail, status)
            if len(self.mails) > 0:
                return self.mails[-1]
            timeout -= 10
            if timeout > 0:
                time.sleep(10)
        raise AssertionError("No mail received within time")

    def get_links_from_email(self, mailNumber):
        '''
        Finds all links in an email body and returns them

        `mailNumber` is the index number of the mail to open
        '''
        body = self.get_email_body(mailNumber)
        return re.findall(r'href=[\'"]?([^\'" >]+)', body)

    def get_matches_from_email(self, mailNumber, regexp):
        """
        Finds all occurrences of a regular expression
        """
        body = self.get_email_body(mailNumber)
        return re.findall(regexp, body)

    def open_link_from_mail(self, mailNumber, linkNumber=0):
        """
        Find a link in an email body and open the link.
        Returns the link's html.

        `mailNumber` is the index number of the mail to open
        `linkNumber` declares which link shall be opened (link
        index in body text)
        """
        urls = self.get_links_from_email(mailNumber)

        if len(urls) > linkNumber:
            resp = urllib2.urlopen(urls[linkNumber])
            content_type = resp.headers.getheader('content-type')
            if content_type:
                enc = content_type.split('charset=')[-1]
                return unicode(resp.read(), enc)
            else:
                return resp.read()
        else:
            raise AssertionError("Link number %i not found!" % linkNumber)

    def close_mailbox(self):
        """
        Close the mailbox after finishing all mail activities of a user.
        """
        self.imap.close()

    def mark_as_read(self):
        """
        Mark all received mails as read
        """
        for mail in self.mails:
            self.imap.store(mail, '+FLAGS', '\SEEN')

    def get_email_body(self, mailNumber):
        """
        Returns an email body

        `mailNumber` is the index number of the mail to open
        """
        if self._is_walking_multipart(mailNumber):
            body = self.get_multipart_payload(decode=True)
        else:
            body = self.imap.fetch(mailNumber, '(BODY[TEXT])')[1][0][1].decode('quoted-printable')
        return body

    def walk_multipart_email(self, mailNumber):
        """
        Returns the number of parts of a multipart email. Content is stored internally
        to be used by other multipart keywords. Subsequent calls iterate over the
        elements, and the various Get Multipart keywords retrieve their contents.

        `mailNumber` is the index number of the mail to open
        """
        if not self._is_walking_multipart(mailNumber):
            data = self.imap.fetch(mailNumber, '(RFC822)')[1][0][1]
            msg = email.message_from_string(data.decode())
            self._start_walking_multipart(mailNumber, msg)

        try:
            self._part = next(self._mp_iter)
        except StopIteration:
            self._init_walking_multipart()
            return False

        # return number of parts
        return len(self._mp_msg.get_payload())

    def _is_walking_multipart(self, mailNumber):
        """
        Check if walking a multipart email is in progress
        """
        return self._mp_msg is not None and self._mailNumber == mailNumber

    def _start_walking_multipart(self, mailNumber, msg):
        self._mailNumber = mailNumber
        self._mp_msg = msg
        self._mp_iter = msg.walk()

    def _init_walking_multipart(self):
        self._mp_msg = None
        self._part = None
        self._mailNumber = None

    def get_multipart_content_type(self):
        """
        Return the content-type for the current part of a multipart email
        """
        return self._part.get_content_type()

    def get_multipart_payload(self, decode=False):
        """
        Return the payload for the current part of a multipart email

        decode is an optional flag that indicates whether to decoding
        """
        s = self._part.get_payload(decode=decode)
        charset = self._part.get_content_charset()
        if charset is not None:
            return s.decode(charset)
        return s

    def get_multipart_field_names(self):
        """
        Return the list of header field names for the current multipart email
        """
        return self._mp_msg.keys()

    def get_multipart_field(self, field):
        """
        Returns the content of a header field

        field is a string such as 'From', 'To', 'Subject', 'Date', etc.
        """
        return self._mp_msg[field]

    def _criteria(self, fromEmail, toEmail, status):
        crit = []
        if fromEmail:
            crit += ['FROM', fromEmail]
        if toEmail:
            crit += ['TO', toEmail]
        if status:
            crit += [status]
        if not crit:
            crit = ['UNSEEN']
        return crit

    def _check_emails(self, fromEmail, toEmail, status):
        crit = self._criteria(fromEmail, toEmail, status)
        type, msgnums = self.imap.search(None, *crit)
        return msgnums[0].split()
<|endoftext|>