<commit_before>"""Settings file

Holds the list of available settings and default values

Usage
-----
    import inselect.settings

    # Run once at the very start of the application
    inselect.settings.init()
    # Get a value
    inselect.settings.get(name)
    # Set a value
    inselect.settings.setValue(name, value)

    # Alternatively, modules can use a QSettings object directly:
    from PySide.QtCore import QSettings
    q = QSettings(*inselect.settings.namespace)
    q.value(name)
    q.setValue(name, value)
"""

import sys
from PySide import QtCore
from inselect import __version__ as inselect_version
from inselect.lib import validators
from inselect.gui.settings import SettingsDialog

# Settings namespace
namespace = ('NHM', 'Inselect')

# Define the available settings. Each entry associates the internal setting
# name to a dictionary defining:
# label : str, required
#     Label as shown to the user.
# description : str, optional
#     Description as shown to the user. This is required if editable is True.
# editable : bool, optional
#     True if the setting can be changed by the user.
# type : str, optional
#     One of 'int', 'float', 'bool', 'list' or 'str' (default).
# validate : function, optional
#     Validation function, returns a boolean.
# default : object, required
#     Default value for the setting
# reset : bool, optional
#     If present and true, reset the value to default value on start up
_settings = {
    'annotation_fields': {
        'label': "Annotation fields",
        'description': 'Comma separated list of fields available in the annotation editor',
        'editable': True,
        'type': 'list',
        'validate': validators.not_empty,
        'default': ['Specimen Number', 'Current Taxon Name', 'Location in Collection']
    },
    'label_field': {
        'label': "Label field",
        'description': 'Field (from the annotation fields) to display for each item (eg. in the sidebar)',
        'editable': True,
        'type': 'str',
        'validate': validators.is_annotation_field,
        'default': 'Specimen Number'
    },
    'export_template': {
        'label': 'Export file name',
        'description': 'Template for image export file names. You can use any of the annotation fields
 between {'
                       'curly brackets}. To insert plain curly brackets, double them.',
        'editable': True,
        'type': 'str',
        'validate': validators.validate_export_template,
        'default': 'BMNHE_{Specimen Number}{Current Taxon Name}'
    },
    'working_directory': {
        'label': 'Working directory',
        'editable': False,
        'type': 'str',
        'default': None
    },
    'about_label': {
        'label': 'About (label)',
        'editable': False,
        'default': "Insect Selector"
    },
    'about_text': {
        'label': 'About',
        'editable': False,
        'reset': True,
        'default': """

Inselect

Insect selector, v.{version}

Contributors

                Stefan van der Walt: Application development
                and segmentation algorithm

                Pieter Holtzhausen: Application development
                and segmentation algorithm

                Alice Heaton: Application development

        """.format(version=inselect_version)
    },
    'icon_size': {
        'label': 'Icon size',
        'editable': False,
        'type': 'int',
        'reset': True,
        'default': 200
    }
}
_q_settings = None

def init():
    """Create the global QSettings object and setup the default values"""
    global _settings, _q_settings
    if sys.platform == 'win32':
        _settings['working_directory']['default'] = QtCore.QCoreApplication.applicationDirPath()
    else:
        _settings['working_directory']['default'] = QtCore.QDir.currentPath()
    _q_settings = QtCore.QSettings(*namespace)
    for name in _settings:
        reset_setting = 'reset' in _settings[name] and _settings[name]['reset']
        if not _q_settings.contains(name) or reset_setting:
            _q_settings.setValue(name, _settings[name]['default'])

def get(name):
    """Return the given setting's value

    Parameters
    ----------
    name: Name of setting

    Returns
    -------
    object
        Setting value
    """
    if 'type' in _settings[name]:
        value_type = _settings[name]['type']
    else:
        value_type = 'str'
    value = _q_settings.value(name)
    # QSettings handles list and strings, but not ints and floats!
    if value_type == 'int':
        return int(value)
    elif value_type == 'float':
        return float(value)
    elif value_type == 'bool':
        return bool(value)
    return value

def set_value(name, value):
    """Set the given settings's value

    Parameters
    ----------
    name: Name of setting
    value: Value of setting
    """
    #TODO: Should components be notified? For instance the sidebar needs
    # to know if the label field is changed.
    return _q_settings.setValue(name, value)

def definition():
    """Returns setting definition list"""
    return _settings

def open_settings_dialog():
    """Open the settings dialog"""
    dialog = SettingsDialog()
    dialog.exec_()

def reset(name=None):
    """Reset settings to default values

    If setting is None, then all settings are reset.

    Parameters
    ----------
    name : str, None
        The setting to reset, or None
    """
    global _settings, _q_settings
    if name:
        _q_settings.setValue(name, _settings[name]['default'])
    else:
        _q_settings.clear()
        init()
<commit_msg>Added myself to contributors in about box<commit_after>"""Settings file

Holds the list of available settings and default values

Usage
-----
    import inselect.settings

    # Run once at the very start of the application
    inselect.settings.init()
    # Get a value
    inselect.settings.get(name)
    # Set a value
    inselect.settings.setValue(name, value)

    # Alternatively, modules can use a QSettings object directly:
    from PySide.QtCore import QSettings
    q = QSettings(*inselect.settings.namespace)
    q.value(name)
    q.setValue(name, value)
"""

import sys
from PySide import QtCore
from inselect import __version__ as inselect_version
from inselect.lib import validators
from inselect.gui.settings import SettingsDialog

# Settings namespace
namespace = ('NHM', 'Inselect')

# Define the available settings. Each entry associates the internal setting
# name to a dictionary defining:
# label : str, required
#     Label as shown to the user.
# description : str, optional
#     Description as shown to the user. This is required if editable is True.
# editable : bool, optional
#     True if the setting can be changed by the user.
# type : str, optional
#     One of 'int', 'float', 'bool', 'list' or 'str' (default).
# validate : function, optional
#     Validation function, returns a boolean.
# default : object, required
#     Default value for the setting
# reset : bool, optional
#     If present and true, reset the value to default value on start up
_settings = {
    'annotation_fields': {
        'label': "Annotation fields",
        'description': 'Comma separated list of fields available in the annotation editor',
        'editable': True,
        'type': 'list',
        'validate': validators.not_empty,
        'default': ['Specimen Number', 'Current Taxon Name', 'Location in Collection']
    },
    'label_field': {
        'label': "Label field",
        'description': 'Field (from the annotation fields) to display for each item (eg. in the sidebar)',
        'editable': True,
        'type': 'str',
        'validate': validators.is_annotation_field,
        'default': 'Specimen Number'
    },
    'export_template': {
        'label': 'Export file name',
        'description': 'Template for image export file names. You can use any of the annotation fields
 between {'
                       'curly brackets}. To insert plain curly brackets, double them.',
        'editable': True,
        'type': 'str',
        'validate': validators.validate_export_template,
        'default': 'BMNHE_{Specimen Number}{Current Taxon Name}'
    },
    'working_directory': {
        'label': 'Working directory',
        'editable': False,
        'type': 'str',
        'default': None
    },
    'about_label': {
        'label': 'About (label)',
        'editable': False,
        'default': "Insect Selector"
    },
    'about_text': {
        'label': 'About',
        'editable': False,
        'reset': True,
        'default': """

Inselect

Insect selector, v.{version}

Contributors

                Stefan van der Walt: Application development
                and segmentation algorithm

                Pieter Holtzhausen: Application development
                and segmentation algorithm

                Alice Heaton: Application development

                Lawrence Hudson: Application development

        """.format(version=inselect_version)
    },
    'icon_size': {
        'label': 'Icon size',
        'editable': False,
        'type': 'int',
        'reset': True,
        'default': 200
    }
}
_q_settings = None

def init():
    """Create the global QSettings object and setup the default values"""
    global _settings, _q_settings
    if sys.platform == 'win32':
        _settings['working_directory']['default'] = QtCore.QCoreApplication.applicationDirPath()
    else:
        _settings['working_directory']['default'] = QtCore.QDir.currentPath()
    _q_settings = QtCore.QSettings(*namespace)
    for name in _settings:
        reset_setting = 'reset' in _settings[name] and _settings[name]['reset']
        if not _q_settings.contains(name) or reset_setting:
            _q_settings.setValue(name, _settings[name]['default'])

def get(name):
    """Return the given setting's value

    Parameters
    ----------
    name: Name of setting

    Returns
    -------
    object
        Setting value
    """
    if 'type' in _settings[name]:
        value_type = _settings[name]['type']
    else:
        value_type = 'str'
    value = _q_settings.value(name)
    # QSettings handles list and strings, but not ints and floats!
    if value_type == 'int':
        return int(value)
    elif value_type == 'float':
        return float(value)
    elif value_type == 'bool':
        return bool(value)
    return value

def set_value(name, value):
    """Set the given settings's value

    Parameters
    ----------
    name: Name of setting
    value: Value of setting
    """
    #TODO: Should components be notified? For instance the sidebar needs
    # to know if the label field is changed.
    return _q_settings.setValue(name, value)

def definition():
    """Returns setting definition list"""
    return _settings

def open_settings_dialog():
    """Open the settings dialog"""
    dialog = SettingsDialog()
    dialog.exec_()

def reset(name=None):
    """Reset settings to default values

    If setting is None, then all settings are reset.

    Parameters
    ----------
    name : str, None
        The setting to reset, or None
    """
    global _settings, _q_settings
    if name:
        _q_settings.setValue(name, _settings[name]['default'])
    else:
        _q_settings.clear()
        init()
<|endoftext|>


<commit_before>#!/usr/bin/env python
from collections import namedtuple
import re

OptTuple = namedtuple('OptTuple', ('flag', 'arg', 'info', 'default', 'dtype'))

options = []
def opt(flag, arg='', info='', default=None, dtype=None):
    info = re.sub('\s+', ' ', info).strip()
    if dtype is not None:
        dtype = dtype.__name__
    options.append(OptTuple(flag, arg, info, default, dtype))

def switch(flag, info=''):
    opt(flag, arg='', info=info, default=False, dtype=bool)

switch('-b', 'dmenu appears at the bottom of the screen.')
switch('-f',
    """dmenu grabs the keyboard before reading stdin.  This is
    faster,  but  will  lock   up  X  until  stdin  reaches
    end-of-file.""")
switch('-i', 'dmenu matches menu items case insensitively.')
opt('-l', 'lines',
    "dmenu lists items vertically,  with the given number of lines.",
    dtype=int,)
opt('-m', 'monitor',
    "dmenu    is   displayed    on   the    monitor   number supplied. Monitor numbers are starting from 0.",
    dtype=int)
opt('-p', 'prompt',
    "defines the prompt  to be displayed to the  left of the input field.",
    dtype=str)

opt('-fn', 'font',
    "defines the font or font set used.")
opt('-nb', 'color',
    "defines  the normal  background color.   #RGB, #RRGGBB, and X color names are supported.")
opt('-nf', 'color', 'defines the normal foreground color.')
opt('-sb', 'color', 'defines the selected background color.')
opt('-sf', 'color', 'defines the selected foreground color.')
switch('-v', 'prints version information to stdout, then exits.')

if __name__ == '__main__':
    import json
    options = [x._asdict() for x in options]
    print(json.dumps(options, indent=2, sort_keys=True))

<commit_msg>Update gen_opts_dmenu<commit_after>#!/usr/bin/env python
import re
import subprocess as sp
from itertools import chain
from class_generator import Flag, Option, MenuType, Assignment

def get_outp(*cmd):
    p = sp.Popen(cmd, stdout=sp.PIPE, stderr=sp.PIPE)
    stdout, stderr = p.communicate()
    return stdout.decode()

def man_dmenu():
    return get_outp('man', '-P', 'cat', 'dmenu')

def get_sections(man_page):
    pat = '^(?P<section>[A-Z]+[\sA-Z]*$)'
    secs = list(re.finditer(pat, man_page, flags=re.MULTILINE))
    locs = list(chain(*(x.span() for x in secs)))
    # print(locs)
    d = {}

    for sec in secs:
        gdict = sec.groupdict()
        section_name = gdict['section']
        sec_begin = sec.end()
        idxend = locs.index(sec_begin) + 1
        # print('idxend = ', idxend)
        try:
            sec_end = locs[idxend]
            d[section_name] = man_page[sec_begin:sec_end]
        except IndexError:
            d[section_name] = man_page[sec_begin:]

    return d

def build_pattern():
    from collections import OrderedDict
    patterns = OrderedDict()
    patterns['flag'] = '-\w+'
    patterns['arg'] = '\s\w+'
    patterns['info'] = '\s+\w+(.|\n)+'

    patts = OrderedDict()
    for k,v in patterns.items():
        patts[k] = '(?P<{}>'.format(k) + v + ')'

    pat = '^' + patts['flag'] + patts['arg'] + '*' + patts['info']

    pc = re.compile(pat, flags=re.MULTILINE)
    return pc

def parse_options(man_page):
    pc = build_pattern()
    def search_str(txt):
        res = pc.search(txt)
        gdict = res.groupdict()
        # print(gdict)
        d = {}
        for k,v in gdict.items():
            if v is not None:
                d[k] = v.strip().replace('\n', ' ')
            else:
                d[k] = v
        return d

    txt = get_sections(man_page)['OPTIONS']
    opts = re.split('^\n', txt, flags=re.MULTILINE)
    opts = (x.strip() for x in opts if x)
    opts = (search_str(x) for x in opts)
    return list(opts)

def make_attribute(option):
    if option['arg'] is None:
        klass = Flag
    else:
        klass = Option
    flag = option['flag']
    info_text = option['info']
    return klass(flag, info_text=info_text)

if __name__ == '__main__':
    import logging
    logr = logging.getLogger(__name__)
    logr.setLevel(logging.DEBUG)

    opts = parse_options(man_dmenu())
    opts = [make_attribute(x) for x in opts]
    dmenu_src = MenuType('DMenu', *opts)

    try:
        dmenu_src.create_class()
    except:
        logr.exception("Couldn't create DMenu class")

    source = str(dmenu_src)
    try:
        from yapf.yapflib.yapf_api import FormatCode
        source, changed = FormatCode(
            source,
            style_config='pep8',
        )
    except ImportError:
        logr.exception("Couldn't import the yapf code formatter!")

    print(source)

<|endoftext|>


<commit_before>#!/usr/bin/env python
"""Prepare GFF transcript files for use as input to Cufflinks.

This supports running RNA-seq pipelines by providing reference annotation and
mask files needed by Cufflinks, ready for use with UCSC named genomes:

http://cufflinks.cbcb.umd.edu/manual.html

Usage:
  prepare_tx_gff.py <org_build>
"""
import os
import sys
import shutil
import platform
import collections
import subprocess

import MySQLdb
import rpy2.robjects as robjects

from bcbio.utils import chdir, safe_makedir

def build_ucsc_map(ensembl_chrs, mt_chr):
    """Build mapping of Ensembl to UCSC names for standard chromosomes.
    """
    out = {}
    for c in ensembl_chrs:
        out[str(c)] = "chr{0}".format(c)
    out[mt_chr] = "chrM"
    return out

# ##  Version and retrieval details for Ensembl and UCSC
ensembl_release = "63"
base_ftp = "ftp://ftp.ensembl.org/pub/release-{release}/gtf"

Build = collections.namedtuple("Build", ["taxname", "fname", "biomart_name",
                                         "ucsc_map"])
build_info = {
    "hg19": Build("homo_sapiens", "Homo_sapiens.GRCh37.{release}.gtf.gz",
                  "hsapiens_gene_ensembl", None),
    "mm9": Build("mus_musculus", "Mus_musculus.NCBIM37.{release}.gtf.gz",
                 "mmusculus_gene_ensembl",
                 build_ucsc_map(range(1, 20) + ["X", "Y"], "MT"))}

ucsc_db= "genome-mysql.cse.ucsc.edu"
ucsc_user="genome"

# ## Main driver functions

def main(org_build):
    work_dir = os.path.join(os.getcwd(), "tmpcbl")
    safe_makedir(work_dir)
    with chdir(work_dir):
        build = build_info[org_build]
        tx_gff = prepare_tx_gff(build, org_build)
        mask_gff = prepare_mask_gff(tx_gff, build)
        upload_to_s3([tx_gff, mask_gff], org_build)

def upload_to_s3(fnames, org_build):
    final_dir = os.path.join(org_build, "rnaseq")
    final_tarball = "{org}-rnaseq.tar.xz".format(org=org_build)
    if not os.path.exists(final_tarball):
        safe_makedir(final_dir)
        for fname in fnames:
            shutil.move(fname, final_dir)
        subprocess.check_call("tar -cvpf - {dir} | xz -zc - > {tarball}".format(
            dir=final_dir, tarball=final_tarball), shell=True)
    python_exe = "python{0}.{1}".format(*platform.python_version_tuple()[:2])
    upload_script = os.path.join(os.path.dirname(__file__), "s3_multipart_upload.py")
    subprocess.check_call([python_exe, upload_script, final_tarball, "biodata",
                           os.path.join("annotation", os.path.basename(final_tarball)),
                           "--public"])

# ## Get transcripts to filter from Ensembl BioMart query

def prepare_mask_gff(base_gff, build):
    """Prepare GFF file with high abundance transcripts to mask.
    """
    out_file = "{0}-mask{1}".format(*os.path.splitext(base_gff))
    if not os.path.exists(out_file):
        tx_ids = high_abudance_transcripts(build)
        with open(base_gff) as in_handle, \
             open(out_file, "w") as out_handle:
            for line in in_handle:
                tx = line.split('transcript_id "')[1].split('"')[0]
                if tx in tx_ids:
                    out_handle.write(line)
    return out_file

def high_abudance_transcripts(build):
    """Get identifiers for high abundance transcriptions using R biomaRt.
    """
    robjects.r.assign("biomart.org", build.biomart_name)
    robjects.r('''
      library(biomaRt)
      mart <- useMart("ensembl", dataset=biomart.org)
      attrs <- c("ensembl_transcript_id")
      filters <- c("biotype")
      filter.types <- c("Mt_rRNA", "rRNA")
      result <- getBM(attributes=attrs, filters=filters, values=filter.types,
                      mart=mart)
      result <- unique(result)
    ''')
    # get first column of data frame: the transcript IDs
    return set(robjects.r["result"][0])

# ## Retrieve GFF file from Ensembl and map to UCSC coordinates

def prepare_tx_gff(build, org_name):
    """Prepare UCSC ready transcript file given build information.
    """
    ensembl_gff = _download_ensembl_gff(build)
    ucsc_name_map = (_query_for_ucsc_ensembl_map(org_name)
                     if build.ucsc_map is None else build.ucsc_map)
    return _remap_gff(ensembl_gff, ucsc_name_map)

def _remap_gff(base_gff, name_map):
    """Remap chromosome names to UCSC instead of Ensembl
    """
    out_file = "ref-transcripts.gtf"
    if not os.path.exists(out_file):
        with open(out_file, "w") as out_handle, \
             open(base_gff) as in_handle:
            for line in in_handle:
                parts = line.split("\t")
                ucsc_name = name_map.get(parts[0], None)
                if ucsc_name:
                    out_handle.write("\t".join([ucsc_name] + parts[1:]))
    return out_file

def _query_for_ucsc_ensembl_map(org_name):
    """Retrieve UCSC to Ensembl name mappings from UCSC MySQL database.
    """
    db = MySQLdb.connect(host=ucsc_db, user=ucsc_user, db=org_name)
    cursor = db.cursor()
    cursor.execute("select * from ucscToEnsembl")
    ucsc_map = {}
    for ucsc, ensembl in cursor.fetchall():
        ucsc_map[ensembl] = ucsc
    return ucsc_map

def _download_ensembl_gff(build):
    """Given build details, download and extract the relevant ensembl GFF.
    """
    dl_url = "/".join([base_ftp, build.taxname, build.fname]).format(release=ensembl_release)
    out_file = os.path.splitext(os.path.basename(dl_url))[0]
    if not os.path.exists(out_file):
        subprocess.check_call(["wget", dl_url])
        subprocess.check_call(["gunzip", os.path.basename(dl_url)])
    return out_file

if __name__ == "__main__":
    main(*sys.argv[1:])
<commit_msg>Only include standard chromosomes in transcript files. TopHat gtf_to_fasta has problems with random and haplotype chromosomes. Fixes #38, thanks to Paul Tang<commit_after>#!/usr/bin/env python
"""Prepare GFF transcript files for use as input to Cufflinks.

This supports running RNA-seq pipelines by providing reference annotation and
mask files needed by Cufflinks, ready for use with UCSC named genomes:

http://cufflinks.cbcb.umd.edu/manual.html

Usage:
  prepare_tx_gff.py <org_build>
"""
import os
import sys
import shutil
import platform
import collections
import subprocess

import MySQLdb
import rpy2.robjects as robjects

from bcbio.utils import chdir, safe_makedir

def build_ucsc_map(ensembl_chrs, mt_chr):
    """Build mapping of Ensembl to UCSC names for standard chromosomes.
    """
    out = {}
    for c in ensembl_chrs:
        out[str(c)] = "chr{0}".format(c)
    out[mt_chr] = "chrM"
    return out

# ##  Version and retrieval details for Ensembl and UCSC
ensembl_release = "66"
base_ftp = "ftp://ftp.ensembl.org/pub/release-{release}/gtf"

Build = collections.namedtuple("Build", ["taxname", "fname", "biomart_name",
                                         "ucsc_map"])
build_info = {
    "hg19": Build("homo_sapiens", "Homo_sapiens.GRCh37.{release}.gtf.gz",
                  "hsapiens_gene_ensembl", None),
    "mm9": Build("mus_musculus", "Mus_musculus.NCBIM37.{release}.gtf.gz",
                 "mmusculus_gene_ensembl",
                 build_ucsc_map(range(1, 20) + ["X", "Y"], "MT"))}

ucsc_db= "genome-mysql.cse.ucsc.edu"
ucsc_user="genome"

# ## Main driver functions

def main(org_build):
    work_dir = os.path.join(os.getcwd(), "tmpcbl")
    safe_makedir(work_dir)
    with chdir(work_dir):
        build = build_info[org_build]
        tx_gff = prepare_tx_gff(build, org_build)
        mask_gff = prepare_mask_gff(tx_gff, build)
        upload_to_s3([tx_gff, mask_gff], org_build)

def upload_to_s3(fnames, org_build):
    final_dir = os.path.join(org_build, "rnaseq")
    final_tarball = "{org}-rnaseq.tar.xz".format(org=org_build)
    if not os.path.exists(final_tarball):
        safe_makedir(final_dir)
        for fname in fnames:
            shutil.move(fname, final_dir)
        subprocess.check_call("tar -cvpf - {dir} | xz -zc - > {tarball}".format(
            dir=final_dir, tarball=final_tarball), shell=True)
    python_exe = "python{0}.{1}".format(*platform.python_version_tuple()[:2])
    upload_script = os.path.join(os.path.dirname(__file__), "s3_multipart_upload.py")
    subprocess.check_call([python_exe, upload_script, final_tarball, "biodata",
                           os.path.join("annotation", os.path.basename(final_tarball)),
                           "--public"])

# ## Get transcripts to filter from Ensembl BioMart query

def prepare_mask_gff(base_gff, build):
    """Prepare GFF file with high abundance transcripts to mask.
    """
    out_file = "{0}-mask{1}".format(*os.path.splitext(base_gff))
    if not os.path.exists(out_file):
        tx_ids = high_abudance_transcripts(build)
        with open(base_gff) as in_handle, \
             open(out_file, "w") as out_handle:
            for line in in_handle:
                tx = line.split('transcript_id "')[1].split('"')[0]
                if tx in tx_ids:
                    out_handle.write(line)
    return out_file

def high_abudance_transcripts(build):
    """Get identifiers for high abundance transcriptions using R biomaRt.
    """
    robjects.r.assign("biomart.org", build.biomart_name)
    robjects.r('''
      library(biomaRt)
      mart <- useMart("ensembl", dataset=biomart.org)
      attrs <- c("ensembl_transcript_id")
      filters <- c("biotype")
      filter.types <- c("Mt_rRNA", "rRNA")
      result <- getBM(attributes=attrs, filters=filters, values=filter.types,
                      mart=mart)
      result <- unique(result)
    ''')
    # get first column of data frame: the transcript IDs
    return set(robjects.r["result"][0])

# ## Retrieve GFF file from Ensembl and map to UCSC coordinates

def prepare_tx_gff(build, org_name):
    """Prepare UCSC ready transcript file given build information.
    """
    ensembl_gff = _download_ensembl_gff(build)
    ucsc_name_map = (_query_for_ucsc_ensembl_map(org_name)
                     if build.ucsc_map is None else build.ucsc_map)
    return _remap_gff(ensembl_gff, ucsc_name_map)

def _remap_gff(base_gff, name_map):
    """Remap chromosome names to UCSC instead of Ensembl
    """
    out_file = "ref-transcripts.gtf"
    if not os.path.exists(out_file):
        with open(out_file, "w") as out_handle, \
             open(base_gff) as in_handle:
            for line in in_handle:
                parts = line.split("\t")
                ucsc_name = name_map.get(parts[0], None)
                if ucsc_name:
                    out_handle.write("\t".join([ucsc_name] + parts[1:]))
    return out_file

def _query_for_ucsc_ensembl_map(org_name):
    """Retrieve UCSC to Ensembl name mappings from UCSC MySQL database.
    """
    db = MySQLdb.connect(host=ucsc_db, user=ucsc_user, db=org_name)
    cursor = db.cursor()
    cursor.execute("select * from ucscToEnsembl")
    ucsc_map = {}
    for ucsc, ensembl in cursor.fetchall():
        # only include standard chromosomes -- not haplotype and random
        # gtf_to_fasta in TopHat doesn't seem to like these
        if len(ucsc) < 6:
            ucsc_map[ensembl] = ucsc
    return ucsc_map

def _download_ensembl_gff(build):
    """Given build details, download and extract the relevant ensembl GFF.
    """
    dl_url = "/".join([base_ftp, build.taxname, build.fname]).format(release=ensembl_release)
    out_file = os.path.splitext(os.path.basename(dl_url))[0]
    if not os.path.exists(out_file):
        subprocess.check_call(["wget", dl_url])
        subprocess.check_call(["gunzip", os.path.basename(dl_url)])
    return out_file

if __name__ == "__main__":
    main(*sys.argv[1:])
<|endoftext|>


<commit_before>import httplib2

from apiclient.discovery import build
from apiclient.http import BatchHttpRequest
from oauth2client.client import flow_from_clientsecrets, OAuth2Credentials
from oauth2client.tools import run

from workflow import Workflow, PasswordNotFound

import config

EMAIL_LIST = []

class PseudoStorage():

    def put(self, value):
        pass

def list_threads(request_id, response, exception):
    if exception is None:
        thread = {
            'Date': None,
            'From': None,
            'id': None,
            'messages_count': None,
            'snippet': None,
            'Subject': None,
            'threadId': None,
            'unread': None,
        }
        if 'messages' in response and len(response['messages']) > 0:
            if 'id' in response['messages'][-1]:
                thread['id'] = response['messages'][-1]['id']

            if 'threadId' in response['messages'][-1]:
                thread['threadId'] = response['messages'][-1]['threadId']

            if 'snippet' in response['messages'][-1]:
                thread['snippet'] = response[
                    'messages'][-1]['snippet'].encode('utf-8')

            for header in response['messages'][-1]['payload']['headers']:
                if header['name'] in thread:
                    thread[header['name']] = header[
                        'value'].encode('utf-8')

            thread['messages_count'] = len(response['messages'])
            thread['unread'] = 'UNREAD' in response['messages'][-1]['labelIds']

        EMAIL_LIST.append(thread)

def get_list(wf, http, service):
    # Retrieve a page of threads
    threads = service.users().threads().list(
        userId='me', labelIds=['INBOX'], maxResults=100).execute()

    batch = BatchHttpRequest()
    if threads['threads']:
        fields = 'messages/id,messages/threadId,messages/labelIds,messages/snippet,messages/payload/headers'
        for thread in threads['threads']:
            batch.add(service.users().threads().get(
                userId='me', id=thread['id'], fields=fields), callback=list_threads)

    batch.execute(http=http)

    return EMAIL_LIST

def refresh_cache():
    wf = Workflow()
    # Start the OAuth flow to retrieve credentials
    flow = flow_from_clientsecrets(
        config.CLIENT_SECRET_FILE, scope=config.OAUTH_SCOPE)
    http = httplib2.Http()

    try:
        credentials = OAuth2Credentials.from_json(
            wf.get_password('gmail_credentials'))
        if credentials is None or credentials.invalid:
            credentials = run(flow, PseudoStorage(), http=http)
            wf.save_password('gmail_credentials', credentials.to_json())
            wf.logger.debug('Credentials securely updated')

        # Authorize the httplib2.Http object with our credentials
        http = credentials.authorize(http)
        # Build the Gmail service from discovery
        gmail_service = build('gmail', 'v1', http=http)

        def wrapper():
            return get_list(wf, http, gmail_service)

        wf.cached_data('gmail_list', data_func=wrapper, max_age=10)

    except PasswordNotFound:
        wf.logger.debug('Credentials not found')
        credentials = run(flow, PseudoStorage(), http=http)
        wf.save_password('gmail_credentials', credentials.to_json())
        wf.logger.debug('New Credentials securely saved')

if __name__ == '__main__':
    refresh_cache()
<commit_msg>Bugfixes and improvements<commit_after>import httplib2

from apiclient.discovery import build
from apiclient.http import BatchHttpRequest
from oauth2client.client import flow_from_clientsecrets, OAuth2Credentials
from oauth2client.tools import run

from workflow import Workflow, PasswordNotFound

import config

EMAIL_LIST = []

class PseudoStorage():

    def put(self, value):
        pass

def list_threads(request_id, response, exception):
    if exception is None:
        thread = {
            'Date': None,
            'From': None,
            'id': None,
            'messages_count': None,
            'snippet': None,
            'Subject': None,
            'threadId': None,
            'unread': None,
        }
        if 'messages' in response and len(response['messages']) > 0:
            if 'id' in response['messages'][-1]:
                thread['id'] = response['messages'][-1]['id']

            if 'threadId' in response['messages'][-1]:
                thread['threadId'] = response['messages'][-1]['threadId']

            if 'snippet' in response['messages'][-1]:
                thread['snippet'] = response[
                    'messages'][-1]['snippet'].encode('utf-8')

            for header in response['messages'][-1]['payload']['headers']:
                if header['name'] in thread:
                    thread[header['name']] = header[
                        'value'].encode('utf-8')

            thread['messages_count'] = len(response['messages'])
            thread['unread'] = 'UNREAD' in response['messages'][-1]['labelIds']

        EMAIL_LIST.append(thread)

def get_list(wf, http, service):
    # Retrieve a page of threads
    threads = service.users().threads().list(
        userId='me', labelIds=['INBOX'], maxResults=100).execute()

    batch = BatchHttpRequest()
    if 'threads' in threads and len(threads['threads']) > 0:
        fields = 'messages/id,messages/threadId,messages/labelIds,messages/snippet,messages/payload/headers'
        for thread in threads['threads']:
            batch.add(service.users().threads().get(
                userId='me', id=thread['id'], fields=fields), callback=list_threads)

        batch.execute(http=http)

    return EMAIL_LIST

def refresh_cache():
    wf = Workflow()
    # Start the OAuth flow to retrieve credentials
    flow = flow_from_clientsecrets(
        config.CLIENT_SECRET_FILE, scope=config.OAUTH_SCOPE)
    http = httplib2.Http()

    try:
        credentials = OAuth2Credentials.from_json(
            wf.get_password('gmail_credentials'))
        if credentials is None or credentials.invalid:
            credentials = run(flow, PseudoStorage(), http=http)
            wf.save_password('gmail_credentials', credentials.to_json())
            wf.logger.debug('Credentials securely updated')

        # Authorize the httplib2.Http object with our credentials
        http = credentials.authorize(http)
        # Build the Gmail service from discovery
        gmail_service = build('gmail', 'v1', http=http)

        wf.cache_data('gmail_list', get_list(wf, http, gmail_service))

    except PasswordNotFound:
        wf.logger.debug('Credentials not found')
        credentials = run(flow, PseudoStorage(), http=http)
        wf.save_password('gmail_credentials', credentials.to_json())
        wf.logger.debug('New Credentials securely saved')

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


<commit_before># -*- coding: utf-8 -*-
# Copyright (c) 2015, Vispy Development Team.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.

import os
import sys
import numpy as np
from nose.tools import assert_true, assert_equal
import warnings

from vispy.testing import assert_in, run_tests_if_main
from vispy.util.fonts import list_fonts, _load_glyph, _vispy_fonts

warnings.simplefilter('always')

def test_font_list():
    """Test font listing"""
    f = list_fonts()
    assert_true(len(f) > 0)
    for font in _vispy_fonts:
        assert_in(font, f)

def test_font_glyph():
    """Test loading glyphs"""
    # try both a vispy and system font
    sys_fonts = set(list_fonts()) - set(_vispy_fonts)
    assert_true(len(sys_fonts) > 0)
    for face in ['OpenSans'] + sorted(list(sys_fonts)):
        font_dict = dict(face=face, size=12, bold=False, italic=False)
        glyphs_dict = dict()
        chars = 'foobar^C&#'
        for char in chars:
            # Warning that Arial might not exist
            with warnings.catch_warnings(record=True):
                warnings.simplefilter('always')
                _load_glyph(font_dict, char, glyphs_dict)
        assert_equal(len(glyphs_dict), np.unique([c for c in chars]).size)

run_tests_if_main()
<commit_msg>FIX: Better msg<commit_after># -*- coding: utf-8 -*-
# Copyright (c) 2015, Vispy Development Team.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.

import numpy as np
from nose.tools import assert_true, assert_equal
import warnings

from vispy.testing import assert_in, run_tests_if_main
from vispy.util.fonts import list_fonts, _load_glyph, _vispy_fonts

warnings.simplefilter('always')

def test_font_list():
    """Test font listing"""
    f = list_fonts()
    assert_true(len(f) > 0)
    for font in _vispy_fonts:
        assert_in(font, f)

def test_font_glyph():
    """Test loading glyphs"""
    # try both a vispy and system font
    sys_fonts = set(list_fonts()) - set(_vispy_fonts)
    assert_true(len(sys_fonts) > 0)
    for face in ['OpenSans'] + sorted(list(sys_fonts)):
        print(face)  # useful for debugging
        font_dict = dict(face=face, size=12, bold=False, italic=False)
        glyphs_dict = dict()
        chars = 'foobar^C&#'
        for char in chars:
            # Warning that Arial might not exist
            with warnings.catch_warnings(record=True):
                warnings.simplefilter('always')
                _load_glyph(font_dict, char, glyphs_dict)
        assert_equal(len(glyphs_dict), np.unique([c for c in chars]).size)

run_tests_if_main()
<|endoftext|>


<commit_before>from django.template import Context
from django.http import Http404, HttpResponse

from core.models import Engine, TrackingForm
from core.decorators import json_response
from core.mail import send_mail

import simplejson

@json_response
def ajax_error(error):

    """Returns a HttpResponse with JSON as a payload

    This function is a simple way of instantiating an error when using
    json_functions. It is decorated with the json_response decorator so that
    the dict that we return is dumped into a json object.

    :param error: :class:`str` which contains the pretty
    error, this will be seen by the user so
    make sure it's understandable.
    :returns: :class:`HttpResponse` with mime/application of json.
    :rtype: :class:`HttpResponse`
    """

    return {
        'success': False,
        'error': str(error)
        }

@json_response
def ajax_multitracking(request):
    json_data = {
        'success': False,
        'error': ''
        }

    try:
        form_data = simplejson.loads(request.POST.get('mass_data'))
        for item in form_data:
            TrackingForm(item).save()

    except Exception as err:
        json_data['error'] = str(err)
        return json_data

    TrackingForm()
    json_data['success'] = True
    return json_data

@json_response
def edit_item(request):
    json_data = {
        'success': False,
        'error': ''
        }

    try:
        object_id = request.POST.get('id')
        db_item = Engine.objects.get(pk=object_id)

    except Exception as err:
        json_data['error'] = str(err)
        return json_data

    json_data['success'] = True
    json_data['itemid'] = object_id
    json_data['market'] = str(db_item.market)
    json_data['ccode'] = str(db_item.ccode)
    json_data['level'] = str(db_item.level)
    json_data['reminderdate'] = str(db_item.reminderdate)
    json_data['remindernumber'] = str(db_item.remindernumber)
    json_data['vendor'] = str(db_item.vendor)
    json_data['mailvendor'] = str(db_item.mailvendor)
    json_data['invoicenumber'] = str(db_item.invoicenumber)
    json_data['invoicestatus'] = str(db_item.invoicestatus)
    json_data['rejectreason'] = str(db_item.rejectreason)
    json_data['paidon'] = str(db_item.paidon)

    return json_data

@json_response
def update_item(request):
    json_data = {
        'success': False,
        'error': ''
        }

    try:
        data = simplejson.loads(request.POST.get('mass_data'))
        object_id = data['itemid']
        db_item = Engine.objects.get(pk=object_id)

    except Exception as err:
        json_data['error'] = str(err)
        return json_data

    db_item.market = data['market']
    db_item.invoicenumber = data['invoicenumber']
    db_item.vendor = data['vendor']
    db_item.mailvendor = data['mailvendor']
    db_item.remindernumber = data['remindernumber']
    db_item.invoicestatus = data['invoicestatus']
    db_item.level = data['level']
    db_item.rejectreason = data['rejectreason']
    db_item.reminderdate = data['reminderdate']

    if data['paidon'] == "None":
        pass
    else:
        db_item.paidon = data['paidon']

    db_item.ccode = data['ccode']

    try:
        db_item.save()
        json_data['success'] = True

    except Exception as err:
        json_data['error'] = str(err)
        return json_data

    return json_data

@json_response
def send_info (request):
    #json_data = {
    #    'success': False,
    #    'error': ''
    #    }

    #mbody = simplejson.loads(request.POST.get('mailbody'))
    #item = simplejson.loads(request.POST.get('id'))

    #mailto = item.mailvendor

    #send_mail('Information rgearding reminder: %s' % item.remindernumber,
    #          mbody,
    #          'dunningteam@dl.com',
    #          mailto
    #          )
    pass
<commit_msg>Created the function to send mail.<commit_after>from django.template import Context
from django.http import Http404, HttpResponse

from core.models import Engine, TrackingForm
from core.decorators import json_response
from django.core.mail import send_mail
from django.core.mail import EmailMultiAlternatives

import simplejson

@json_response
def ajax_error(error):

    """Returns a HttpResponse with JSON as a payload

    This function is a simple way of instantiating an error when using
    json_functions. It is decorated with the json_response decorator so that
    the dict that we return is dumped into a json object.

    :param error: :class:`str` which contains the pretty
    error, this will be seen by the user so
    make sure it's understandable.
    :returns: :class:`HttpResponse` with mime/application of json.
    :rtype: :class:`HttpResponse`
    """

    return {
        'success': False,
        'error': str(error)
        }

@json_response
def ajax_multitracking(request):
    json_data = {
        'success': False,
        'error': ''
        }

    try:
        form_data = simplejson.loads(request.POST.get('mass_data'))
        for item in form_data:
            TrackingForm(item).save()

    except Exception as err:
        json_data['error'] = str(err)
        return json_data

    TrackingForm()
    json_data['success'] = True
    return json_data

@json_response
def edit_item(request):
    json_data = {
        'success': False,
        'error': ''
        }

    try:
        object_id = request.POST.get('id')
        db_item = Engine.objects.get(pk=object_id)

    except Exception as err:
        json_data['error'] = str(err)
        return json_data

    json_data['success'] = True
    json_data['itemid'] = object_id
    json_data['market'] = str(db_item.market)
    json_data['ccode'] = str(db_item.ccode)
    json_data['level'] = str(db_item.level)
    json_data['reminderdate'] = str(db_item.reminderdate)
    json_data['remindernumber'] = str(db_item.remindernumber)
    json_data['vendor'] = str(db_item.vendor)
    json_data['mailvendor'] = str(db_item.mailvendor)
    json_data['invoicenumber'] = str(db_item.invoicenumber)
    json_data['invoicestatus'] = str(db_item.invoicestatus)
    json_data['rejectreason'] = str(db_item.rejectreason)
    json_data['paidon'] = str(db_item.paidon)

    return json_data

@json_response
def update_item(request):
    json_data = {
        'success': False,
        'error': ''
        }

    try:
        data = simplejson.loads(request.POST.get('mass_data'))
        object_id = data['itemid']
        db_item = Engine.objects.get(pk=object_id)

    except Exception as err:
        json_data['error'] = str(err)
        return json_data

    db_item.market = data['market']
    db_item.invoicenumber = data['invoicenumber']
    db_item.vendor = data['vendor']
    db_item.mailvendor = data['mailvendor']
    db_item.remindernumber = data['remindernumber']
    db_item.invoicestatus = data['invoicestatus']
    db_item.level = data['level']
    db_item.rejectreason = data['rejectreason']
    db_item.reminderdate = data['reminderdate']

    if data['paidon'] == "None":
        pass
    else:
        db_item.paidon = data['paidon']

    db_item.ccode = data['ccode']

    try:
        db_item.save()
        json_data['success'] = True

    except Exception as err:
        json_data['error'] = str(err)
        return json_data

    return json_data

@json_response
def send_info(request):
    json_data = {
        'success': False,
        'error': ''
        }

    mbody = simplejson.loads(request.POST.get('mailbody'))

    itemid = simplejson.loads(request.POST.get('id'))

    dbitem = Engine.objects.get(pk=itemid)

    mailto = dbitem.mailvendor

    try:
        subject = 'Your subject %s: %s%s' % (dbitem.market, dbitem.market, dbitem.ccode)
        from_email = 'Your address'
        html_content = mbody
        msg = EmailMultiAlternatives(subject, from_email, ['promat85@gmail.com', mailto])
        msg.attach_alternative(html_content, "text/html")
        msg.attach_file('your/pathto/file')
        msg.send()
        #send_mail('Information regarding reminder: %s' % dbitem.remindernumber,
        #      mbody,
        #      'dunningteam@dl.com',
        #      [mailto]
        #      )
        json_data['success'] = True

    except Exception as err:
        json_data['error'] = str(err)
        return json_data

    return json_data
<|endoftext|>


<commit_before># Collection of experimental rhodopsin data sets

# http://peak.telecommunity.com/DevCenter/PythonEggs#accessing-package-resources
# foo_config = open(os.path.join(os.path.dirname(__file__),'foo.conf').read() # Old way

# from pkg_resources import resource_string # New way
# foo_config = resource_string(__name__, 'foo.conf')

from pkg_resources import resource_stream #, resource_string
import pickle

#from pyrho.dataSets import load_ChR2
# c.f. https://github.com/scikit-learn/scikit-learn/blob/master/sklearn/datasets/base.py

#def load(resource='test.txt'):
#    return resource_string(__name__, resource)

def loadChR2():
    """Return dictionary of ChR2 protocol data.
        With thanks to Nir Grossman, Juan Burrone and Matthew Grub"""
    return pickle.load(resource_stream(__name__, 'ChR2.pkl'))

'''
def loadpkl(pkl, path=None):
    # with, try, finaly etc...
    #import os
    pklFile = pkl+".pkl"
    if path is None:
        dataSet = resource_string(__name__, pklFile)

        cwd = os.getcwd()
        if pkl in cwd: ### Finish!!!
            pass #pklFile = pklFile
            #fh = open(pklFile, "rb")
        else:
            pklFile = os.path.join(dDir, pklFile)
            #fh = open(os.path.join(dDir, pklFile), "rb") #dDir+'expData'+".pkl"
    else:
        pklFile = os.path.join(path, pklFile)
    with open(pklFile, "rb") as fh :
        dataSet = pickle.load(fh)
    #fh.close()
    return dataSet
'''<commit_msg>Add ChR2 data description<commit_after># Collection of experimental rhodopsin data sets

# http://peak.telecommunity.com/DevCenter/PythonEggs#accessing-package-resources
# foo_config = open(os.path.join(os.path.dirname(__file__),'foo.conf').read() # Old way

# from pkg_resources import resource_string # New way
# foo_config = resource_string(__name__, 'foo.conf')

from pkg_resources import resource_stream #, resource_string
import pickle

#from pyrho.dataSets import load_ChR2
# c.f. https://github.com/scikit-learn/scikit-learn/blob/master/sklearn/datasets/base.py

#def load(resource='test.txt'):
#    return resource_string(__name__, resource)

def loadChR2():
    """Return dictionary of ChR2 protocol data.
        The dataset is a dictionary containing three ProtocolData objects:
        'step'      : Six PhotoCurrents recorded over flux values from
                      2.21e15 to 2.65e17 photons/mm2/s to steady-state
        'delta'     : The PhotoCurrent from the 'step' set with the strongest peak
        'shortPulse': Ten PhotoCurrents recorded with pulse durations of
                      [1, 2, 3, 4, 5, 6, 8, 10, 20, 30] ms
        With thanks to Nir Grossman, Juan Burrone and Matthew Grub.
    """
    return pickle.load(resource_stream(__name__, 'ChR2.pkl'))

'''
def loadpkl(pkl, path=None):
    # with, try, finaly etc...
    #import os
    pklFile = pkl+".pkl"
    if path is None:
        dataSet = resource_string(__name__, pklFile)

        cwd = os.getcwd()
        if pkl in cwd: ### Finish!!!
            pass #pklFile = pklFile
            #fh = open(pklFile, "rb")
        else:
            pklFile = os.path.join(dDir, pklFile)
            #fh = open(os.path.join(dDir, pklFile), "rb") #dDir+'expData'+".pkl"
    else:
        pklFile = os.path.join(path, pklFile)
    with open(pklFile, "rb") as fh :
        dataSet = pickle.load(fh)
    #fh.close()
    return dataSet
'''<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from django.test import TestCase
from django.conf import settings
from models import GroupMigration, User
from vkontakte_users.factories import UserFactory
from vkontakte_groups.factories import GroupFactory
from factories import GroupMigrationFactory, GroupMembership
from datetime import datetime, timedelta

GROUP_ID = 30221121

class VkontakteGroupsMigrationTest(TestCase):

    def test_user_memberships(self):

        migration1 = GroupMigrationFactory(time=datetime.now() - timedelta(3), members_ids=range(30, 100))
        migration1.update()
        migration1.save()

        migration2 = GroupMigrationFactory(group=migration1.group, time=datetime.now() - timedelta(2), members_ids=range(0, 50))
        migration2.update()
        migration2.save()

        migration3 = GroupMigrationFactory(group=migration1.group, time=datetime.now() - timedelta(1), members_ids=range(30, 110))
        migration3.update()
        migration3.save()

        migration4 = GroupMigrationFactory(group=migration1.group, members_ids=range(15, 100))
        migration4.update()
        migration4.save()

        def membership(id):
            return GroupMembership.objects.get(user_id=id)

        def memberships(id):
            return GroupMembership.objects.filter(user_id=id)

        def id90_state1():
            self.assertEqual(membership(90).time_entered, None)
            self.assertEqual(membership(90).time_left, None)

        def id90_state2():
            self.assertEqual(membership(90).time_entered, None)
            self.assertEqual(membership(90).time_left, migration2.time)

        def id90_state3():
            self.assertEqual(memberships(90)[0].time_entered, None)
            self.assertEqual(memberships(90)[0].time_left, migration2.time)
            self.assertEqual(memberships(90)[1].time_entered, migration3.time)
            self.assertEqual(memberships(90)[1].time_left, None)

        def id90_state4():
            self.assertEqual(memberships(90)[0].time_entered, None)
            self.assertEqual(memberships(90)[0].time_left, migration2.time)
            self.assertEqual(memberships(90)[1].time_entered, migration3.time)
            self.assertEqual(memberships(90)[1].time_left, None)

        def id90_state4_corrected():
            self.assertEqual(memberships(90)[0].time_entered, None)
            self.assertEqual(memberships(90)[0].time_left, migration2.time)
            self.assertEqual(memberships(90)[1].time_entered, migration4.time)
            self.assertEqual(memberships(90)[1].time_left, None)

        def id0_state1():
            self.assertEqual(memberships(0).count(), 0)

        def id0_state2():
            self.assertEqual(membership(0).time_entered, migration2.time)
            self.assertEqual(membership(0).time_left, None)

        def id0_state3():
            self.assertEqual(membership(0).time_entered, migration2.time)
            self.assertEqual(membership(0).time_left, migration3.time)

        def id0_state3_corrected():
            self.assertEqual(membership(0).time_entered, migration2.time)
            self.assertEqual(membership(0).time_left, migration4.time)

        def id20_state1():
            self.assertEqual(memberships(20).count(), 0)

        def id20_state2():
            self.assertEqual(membership(20).time_entered, migration2.time)
            self.assertEqual(membership(20).time_left, None)

        def id20_state3():
            self.assertEqual(membership(20).time_entered, migration2.time)
            self.assertEqual(membership(20).time_left, migration3.time)

        def id20_state4():
            self.assertEqual(memberships(20)[0].time_entered, migration2.time)
            self.assertEqual(memberships(20)[0].time_left, migration3.time)
            self.assertEqual(memberships(20)[1].time_entered, migration4.time)
            self.assertEqual(memberships(20)[1].time_left, None)

        def id40_state1():
            self.assertEqual(membership(40).time_entered, None)
            self.assertEqual(membership(40).time_left, None)

        def id40_state3():
            self.assertEqual(membership(40).time_entered, None)
            self.assertEqual(membership(40).time_left, None)

        def id105_state1():
            self.assertEqual(memberships(105).count(), 0)

        def id105_state3():
            self.assertEqual(membership(105).time_entered, migration3.time)
            self.assertEqual(membership(105).time_left, None)

        def id105_state4():
            self.assertEqual(membership(105).time_entered, migration3.time)
            self.assertEqual(membership(105).time_left, migration4.time)

        migration1.update_users_memberships()
        id0_state1()
        id20_state1()
        id40_state1()
        id90_state1()
        id105_state1()

        migration2.update_users_memberships()
        id0_state2()
        id20_state2()
        id40_state1()
        id90_state2()
        id105_state1()

        migration3.update_users_memberships()
        id0_state3()
        id20_state3()
        id40_state3()
        id90_state3()
        id105_state3()

        migration4.update_users_memberships()
        id0_state3()
        id20_state4()
        id40_state3()
        id90_state4()
        id105_state4()

        # hide migration
        migration3.hide()
        id0_state3_corrected()
        id20_state2()
        id40_state1()
        id90_state4_corrected()
        id105_state1()

    def test_comparing_with_statistic(self):

        if 'vkontakte_groups_statistic' not in settings.INSTALLED_APPS:
            return False

        from vkontakte_groups_statistic.factories import GroupStatFactory

        migration = GroupMigrationFactory(members_ids=range(0, 10000))
        migration.update()
        migration.save()

        stat = GroupStatFactory(group=migration.group, members=9900, date=migration.time.date())

        self.assertEqual(migration.hidden, False)
        migration.compare_with_statistic()
        self.assertEqual(migration.hidden, True)

        stat.members = 9901
        stat.save()
        migration.hidden = False

        self.assertEqual(migration.hidden, False)
        migration.compare_with_statistic()
        self.assertEqual(migration.hidden, False)

    def test_comparing_with_previous(self):

        migration1 = GroupMigrationFactory(time=datetime.now() - timedelta(2), members_ids=range(0, 10000))
        migration1.update()
        migration1.save()

        migration2 = GroupMigrationFactory(group=migration1.group, time=datetime.now() - timedelta(1), members_ids=range(0, 9000))
        migration2.update()
        migration2.save()

        self.assertEqual(migration2.hidden, False)
        migration2.compare_with_previous()
        self.assertEqual(migration2.hidden, True)

        migration3 = GroupMigrationFactory(group=migration1.group, members_ids=range(0, 9001))
        migration3.update()
        migration3.save()

        self.assertEqual(migration3.hidden, False)
        migration3.compare_with_previous()
        self.assertEqual(migration3.hidden, False)

    def test_m2m_relations(self):

        [UserFactory(remote_id=i, fetched=datetime.now()) for i in range(0, 2000)]
        migration = GroupMigrationFactory(members_ids=list(range(1000, 2000)))
        for i in range(0, 1500):
            migration.group.users.add(User.objects.get(remote_id=i))

        self.assertListEqual(list(migration.group.users.values_list('remote_id', flat=True)), range(0, 1500))
        migration.update_users_relations()
        self.assertListEqual(list(migration.group.users.values_list('remote_id', flat=True)), range(1000, 2000))

    def test_deleting_hiding_migration(self):

        for i in range(1,7):
            UserFactory.create(remote_id=i)

        group = GroupFactory.create(remote_id=GROUP_ID)
        stat1 = GroupMigration.objects.create(group=group, members_ids=[1,2,3,4,5])
        stat1.save_final()
        stat2 = GroupMigration.objects.create(group=group, members_ids=[1,2,3,4,6])
        stat2.save_final()
        stat3 = GroupMigration.objects.create(group=group, members_ids=[1,2,3,5,7])
        stat3.save_final()

        # difference between stat2 and stat1
        self.assertItemsEqual(stat2.members_entered_ids, [6])
        self.assertItemsEqual(stat2.members_left_ids, [5])
        # difference between stat3 and stat2
        self.assertItemsEqual(stat3.members_entered_ids, [5,7])
        self.assertItemsEqual(stat3.members_left_ids, [4,6])

        stat2.delete()
        stat3 = GroupMigration.objects.get(id=stat3.id)

        # difference between stat3 and stat1
        self.assertItemsEqual(stat3.members_entered_ids, [7])
        self.assertItemsEqual(stat3.members_left_ids, [4])

        stat4 = GroupMigration.objects.create(group=group, members_ids=[1,2,3,4,6])
        stat4.save_final()

        # difference between stat4 and stat3
        self.assertItemsEqual(stat4.members_entered_ids, [4,6])
        self.assertItemsEqual(stat4.members_left_ids, [5,7])

        stat3.hide()
        stat4 = GroupMigration.objects.get(id=stat4.id)

        # difference between stat4 and stat1
        self.assertItemsEqual(stat4.members_entered_ids, [6])
        self.assertItemsEqual(stat4.members_left_ids, [5])

        stat5 = GroupMigration.objects.create(group=group, members_ids=[1,2,3,5,7])
        stat5.save_final()

        # difference between stat5 and stat4
        self.assertItemsEqual(stat5.members_entered_ids, [5,7])
        self.assertItemsEqual(stat5.members_left_ids, [4,6])

        stat4.hidden = True
        stat4.save()
        stat5 = GroupMigration.objects.get(id=stat5.id)

        # difference between stat5 and stat1
        self.assertItemsEqual(stat5.members_entered_ids, [7])
        self.assertItemsEqual(stat5.members_left_ids, [4])<commit_msg>enrich tests<commit_after># -*- coding: utf-8 -*-
from django.test import TestCase
from django.conf import settings
from models import GroupMigration, User
from vkontakte_users.factories import UserFactory
from vkontakte_groups.factories import GroupFactory
from factories import GroupMigrationFactory, GroupMembership
from datetime import datetime, timedelta

GROUP_ID = 30221121

class VkontakteGroupsMigrationTest(TestCase):

    def test_user_memberships(self):

        migration1 = GroupMigrationFactory(time=datetime.now() - timedelta(3), members_ids=range(30, 100))
        migration1.update()
        migration1.save()

        migration2 = GroupMigrationFactory(group=migration1.group, time=datetime.now() - timedelta(2), members_ids=range(0, 50))
        migration2.update()
        migration2.save()

        migration3 = GroupMigrationFactory(group=migration1.group, time=datetime.now() - timedelta(1), members_ids=range(30, 110))
        migration3.update()
        migration3.save()

        migration4 = GroupMigrationFactory(group=migration1.group, members_ids=range(15, 100))
        migration4.update()
        migration4.save()

        def membership(id):
            return GroupMembership.objects.get(user_id=id)

        def memberships(id):
            return GroupMembership.objects.filter(user_id=id)

        def id90_state1():
            self.assertEqual(membership(90).time_entered, None)
            self.assertEqual(membership(90).time_left, None)

        def id90_state2():
            self.assertEqual(membership(90).time_entered, None)
            self.assertEqual(membership(90).time_left, migration2.time)

        def id90_state3():
            self.assertEqual(memberships(90)[0].time_entered, None)
            self.assertEqual(memberships(90)[0].time_left, migration2.time)
            self.assertEqual(memberships(90)[1].time_entered, migration3.time)
            self.assertEqual(memberships(90)[1].time_left, None)

        def id90_state4():
            self.assertEqual(memberships(90)[0].time_entered, None)
            self.assertEqual(memberships(90)[0].time_left, migration2.time)
            self.assertEqual(memberships(90)[1].time_entered, migration3.time)
            self.assertEqual(memberships(90)[1].time_left, None)

        def id90_state4_corrected():
            self.assertEqual(memberships(90)[0].time_entered, None)
            self.assertEqual(memberships(90)[0].time_left, migration2.time)
            self.assertEqual(memberships(90)[1].time_entered, migration4.time)
            self.assertEqual(memberships(90)[1].time_left, None)

        def id0_state1():
            self.assertEqual(memberships(0).count(), 0)

        def id0_state2():
            self.assertEqual(membership(0).time_entered, migration2.time)
            self.assertEqual(membership(0).time_left, None)

        def id0_state3():
            self.assertEqual(membership(0).time_entered, migration2.time)
            self.assertEqual(membership(0).time_left, migration3.time)

        def id0_state3_corrected():
            self.assertEqual(membership(0).time_entered, migration2.time)
            self.assertEqual(membership(0).time_left, migration4.time)

        def id20_state1():
            self.assertEqual(memberships(20).count(), 0)

        def id20_state2():
            self.assertEqual(membership(20).time_entered, migration2.time)
            self.assertEqual(membership(20).time_left, None)

        def id20_state3():
            self.assertEqual(membership(20).time_entered, migration2.time)
            self.assertEqual(membership(20).time_left, migration3.time)

        def id20_state4():
            self.assertEqual(memberships(20)[0].time_entered, migration2.time)
            self.assertEqual(memberships(20)[0].time_left, migration3.time)
            self.assertEqual(memberships(20)[1].time_entered, migration4.time)
            self.assertEqual(memberships(20)[1].time_left, None)

        def id40_state1():
            self.assertEqual(membership(40).time_entered, None)
            self.assertEqual(membership(40).time_left, None)

        def id40_state3():
            self.assertEqual(membership(40).time_entered, None)
            self.assertEqual(membership(40).time_left, None)

        def id105_state1():
            self.assertEqual(memberships(105).count(), 0)

        def id105_state3():
            self.assertEqual(membership(105).time_entered, migration3.time)
            self.assertEqual(membership(105).time_left, None)

        def id105_state4():
            self.assertEqual(membership(105).time_entered, migration3.time)
            self.assertEqual(membership(105).time_left, migration4.time)

        migration1.update_users_memberships()
        self.assertEqual(GroupMembership.objects.count(), 70)
        id0_state1()
        id20_state1()
        id40_state1()
        id90_state1()
        id105_state1()

        migration2.update_users_memberships()
        self.assertEqual(GroupMembership.objects.count(), 100)
        id0_state2()
        id20_state2()
        id40_state1()
        id90_state2()
        id105_state1()

        migration3.update_users_memberships()
        self.assertEqual(GroupMembership.objects.count(), 160)
        id0_state3()
        id20_state3()
        id40_state3()
        id90_state3()
        id105_state3()

        migration4.update_users_memberships()
        self.assertEqual(GroupMembership.objects.count(), 175)
        id0_state3()
        id20_state4()
        id40_state3()
        id90_state4()
        id105_state4()

        # hide migration3
        migration3 = GroupMigration.objects.get(id=migration3.id)
        migration3.hide()
        self.assertEqual(GroupMembership.objects.count(), 150)
        id0_state3_corrected()
        id20_state2()
        id40_state1()
        id90_state4_corrected()
        id105_state1()

        # hide migration4 -> back to state2
        migration4 = GroupMigration.objects.get(id=migration4.id)
        migration4.hide()
        self.assertEqual(GroupMembership.objects.count(), 100)
        id0_state2()
        id20_state2()
        id40_state1()
        id90_state2()
        id105_state1()

        # hide migration2 -> back to state1
        migration2 = GroupMigration.objects.get(id=migration2.id)
        migration2.hide()
        self.assertEqual(GroupMembership.objects.count(), 70)
        id0_state1()
        id20_state1()
        id40_state1()
        id90_state1()
        id105_state1()

    def test_comparing_with_statistic(self):

        if 'vkontakte_groups_statistic' not in settings.INSTALLED_APPS:
            return False

        from vkontakte_groups_statistic.factories import GroupStatFactory

        migration = GroupMigrationFactory(members_ids=range(0, 10000))
        migration.update()
        migration.save()

        stat = GroupStatFactory(group=migration.group, members=9900, date=migration.time.date())

        self.assertEqual(migration.hidden, False)
        migration.compare_with_statistic()
        self.assertEqual(migration.hidden, True)

        stat.members = 9901
        stat.save()
        migration.hidden = False

        self.assertEqual(migration.hidden, False)
        migration.compare_with_statistic()
        self.assertEqual(migration.hidden, False)

    def test_comparing_with_previous(self):

        migration1 = GroupMigrationFactory(time=datetime.now() - timedelta(2), members_ids=range(0, 10000))
        migration1.update()
        migration1.save()

        migration2 = GroupMigrationFactory(group=migration1.group, time=datetime.now() - timedelta(1), members_ids=range(0, 9000))
        migration2.update()
        migration2.save()

        self.assertEqual(migration2.hidden, False)
        migration2.compare_with_previous()
        self.assertEqual(migration2.hidden, True)

        migration3 = GroupMigrationFactory(group=migration1.group, members_ids=range(0, 9001))
        migration3.update()
        migration3.save()

        self.assertEqual(migration3.hidden, False)
        migration3.compare_with_previous()
        self.assertEqual(migration3.hidden, False)

    def test_m2m_relations(self):

        [UserFactory(remote_id=i, fetched=datetime.now()) for i in range(0, 2000)]
        migration = GroupMigrationFactory(members_ids=list(range(1000, 2000)))
        for i in range(0, 1500):
            migration.group.users.add(User.objects.get(remote_id=i))

        self.assertListEqual(list(migration.group.users.values_list('remote_id', flat=True)), range(0, 1500))
        migration.update_users_relations()
        self.assertListEqual(list(migration.group.users.values_list('remote_id', flat=True)), range(1000, 2000))

    def test_deleting_hiding_migration(self):

        for i in range(1,7):
            UserFactory.create(remote_id=i)

        group = GroupFactory.create(remote_id=GROUP_ID)
        stat1 = GroupMigration.objects.create(group=group, members_ids=[1,2,3,4,5])
        stat1.save_final()
        stat2 = GroupMigration.objects.create(group=group, members_ids=[1,2,3,4,6])
        stat2.save_final()
        stat3 = GroupMigration.objects.create(group=group, members_ids=[1,2,3,5,7])
        stat3.save_final()

        # difference between stat2 and stat1
        self.assertItemsEqual(stat2.members_entered_ids, [6])
        self.assertItemsEqual(stat2.members_left_ids, [5])
        # difference between stat3 and stat2
        self.assertItemsEqual(stat3.members_entered_ids, [5,7])
        self.assertItemsEqual(stat3.members_left_ids, [4,6])

        stat2.delete()
        stat3 = GroupMigration.objects.get(id=stat3.id)

        # difference between stat3 and stat1
        self.assertItemsEqual(stat3.members_entered_ids, [7])
        self.assertItemsEqual(stat3.members_left_ids, [4])

        stat4 = GroupMigration.objects.create(group=group, members_ids=[1,2,3,4,6])
        stat4.save_final()

        # difference between stat4 and stat3
        self.assertItemsEqual(stat4.members_entered_ids, [4,6])
        self.assertItemsEqual(stat4.members_left_ids, [5,7])

        stat3.hide()
        stat4 = GroupMigration.objects.get(id=stat4.id)

        # difference between stat4 and stat1
        self.assertItemsEqual(stat4.members_entered_ids, [6])
        self.assertItemsEqual(stat4.members_left_ids, [5])

        stat5 = GroupMigration.objects.create(group=group, members_ids=[1,2,3,5,7])
        stat5.save_final()

        # difference between stat5 and stat4
        self.assertItemsEqual(stat5.members_entered_ids, [5,7])
        self.assertItemsEqual(stat5.members_left_ids, [4,6])

        stat4.hidden = True
        stat4.save()
        stat5 = GroupMigration.objects.get(id=stat5.id)

        # difference between stat5 and stat1
        self.assertItemsEqual(stat5.members_entered_ids, [7])
        self.assertItemsEqual(stat5.members_left_ids, [4])<|endoftext|>


<commit_before>from __future__ import absolute_import, division, print_function

import weakref

import torch
from torch.distributions import constraints, transform_to

class ParamStoreDict(object):
    """
    Global store for parameters in Pyro. This is basically a key-value store.
    The typical user interacts with the ParamStore primarily through the
    primitive `pyro.param`.

    See `Intro Part II <http://pyro.ai/examples/intro_part_ii.html>`_ for further discussion
    and `SVI Part I <http://pyro.ai/examples/svi_part_i.html>`_ for some examples.

    Some things to bear in mind when using parameters in Pyro:

    - parameters must be assigned unique names
    - the `init_tensor` argument to `pyro.param` is only used the first time that a given (named)
      parameter is registered with Pyro.
    - for this reason, a user may need to use the `clear()` method if working in a REPL in order to
      get the desired behavior. this method can also be invoked with `pyro.clear_param_store()`.
    - the internal name of a parameter within a PyTorch `nn.Module` that has been registered with
      Pyro is prepended with the Pyro name of the module. so nothing prevents the user from having
      two different modules each of which contains a parameter named `weight`. by contrast, a user
      can only have one top-level parameter named `weight` (outside of any module).
    - parameters can be saved and loaded from disk using `save` and `load`.
    """

    def __init__(self):
        """
        initialize ParamStore data structures
        """
        self._params = {}  # dictionary from param name to param
        self._param_to_name = {}  # dictionary from unconstrained param to param name
        self._constraints = {}  # dictionary from param name to constraint object

    def clear(self):
        """
        Clear the ParamStore
        """
        self._params = {}
        self._param_to_name = {}
        self._constraints = {}

    def named_parameters(self):
        """
        Returns an iterator over tuples of the form (name, parameter) for each parameter in the ParamStore
        """
        # TODO consider returing constrained
        return self._params.items()

    def get_all_param_names(self):
        """
        Get all parameter names in the ParamStore
        """
        return self._params.keys()

    def replace_param(self, param_name, new_param, old_param):
        """
        Replace the param param_name with current value old_param with the new value new_param

        :param param_name: parameter name
        :type param_name: str
        :param new_param: the paramater to be put into the ParamStore
        :type new_param: torch.Tensor
        :param old_param: the paramater to be removed from the ParamStore
        :type new_param: torch.Tensor
        """
        assert self._params[param_name] is old_param.unconstrained()
        del self._params[param_name]
        del self._param_to_name[old_param.unconstrained()]
        self.get_param(param_name, new_param, constraint=self._constraints[param_name])

    def get_param(self, name, init_tensor=None, constraint=constraints.real):
        """
        Get parameter from its name. If it does not yet exist in the
        ParamStore, it will be created and stored.
        The Pyro primitive `pyro.param` dispatches to this method.

        :param name: parameter name
        :type name: str
        :param init_tensor: initial tensor
        :type init_tensor: torch.Tensor
        :returns: parameter
        :rtype: torch.Tensor
        """
        if name not in self._params:
            # if not create the init tensor through
            assert init_tensor is not None,\
                "cannot initialize a parameter '{}' with None. Did you get the param name right?".format(name)

            # a function
            if callable(init_tensor):
                init_tensor = init_tensor()

            # store the unconstrained value and constraint
            with torch.no_grad():
                unconstrained_param = transform_to(constraint).inv(init_tensor)
            unconstrained_param.requires_grad_(True)
            self._params[name] = unconstrained_param
            self._constraints[name] = constraint

            # keep track of each tensor and it's name
            self._param_to_name[unconstrained_param] = name

        elif init_tensor is not None and not callable(init_tensor):
            if self._params[name].shape != init_tensor.shape:
                raise ValueError("param {} init tensor shape does not match existing value: {} vs {}".format(
                    name, init_tensor.shape, self._params[name].shape))

        # get the guaranteed to exist param
        unconstrained_param = self._params[name]

        # compute the constrained value
        param = transform_to(self._constraints[name])(unconstrained_param)
        param.unconstrained = weakref.ref(unconstrained_param)

        return param

    def param_name(self, p):
        """
        Get parameter name from parameter

        :param p: parameter
        :returns: parameter name
        """
        if p not in self._param_to_name:
            return None

        return self._param_to_name[p]

    def get_state(self):
        """
        Get the ParamStore state.
        """
        state = {
            'params': self._params,
            'constraints': self._constraints,
        }
        return state

    def set_state(self, state):
        """
        Set the ParamStore state using state from a previous get_state() call
        """
        assert isinstance(state, dict), "malformed ParamStore state"
        assert set(state.keys()) == set(['params', 'constraints']), \
            "malformed ParamStore keys {}".format(state.keys())

        for param_name, param in state['params'].items():
            self._params[param_name] = param
            self._param_to_name[param] = param_name

        for param_name, constraint in state['constraints'].items():
            if isinstance(constraint, type(constraints.real)):
                # Work around lack of hash & equality comparison on constraints.
                constraint = constraints.real
            self._constraints[param_name] = constraint

    def save(self, filename):
        """
        Save parameters to disk

        :param filename: file name to save to
        :type name: str
        """
        with open(filename, "wb") as output_file:
            torch.save(self.get_state(), output_file)

    def load(self, filename):
        """
        Loads parameters from disk

        .. note::

           If using :meth:`pyro.module` on parameters loaded from
           disk, be sure to set the ``update_module_params`` flag::

               pyro.get_param_store().load('saved_params.save')
               pyro.module('module', nn, update_module_params=True)

        :param filename: file name to load from
        :type name: str
        """
        with open(filename, "rb") as input_file:
            state = torch.load(input_file)
        self.set_state(state)

# used to create fully-formed param names, e.g. mymodule$$$mysubmodule.weight
_MODULE_NAMESPACE_DIVIDER = "$$$"

def param_with_module_name(pyro_name, param_name):
    return _MODULE_NAMESPACE_DIVIDER.join([pyro_name, param_name])

def module_from_param_with_module_name(param_name):
    return param_name.split(_MODULE_NAMESPACE_DIVIDER)[0]

def user_param_name(param_name):
    if _MODULE_NAMESPACE_DIVIDER in param_name:
        return param_name.split(_MODULE_NAMESPACE_DIVIDER)[1]
    return param_name
<commit_msg>add regex search support for param store (#1259)<commit_after>from __future__ import absolute_import, division, print_function

import re
import weakref

import torch
from torch.distributions import constraints, transform_to

class ParamStoreDict(object):
    """
    Global store for parameters in Pyro. This is basically a key-value store.
    The typical user interacts with the ParamStore primarily through the
    primitive `pyro.param`.

    See `Intro Part II <http://pyro.ai/examples/intro_part_ii.html>`_ for further discussion
    and `SVI Part I <http://pyro.ai/examples/svi_part_i.html>`_ for some examples.

    Some things to bear in mind when using parameters in Pyro:

    - parameters must be assigned unique names
    - the `init_tensor` argument to `pyro.param` is only used the first time that a given (named)
      parameter is registered with Pyro.
    - for this reason, a user may need to use the `clear()` method if working in a REPL in order to
      get the desired behavior. this method can also be invoked with `pyro.clear_param_store()`.
    - the internal name of a parameter within a PyTorch `nn.Module` that has been registered with
      Pyro is prepended with the Pyro name of the module. so nothing prevents the user from having
      two different modules each of which contains a parameter named `weight`. by contrast, a user
      can only have one top-level parameter named `weight` (outside of any module).
    - parameters can be saved and loaded from disk using `save` and `load`.
    """

    def __init__(self):
        """
        initialize ParamStore data structures
        """
        self._params = {}  # dictionary from param name to param
        self._param_to_name = {}  # dictionary from unconstrained param to param name
        self._constraints = {}  # dictionary from param name to constraint object

    def clear(self):
        """
        Clear the ParamStore
        """
        self._params = {}
        self._param_to_name = {}
        self._constraints = {}

    def named_parameters(self):
        """
        Returns an iterator over tuples of the form (name, parameter) for each parameter in the ParamStore
        """
        # TODO consider returing constrained
        return self._params.items()

    def get_all_param_names(self):
        """
        Get all parameter names in the ParamStore
        """
        return self._params.keys()

    def replace_param(self, param_name, new_param, old_param):
        """
        Replace the param param_name with current value old_param with the new value new_param

        :param param_name: parameter name
        :type param_name: str
        :param new_param: the paramater to be put into the ParamStore
        :type new_param: torch.Tensor
        :param old_param: the paramater to be removed from the ParamStore
        :type new_param: torch.Tensor
        """
        assert self._params[param_name] is old_param.unconstrained()
        del self._params[param_name]
        del self._param_to_name[old_param.unconstrained()]
        self.get_param(param_name, new_param, constraint=self._constraints[param_name])

    def get_param(self, name, init_tensor=None, constraint=constraints.real):
        """
        Get parameter from its name. If it does not yet exist in the
        ParamStore, it will be created and stored.
        The Pyro primitive `pyro.param` dispatches to this method.

        :param name: parameter name
        :type name: str
        :param init_tensor: initial tensor
        :type init_tensor: torch.Tensor
        :param constraint: torch constraint
        :type constraint: torch.distributions.constraints.Constraint
        :returns: parameter
        :rtype: torch.Tensor
        """
        if name not in self._params:
            # if not create the init tensor through
            assert init_tensor is not None,\
                "cannot initialize a parameter '{}' with None. Did you get the param name right?".format(name)

            # a function
            if callable(init_tensor):
                init_tensor = init_tensor()

            # store the unconstrained value and constraint
            with torch.no_grad():
                unconstrained_param = transform_to(constraint).inv(init_tensor)
            unconstrained_param.requires_grad_(True)
            self._params[name] = unconstrained_param
            self._constraints[name] = constraint

            # keep track of each tensor and it's name
            self._param_to_name[unconstrained_param] = name

        elif init_tensor is not None and not callable(init_tensor):
            if self._params[name].shape != init_tensor.shape:
                raise ValueError("param {} init tensor shape does not match existing value: {} vs {}".format(
                    name, init_tensor.shape, self._params[name].shape))

        # get the guaranteed to exist param
        unconstrained_param = self._params[name]

        # compute the constrained value
        param = transform_to(self._constraints[name])(unconstrained_param)
        param.unconstrained = weakref.ref(unconstrained_param)

        return param

    def match(self, name):
        """
        Get all parameters that match regex. The parameter must exist.

        :param name: regular expression
        :type name: str
        :returns: dict with key param name and value torch Tensor
        """
        pattern = re.compile(name)
        params_dict = {key: self.get_param(key) for key in self._params.keys()
                       if pattern.match(key)}
        return params_dict

    def param_name(self, p):
        """
        Get parameter name from parameter

        :param p: parameter
        :returns: parameter name
        """
        if p not in self._param_to_name:
            return None

        return self._param_to_name[p]

    def get_state(self):
        """
        Get the ParamStore state.
        """
        state = {
            'params': self._params,
            'constraints': self._constraints,
        }
        return state

    def set_state(self, state):
        """
        Set the ParamStore state using state from a previous get_state() call
        """
        assert isinstance(state, dict), "malformed ParamStore state"
        assert set(state.keys()) == set(['params', 'constraints']), \
            "malformed ParamStore keys {}".format(state.keys())

        for param_name, param in state['params'].items():
            self._params[param_name] = param
            self._param_to_name[param] = param_name

        for param_name, constraint in state['constraints'].items():
            if isinstance(constraint, type(constraints.real)):
                # Work around lack of hash & equality comparison on constraints.
                constraint = constraints.real
            self._constraints[param_name] = constraint

    def save(self, filename):
        """
        Save parameters to disk

        :param filename: file name to save to
        :type name: str
        """
        with open(filename, "wb") as output_file:
            torch.save(self.get_state(), output_file)

    def load(self, filename):
        """
        Loads parameters from disk

        .. note::

           If using :meth:`pyro.module` on parameters loaded from
           disk, be sure to set the ``update_module_params`` flag::

               pyro.get_param_store().load('saved_params.save')
               pyro.module('module', nn, update_module_params=True)

        :param filename: file name to load from
        :type name: str
        """
        with open(filename, "rb") as input_file:
            state = torch.load(input_file)
        self.set_state(state)

# used to create fully-formed param names, e.g. mymodule$$$mysubmodule.weight
_MODULE_NAMESPACE_DIVIDER = "$$$"

def param_with_module_name(pyro_name, param_name):
    return _MODULE_NAMESPACE_DIVIDER.join([pyro_name, param_name])

def module_from_param_with_module_name(param_name):
    return param_name.split(_MODULE_NAMESPACE_DIVIDER)[0]

def user_param_name(param_name):
    if _MODULE_NAMESPACE_DIVIDER in param_name:
        return param_name.split(_MODULE_NAMESPACE_DIVIDER)[1]
    return param_name
<|endoftext|>


<commit_before>#!/usr/bin/python3
import gi, time
import socket

# import GStreamer and GTK-Helper classes
gi.require_version('Gst', '1.0')
from gi.repository import GLib, Gst, GObject

# init GObject before importing local classes
GObject.threads_init()
Gst.init(None)

class Example:
	def __init__(self):
		self.mainloop = GObject.MainLoop()
		self.vsink   = Gst.parse_launch('intervideosrc channel=video ! video/x-raw,height=600,width=800,format=I420,framerate=25/1 ! timeoverlay ! videoconvert ! ximagesink')
		self.vsource = None

		self.asink   = Gst.parse_launch('interaudiosrc channel=audio ! audio/x-raw,format=S16LE,layout=interleaved,rate=48000,channels=2 ! autoaudiosink')
		self.asource = None

		# Create the server, binding to localhost on port 5000
		vsock = socket.socket(socket.AF_INET6)
		vsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
		vsock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, False)
		vsock.bind(('::', 5000))
		vsock.listen(1)

		# register socket for callback inside the GTK-Mainloop
		GObject.io_add_watch(vsock, GObject.IO_IN, self.connection_handler_video)

		# Create the server, binding to localhost on port 6000
		asock = socket.socket(socket.AF_INET6)
		asock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
		asock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, False)
		asock.bind(('::', 6000))
		asock.listen(1)

		# register socket for callback inside the GTK-Mainloop
		GObject.io_add_watch(asock, GObject.IO_IN, self.connection_handler_audio)

	def connection_handler_video(self, sock, *args):
		'''Asynchronous connection listener. Starts a handler for each connection.'''
		if self.vsource:
			return False

		conn, addr = sock.accept()
		print("Connection from", addr)

		self.vsource = Gst.parse_launch('appsrc name=a ! gdpdepay ! video/x-raw,height=600,width=800,format=I420,framerate=25/1 ! timeoverlay halignment=right ! intervideosink channel=video')
		self.vsource.set_state(Gst.State.PLAYING)

		# register data-received handler inside the GTK-Mainloop
		GObject.io_add_watch(conn, GObject.IO_IN, self.data_handler_video)
		return True

	def data_handler_video(self, conn, *args):
		'''Asynchronous data handler. Processes data-blocks line from the socket.'''
		blob = conn.recv(10000000) # >1920x1080x3
		if not len(blob):
			print("Connection closed.")
			self.vsource.set_state(Gst.State.NULL)
			self.vsource = None
			return False

		print("Video-Blob of %u bytes" % len(blob))
		buf = Gst.Buffer.new_wrapped(blob)
		self.vsource.get_by_name('a').emit('push-buffer', buf)
		return True

	def connection_handler_audio(self, sock, *args):
		'''Asynchronous connection listener. Starts a handler for each connection.'''
		if self.asource:
			return False

		conn, addr = sock.accept()
		print("Connection from", addr)

		self.asource = Gst.parse_launch('appsrc name=a ! gdpdepay ! audio/x-raw,format=S16LE,layout=interleaved,rate=48000,channels=2 ! interaudiosink channel=audio')
		self.asource.set_state(Gst.State.PLAYING)

		# register data-received handler inside the GTK-Mainloop
		GObject.io_add_watch(conn, GObject.IO_IN, self.data_handler_audio)
		return True

	def data_handler_audio(self, conn, *args):
		'''Asynchronous data handler. Processes data-blocks line from the socket.'''
		blob = conn.recv(10000000) # >1920x1080x3
		if not len(blob):
			print("Connection closed.")
			self.asource.set_state(Gst.State.NULL)
			self.asource = None
			return False

		print("Audio-Blob of %u bytes" % len(blob))
		buf = Gst.Buffer.new_wrapped(blob)
		self.asource.get_by_name('a').emit('push-buffer', buf)
		return True

	def run(self):
		self.vsink.set_state(Gst.State.PLAYING)
		self.asink.set_state(Gst.State.PLAYING)
		self.mainloop.run()

	def kill(self):
		self.vsink.set_state(Gst.State.NULL)
		self.asink.set_state(Gst.State.NULL)
		self.mainloop.quit()

	def on_eos(self, bus, msg):
		print('on_eos()')
		#self.kill()

	def on_error(self, bus, msg):
		print('on_error():', msg.parse_error())
		#self.kill()

example = Example()
example.run()
<commit_msg>Experiment with Muxing and use socket-fds to transport data into gstreamer<commit_after>#!/usr/bin/python3
import gi, time
import socket

# import GStreamer and GTK-Helper classes
gi.require_version('Gst', '1.0')
from gi.repository import GLib, Gst, GObject

# init GObject before importing local classes
GObject.threads_init()
Gst.init(None)

class Example:
	def __init__(self):
		self.mainloop = GObject.MainLoop()
		#self.vsink   = Gst.parse_launch('intervideosrc channel=video ! video/x-raw,height=600,width=800,format=I420,framerate=25/1 ! timeoverlay ! videoconvert ! ximagesink')
		self.vsource = None

		#self.asink   = Gst.parse_launch('interaudiosrc channel=audio ! audio/x-raw,format=S16LE,layout=interleaved,rate=48000,channels=2 ! autoaudiosink')
		self.asource = None

		self.sink   = Gst.parse_launch("""
			intervideosrc channel=video !
				queue !
				video/x-raw,height=600,width=800,format=I420,framerate=25/1 !
				timeoverlay !
				tee name=vtee !
				queue !
				videoconvert !
				avenc_mpeg2video bitrate=50000 max-key-interval=0 !
				queue !
				mux.

			interaudiosrc blocksize=4096 channel=audio !
				queue !
				audio/x-raw,format=S16LE,layout=interleaved,rate=48000,channels=2 !
				tee name=atee !
				queue !
				avenc_mp2 bitrate=192000 !
				queue !
				mux.

			mpegtsmux name=mux !
				filesink location=foo.ts

			vtee. ! queue ! videoconvert ! xvimagesink
			atee. ! queue ! audioconvert ! alsasink
		""")

		# Create the server, binding to localhost on port 5000
		vsock = socket.socket(socket.AF_INET6)
		vsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
		vsock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, False)
		vsock.bind(('::', 5000))
		vsock.listen(1)

		# Create the server, binding to localhost on port 5000
		asock = socket.socket(socket.AF_INET6)
		asock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
		asock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, False)
		asock.bind(('::', 6000))
		asock.listen(1)

		# register socket for callback inside the GTK-Mainloop
		GObject.io_add_watch(vsock, GObject.IO_IN, self.video_connect)
		GObject.io_add_watch(asock, GObject.IO_IN, self.audio_connect)

	def video_connect(self, sock, *args):
		'''Asynchronous connection listener. Starts a handler for each connection.'''
		if self.vsource:
			return False

		conn, addr = sock.accept()
		print("Connection from", addr)

		self.vsource = Gst.parse_launch('fdsrc name=a fd=%u ! gdpdepay ! video/x-raw,height=600,width=800,format=I420,framerate=25/1 ! timeoverlay halignment=right ! intervideosink channel=video' % conn.fileno())

		self.vsource.bus.add_signal_watch()
		self.vsource.bus.connect("message::eos", self.video_disconnect)

		self.vsource.set_state(Gst.State.PLAYING)

		self.vconn = conn
		return True

	def video_disconnect(self, bus, message):
		self.vsource.set_state(Gst.State.NULL)
		self.vsource = None
		self.vconn = None
		return True

	def audio_connect(self, sock, *args):
		'''Asynchronous connection listener. Starts a handler for each connection.'''
		if self.asource:
			return False

		conn, addr = sock.accept()
		print("Connection from", addr)

		self.asource = Gst.parse_launch('fdsrc name=a fd=%u ! gdpdepay ! audio/x-raw,format=S16LE,layout=interleaved,rate=48000,channels=2 ! interaudiosink channel=audio' % conn.fileno())

		self.asource.bus.add_signal_watch()
		self.asource.bus.connect("message::eos", self.audio_disconnect)

		self.asource.set_state(Gst.State.PLAYING)

		self.aconn = conn
		return True

	def audio_disconnect(self, bus, message):
		self.asource.set_state(Gst.State.NULL)
		self.asource = None
		self.aconn = None
		return True

	def run(self):
		self.sink.set_state(Gst.State.PLAYING)
		self.mainloop.run()

	def kill(self):
		self.sink.set_state(Gst.State.NULL)
		self.mainloop.quit()

example = Example()
example.run()
<|endoftext|>