<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
FreeSurfer tools interfaces
~~~~~~~~~~~~~~~~~~~~~~~~~~~

Fetch some example data:

    >>> import os
    >>> from niworkflows import data
    >>> data_root = data.get_bids_examples(variant='BIDS-examples-1-enh-ds054')
    >>> os.chdir(data_root)

Disable warnings:

    >>> import niworkflows.nipype as nn
    >>> nn.logging.getLogger('interface').setLevel('ERROR')

"""

import os.path as op
import nibabel as nb
import numpy as np

from skimage import morphology as sim
from scipy.ndimage.morphology import binary_fill_holes
from nilearn.image import resample_to_img, new_img_like

from niworkflows.nipype.utils.filemanip import copyfile, filename_to_list, fname_presuffix
from niworkflows.nipype.interfaces.base import (
    isdefined, InputMultiPath, BaseInterfaceInputSpec, TraitedSpec, File, traits, Directory
)
from niworkflows.nipype.interfaces import freesurfer as fs
from niworkflows.nipype.interfaces.base import SimpleInterface
from niworkflows.nipype.interfaces.freesurfer.preprocess import ConcatenateLTA

class StructuralReference(fs.RobustTemplate):
    """ Variation on RobustTemplate that simply copies the source if a single
    volume is provided.

    >>> from fmriprep.utils.bids import collect_data
    >>> t1w = collect_data('ds114', '01')[0]['t1w']
    >>> template = StructuralReference()
    >>> template.inputs.in_files = t1w
    >>> template.inputs.auto_detect_sensitivity = True
    >>> template.cmdline  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
    'mri_robust_template --satit --mov .../sub-01_ses-retest_T1w.nii.gz
        .../sub-01_ses-test_T1w.nii.gz --template mri_robust_template_out.mgz'

    """

    def _num_vols(self):
        n_files = len(self.inputs.in_files)
        if n_files != 1:
            return n_files

        img = nb.load(self.inputs.in_files[0])
        if len(img.shape) == 3:
            return 1

        return img.shape[3]

    @property
    def cmdline(self):
        if self._num_vols() == 1:
            return "echo Only one time point!"
        return super(StructuralReference, self).cmdline

    def _list_outputs(self):
        outputs = super(StructuralReference, self)._list_outputs()
        if self._num_vols() == 1:
            in_file = self.inputs.in_files[0]
            outputs['out_file'] = in_file
            if isdefined(outputs['transform_outputs']):
                transform_file = outputs['transform_outputs'][0]
                fs.utils.LTAConvert(in_lta='identity.nofile', source_file=in_file,
                                    target_file=in_file, out_lta=transform_file).run()
        return outputs

class MakeMidthicknessInputSpec(fs.utils.MRIsExpandInputSpec):
    graymid = InputMultiPath(desc='Existing graymid/midthickness file')

class MakeMidthickness(fs.MRIsExpand):
    """ Variation on MRIsExpand that checks for an existing midthickness/graymid
    surface, and copies if available.

    mris_expand is an expensive operation, so this avoids re-running it when the
    working directory is lost.
    If users provide their own midthickness/graymid file, we assume they have
    created it correctly.
    """
    input_spec = MakeMidthicknessInputSpec

    @property
    def cmdline(self):
        cmd = super(MakeMidthickness, self).cmdline
        if not isdefined(self.inputs.graymid) or len(self.inputs.graymid) < 1:
            return cmd

        # Possible graymid values inclue {l,r}h.{graymid,midthickness}
        # Prefer midthickness to graymid, require to be of the same hemisphere
        # as input
        source = None
        in_base = op.basename(self.inputs.in_file)
        mt = self._associated_file(in_base, 'midthickness')
        gm = self._associated_file(in_base, 'graymid')

        for surf in self.inputs.graymid:
            if op.basename(surf) == mt:
                source = surf
                break
            if op.basename(surf) == gm:
                source = surf

        if source is None:
            return cmd

        return "cp {} {}".format(source, self._list_outputs()['out_file'])

class FSInjectBrainExtractedInputSpec(BaseInterfaceInputSpec):
    subjects_dir = Directory(mandatory=True, desc='FreeSurfer SUBJECTS_DIR')
    subject_id = traits.Str(mandatory=True, desc='Subject ID')
    in_brain = File(mandatory=True, exists=True, desc='input file, part of a BIDS tree')

class FSInjectBrainExtractedOutputSpec(TraitedSpec):
    subjects_dir = Directory(desc='FreeSurfer SUBJECTS_DIR')
    subject_id = traits.Str(desc='Subject ID')

class FSInjectBrainExtracted(SimpleInterface):
    input_spec = FSInjectBrainExtractedInputSpec
    output_spec = FSInjectBrainExtractedOutputSpec

    def _run_interface(self, runtime):
        subjects_dir, subject_id = inject_skullstripped(
            self.inputs.subjects_dir,
            self.inputs.subject_id,
            self.inputs.in_brain)
        self._results['subjects_dir'] = subjects_dir
        self._results['subject_id'] = subject_id
        return runtime

class FSDetectInputsInputSpec(BaseInterfaceInputSpec):
    t1w_list = InputMultiPath(File(exists=True), mandatory=True,
                              desc='input file, part of a BIDS tree')
    t2w_list = InputMultiPath(File(exists=True), desc='input file, part of a BIDS tree')
    hires_enabled = traits.Bool(True, usedefault=True, desc='enable hi-resolution processing')

class FSDetectInputsOutputSpec(TraitedSpec):
    t2w = File(desc='reference T2w image')
    use_t2w = traits.Bool(desc='enable use of T2w downstream computation')
    hires = traits.Bool(desc='enable hi-res processing')
    mris_inflate = traits.Str(desc='mris_inflate argument')

class FSDetectInputs(SimpleInterface):
    input_spec = FSDetectInputsInputSpec
    output_spec = FSDetectInputsOutputSpec

    def _run_interface(self, runtime):
        t2w, self._results['hires'], mris_inflate = detect_inputs(
            self.inputs.t1w_list,
            hires_enabled=self.inputs.hires_enabled,
            t2w_list=self.inputs.t2w_list if isdefined(self.inputs.t2w_list) else None)

        self._results['use_t2w'] = t2w is not None
        if self._results['use_t2w']:
            self._results['t2w'] = t2w

        if self._results['hires']:
            self._results['mris_inflate'] = mris_inflate

        return runtime

class PatchedConcatenateLTA(ConcatenateLTA):
    """
    A temporarily patched version of ``fs.ConcatenateLTA`` to recover from
    `this bug <https://www.mail-archive.com/freesurfer@nmr.mgh.harvard.edu/msg55520.html>`_
    in FreeSurfer, that was
    `fixed here <https://github.com/freesurfer/freesurfer/pull/180>`_.

    The original FMRIPREP's issue is found
    `here <https://github.com/poldracklab/fmriprep/issues/768>`_.
    """

    def _list_outputs(self):
        outputs = super(ConcatenateLTA, self)._list_outputs()

        with open(outputs['out_file'], 'r') as f:
            lines = f.readlines()

        fixed = False
        newfile = []
        for line in lines:
            if line.startswith('filename = ') and len(line.strip("\n")) >= 255:
                fixed = True
                newfile.append('filename = path_too_long\n')
            else:
                newfile.append(line)

        if fixed:
            with open(outputs['out_file'], 'w') as f:
                f.write(''.join(newfile))
        return outputs

class RefineBrainMaskInputSpec(BaseInterfaceInputSpec):
    in_anat = File(exists=True, mandatory=True,
                   desc='input anatomical reference (INU corrected)')
    in_aseg = File(exists=True, mandatory=True,
                   desc='input ``aseg`` file, in NifTi format.')
    in_ants = File(exists=True, mandatory=True,
                   desc='brain tissue segmentation generated with antsBrainExtraction.sh')

class RefineBrainMaskOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc='new mask')

class RefineBrainMask(SimpleInterface):
    input_spec = RefineBrainMaskInputSpec
    output_spec = RefineBrainMaskOutputSpec

    def _run_interface(self, runtime):

        self._results['out_file'] = fname_presuffix(
            self.inputs.in_anat, suffix='_rbrainmask', newpath=runtime.cwd)

        anatnii = nb.load(self.inputs.in_anat)
        msknii = nb.Nifti1Image(
            grow_mask(anatnii.get_data(),
                      nb.load(self.inputs.in_aseg).get_data(),
                      nb.load(self.inputs.in_ants).get_data()),
            anatnii.affine,
            anatnii.header
        )
        msknii.set_data_dtype(np.uint8)
        msknii.to_filename(self._results['out_file'])

        return runtime

def inject_skullstripped(subjects_dir, subject_id, skullstripped):
    mridir = op.join(subjects_dir, subject_id, 'mri')
    t1 = op.join(mridir, 'T1.mgz')
    bm_auto = op.join(mridir, 'brainmask.auto.mgz')
    bm = op.join(mridir, 'brainmask.mgz')

    if not op.exists(bm_auto):
        img = nb.load(t1)
        mask = nb.load(skullstripped)
        bmask = new_img_like(mask, mask.get_data() > 0)
        resampled_mask = resample_to_img(bmask, img, 'nearest')
        masked_image = new_img_like(img, img.get_data() * resampled_mask.get_data())
        masked_image.to_filename(bm_auto)

    if not op.exists(bm):
        copyfile(bm_auto, bm, copy=True, use_hardlink=True)

    return subjects_dir, subject_id

def detect_inputs(t1w_list, t2w_list=None, hires_enabled=True):
    t1w_list = filename_to_list(t1w_list)
    t2w_list = filename_to_list(t2w_list) if t2w_list is not None else []
    t1w_ref = nb.load(t1w_list[0])
    # Use high resolution preprocessing if voxel size < 1.0mm
    # Tolerance of 0.05mm requires that rounds down to 0.9mm or lower
    hires = hires_enabled and max(t1w_ref.header.get_zooms()) < 1 - 0.05

    t2w = None
    if t2w_list and max(nb.load(t2w_list[0]).header.get_zooms()) < 1.2:
        t2w = t2w_list[0]

    # https://surfer.nmr.mgh.harvard.edu/fswiki/SubmillimeterRecon
    mris_inflate = '-n 50' if hires else None
    return (t2w, hires, mris_inflate)

def refine_aseg(aseg, ball_size=4):
    """
    First step to reconcile ANTs' and FreeSurfer's brain masks.

    Here, the ``aseg.mgz`` mask from FreeSurfer is refined in two
    steps, using binary morphological operations:

      1. With a binary closing operation the sulci are included
         into the mask. This results in a smoother brain mask
         that does not exclude deep, wide sulci.

      2. Fill any holes (typically, there could be a hole next to
         the pineal gland and the corpora quadrigemina if the great
         cerebral brain is segmented out).

    """
    # Read aseg data
    bmask = aseg.copy()
    bmask[bmask > 0] = 1
    bmask = bmask.astype(np.uint8)

    # Morphological operations
    selem = sim.ball(ball_size)
    newmask = sim.binary_closing(bmask, selem)
    newmask = binary_fill_holes(newmask.astype(np.uint8), selem).astype(np.uint8)

    return newmask.astype(np.uint8)

def grow_mask(anat, aseg, ants_segs, ww=7, zval=2.0, bw=4):
    """
    Grow mask including pixels that have a high likelihood.
    GM tissue parameters are sampled in image patches of ``ww`` size.

    This is inspired on mindboggle's solution to the problem:
    https://github.com/nipy/mindboggle/blob/master/mindboggle/guts/segment.py#L1660

    """
    selem = sim.ball(bw)

    aseg[aseg == 42] = 3  # Collapse both hemispheres
    gm = anat.copy()
    gm[aseg != 3] = 0

    refined = refine_aseg(aseg)
    newrefmask = sim.binary_dilation(refined, selem) - refined
    indices = np.argwhere(newrefmask > 0)
    for pixel in indices:
        # When ATROPOS identified the pixel as GM, set and carry on
        if ants_segs[tuple(pixel)] == 2:
            refined[tuple(pixel)] = 1
            continue

        window = gm[
            pixel[0] - ww:pixel[0] + ww,
            pixel[1] - ww:pixel[1] + ww,
            pixel[2] - ww:pixel[2] + ww
        ]
        if np.any(window > 0):
            zstat = abs(anat[tuple(pixel)] - window[window > 0].mean()) / window[window > 0].std()
            refined[tuple(pixel)] = int(zstat < zval)

    refined = sim.binary_opening(refined, selem)
    return refined
<commit_msg>[docs only] Make ATROPOS segmentation optional<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
FreeSurfer tools interfaces
~~~~~~~~~~~~~~~~~~~~~~~~~~~

Fetch some example data:

    >>> import os
    >>> from niworkflows import data
    >>> data_root = data.get_bids_examples(variant='BIDS-examples-1-enh-ds054')
    >>> os.chdir(data_root)

Disable warnings:

    >>> import niworkflows.nipype as nn
    >>> nn.logging.getLogger('interface').setLevel('ERROR')

"""

import os.path as op
import nibabel as nb
import numpy as np

from skimage import morphology as sim
from scipy.ndimage.morphology import binary_fill_holes
from nilearn.image import resample_to_img, new_img_like

from niworkflows.nipype.utils.filemanip import copyfile, filename_to_list, fname_presuffix
from niworkflows.nipype.interfaces.base import (
    isdefined, InputMultiPath, BaseInterfaceInputSpec, TraitedSpec, File, traits, Directory
)
from niworkflows.nipype.interfaces import freesurfer as fs
from niworkflows.nipype.interfaces.base import SimpleInterface
from niworkflows.nipype.interfaces.freesurfer.preprocess import ConcatenateLTA

class StructuralReference(fs.RobustTemplate):
    """ Variation on RobustTemplate that simply copies the source if a single
    volume is provided.

    >>> from fmriprep.utils.bids import collect_data
    >>> t1w = collect_data('ds114', '01')[0]['t1w']
    >>> template = StructuralReference()
    >>> template.inputs.in_files = t1w
    >>> template.inputs.auto_detect_sensitivity = True
    >>> template.cmdline  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
    'mri_robust_template --satit --mov .../sub-01_ses-retest_T1w.nii.gz
        .../sub-01_ses-test_T1w.nii.gz --template mri_robust_template_out.mgz'

    """

    def _num_vols(self):
        n_files = len(self.inputs.in_files)
        if n_files != 1:
            return n_files

        img = nb.load(self.inputs.in_files[0])
        if len(img.shape) == 3:
            return 1

        return img.shape[3]

    @property
    def cmdline(self):
        if self._num_vols() == 1:
            return "echo Only one time point!"
        return super(StructuralReference, self).cmdline

    def _list_outputs(self):
        outputs = super(StructuralReference, self)._list_outputs()
        if self._num_vols() == 1:
            in_file = self.inputs.in_files[0]
            outputs['out_file'] = in_file
            if isdefined(outputs['transform_outputs']):
                transform_file = outputs['transform_outputs'][0]
                fs.utils.LTAConvert(in_lta='identity.nofile', source_file=in_file,
                                    target_file=in_file, out_lta=transform_file).run()
        return outputs

class MakeMidthicknessInputSpec(fs.utils.MRIsExpandInputSpec):
    graymid = InputMultiPath(desc='Existing graymid/midthickness file')

class MakeMidthickness(fs.MRIsExpand):
    """ Variation on MRIsExpand that checks for an existing midthickness/graymid
    surface, and copies if available.

    mris_expand is an expensive operation, so this avoids re-running it when the
    working directory is lost.
    If users provide their own midthickness/graymid file, we assume they have
    created it correctly.
    """
    input_spec = MakeMidthicknessInputSpec

    @property
    def cmdline(self):
        cmd = super(MakeMidthickness, self).cmdline
        if not isdefined(self.inputs.graymid) or len(self.inputs.graymid) < 1:
            return cmd

        # Possible graymid values inclue {l,r}h.{graymid,midthickness}
        # Prefer midthickness to graymid, require to be of the same hemisphere
        # as input
        source = None
        in_base = op.basename(self.inputs.in_file)
        mt = self._associated_file(in_base, 'midthickness')
        gm = self._associated_file(in_base, 'graymid')

        for surf in self.inputs.graymid:
            if op.basename(surf) == mt:
                source = surf
                break
            if op.basename(surf) == gm:
                source = surf

        if source is None:
            return cmd

        return "cp {} {}".format(source, self._list_outputs()['out_file'])

class FSInjectBrainExtractedInputSpec(BaseInterfaceInputSpec):
    subjects_dir = Directory(mandatory=True, desc='FreeSurfer SUBJECTS_DIR')
    subject_id = traits.Str(mandatory=True, desc='Subject ID')
    in_brain = File(mandatory=True, exists=True, desc='input file, part of a BIDS tree')

class FSInjectBrainExtractedOutputSpec(TraitedSpec):
    subjects_dir = Directory(desc='FreeSurfer SUBJECTS_DIR')
    subject_id = traits.Str(desc='Subject ID')

class FSInjectBrainExtracted(SimpleInterface):
    input_spec = FSInjectBrainExtractedInputSpec
    output_spec = FSInjectBrainExtractedOutputSpec

    def _run_interface(self, runtime):
        subjects_dir, subject_id = inject_skullstripped(
            self.inputs.subjects_dir,
            self.inputs.subject_id,
            self.inputs.in_brain)
        self._results['subjects_dir'] = subjects_dir
        self._results['subject_id'] = subject_id
        return runtime

class FSDetectInputsInputSpec(BaseInterfaceInputSpec):
    t1w_list = InputMultiPath(File(exists=True), mandatory=True,
                              desc='input file, part of a BIDS tree')
    t2w_list = InputMultiPath(File(exists=True), desc='input file, part of a BIDS tree')
    hires_enabled = traits.Bool(True, usedefault=True, desc='enable hi-resolution processing')

class FSDetectInputsOutputSpec(TraitedSpec):
    t2w = File(desc='reference T2w image')
    use_t2w = traits.Bool(desc='enable use of T2w downstream computation')
    hires = traits.Bool(desc='enable hi-res processing')
    mris_inflate = traits.Str(desc='mris_inflate argument')

class FSDetectInputs(SimpleInterface):
    input_spec = FSDetectInputsInputSpec
    output_spec = FSDetectInputsOutputSpec

    def _run_interface(self, runtime):
        t2w, self._results['hires'], mris_inflate = detect_inputs(
            self.inputs.t1w_list,
            hires_enabled=self.inputs.hires_enabled,
            t2w_list=self.inputs.t2w_list if isdefined(self.inputs.t2w_list) else None)

        self._results['use_t2w'] = t2w is not None
        if self._results['use_t2w']:
            self._results['t2w'] = t2w

        if self._results['hires']:
            self._results['mris_inflate'] = mris_inflate

        return runtime

class PatchedConcatenateLTA(ConcatenateLTA):
    """
    A temporarily patched version of ``fs.ConcatenateLTA`` to recover from
    `this bug <https://www.mail-archive.com/freesurfer@nmr.mgh.harvard.edu/msg55520.html>`_
    in FreeSurfer, that was
    `fixed here <https://github.com/freesurfer/freesurfer/pull/180>`_.

    The original FMRIPREP's issue is found
    `here <https://github.com/poldracklab/fmriprep/issues/768>`_.
    """

    def _list_outputs(self):
        outputs = super(ConcatenateLTA, self)._list_outputs()

        with open(outputs['out_file'], 'r') as f:
            lines = f.readlines()

        fixed = False
        newfile = []
        for line in lines:
            if line.startswith('filename = ') and len(line.strip("\n")) >= 255:
                fixed = True
                newfile.append('filename = path_too_long\n')
            else:
                newfile.append(line)

        if fixed:
            with open(outputs['out_file'], 'w') as f:
                f.write(''.join(newfile))
        return outputs

class RefineBrainMaskInputSpec(BaseInterfaceInputSpec):
    in_anat = File(exists=True, mandatory=True,
                   desc='input anatomical reference (INU corrected)')
    in_aseg = File(exists=True, mandatory=True,
                   desc='input ``aseg`` file, in NifTi format.')
    in_ants = File(exists=True, mandatory=True,
                   desc='brain tissue segmentation generated with antsBrainExtraction.sh')

class RefineBrainMaskOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc='new mask')

class RefineBrainMask(SimpleInterface):
    input_spec = RefineBrainMaskInputSpec
    output_spec = RefineBrainMaskOutputSpec

    def _run_interface(self, runtime):

        self._results['out_file'] = fname_presuffix(
            self.inputs.in_anat, suffix='_rbrainmask', newpath=runtime.cwd)

        anatnii = nb.load(self.inputs.in_anat)
        msknii = nb.Nifti1Image(
            grow_mask(anatnii.get_data(),
                      nb.load(self.inputs.in_aseg).get_data(),
                      nb.load(self.inputs.in_ants).get_data()),
            anatnii.affine,
            anatnii.header
        )
        msknii.set_data_dtype(np.uint8)
        msknii.to_filename(self._results['out_file'])

        return runtime

def inject_skullstripped(subjects_dir, subject_id, skullstripped):
    mridir = op.join(subjects_dir, subject_id, 'mri')
    t1 = op.join(mridir, 'T1.mgz')
    bm_auto = op.join(mridir, 'brainmask.auto.mgz')
    bm = op.join(mridir, 'brainmask.mgz')

    if not op.exists(bm_auto):
        img = nb.load(t1)
        mask = nb.load(skullstripped)
        bmask = new_img_like(mask, mask.get_data() > 0)
        resampled_mask = resample_to_img(bmask, img, 'nearest')
        masked_image = new_img_like(img, img.get_data() * resampled_mask.get_data())
        masked_image.to_filename(bm_auto)

    if not op.exists(bm):
        copyfile(bm_auto, bm, copy=True, use_hardlink=True)

    return subjects_dir, subject_id

def detect_inputs(t1w_list, t2w_list=None, hires_enabled=True):
    t1w_list = filename_to_list(t1w_list)
    t2w_list = filename_to_list(t2w_list) if t2w_list is not None else []
    t1w_ref = nb.load(t1w_list[0])
    # Use high resolution preprocessing if voxel size < 1.0mm
    # Tolerance of 0.05mm requires that rounds down to 0.9mm or lower
    hires = hires_enabled and max(t1w_ref.header.get_zooms()) < 1 - 0.05

    t2w = None
    if t2w_list and max(nb.load(t2w_list[0]).header.get_zooms()) < 1.2:
        t2w = t2w_list[0]

    # https://surfer.nmr.mgh.harvard.edu/fswiki/SubmillimeterRecon
    mris_inflate = '-n 50' if hires else None
    return (t2w, hires, mris_inflate)

def refine_aseg(aseg, ball_size=4):
    """
    First step to reconcile ANTs' and FreeSurfer's brain masks.

    Here, the ``aseg.mgz`` mask from FreeSurfer is refined in two
    steps, using binary morphological operations:

      1. With a binary closing operation the sulci are included
         into the mask. This results in a smoother brain mask
         that does not exclude deep, wide sulci.

      2. Fill any holes (typically, there could be a hole next to
         the pineal gland and the corpora quadrigemina if the great
         cerebral brain is segmented out).

    """
    # Read aseg data
    bmask = aseg.copy()
    bmask[bmask > 0] = 1
    bmask = bmask.astype(np.uint8)

    # Morphological operations
    selem = sim.ball(ball_size)
    newmask = sim.binary_closing(bmask, selem)
    newmask = binary_fill_holes(newmask.astype(np.uint8), selem).astype(np.uint8)

    return newmask.astype(np.uint8)

def grow_mask(anat, aseg, ants_segs=None, ww=7, zval=2.0, bw=4):
    """
    Grow mask including pixels that have a high likelihood.
    GM tissue parameters are sampled in image patches of ``ww`` size.

    This is inspired on mindboggle's solution to the problem:
    https://github.com/nipy/mindboggle/blob/master/mindboggle/guts/segment.py#L1660

    """
    selem = sim.ball(bw)

    if ants_segs is None:
        ants_segs = np.zeros_like(aseg, dtype=np.uint8)

    aseg[aseg == 42] = 3  # Collapse both hemispheres
    gm = anat.copy()
    gm[aseg != 3] = 0

    refined = refine_aseg(aseg)
    newrefmask = sim.binary_dilation(refined, selem) - refined
    indices = np.argwhere(newrefmask > 0)
    for pixel in indices:
        # When ATROPOS identified the pixel as GM, set and carry on
        if ants_segs[tuple(pixel)] == 2:
            refined[tuple(pixel)] = 1
            continue

        window = gm[
            pixel[0] - ww:pixel[0] + ww,
            pixel[1] - ww:pixel[1] + ww,
            pixel[2] - ww:pixel[2] + ww
        ]
        if np.any(window > 0):
            zstat = abs(anat[tuple(pixel)] - window[window > 0].mean()) / window[window > 0].std()
            refined[tuple(pixel)] = int(zstat < zval)

    refined = sim.binary_opening(refined, selem)
    return refined
<|endoftext|>


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

# This script can check that an expected json blob is a subset of what actually gets produced.
# The comparison is independent of the value of IDs (which are unstable) and instead uses their
# relative ordering to check them against eachother by looking them up in their respective blob's
# `index` or `paths` mappings. To add a new test run `rustdoc --output-format json -o . yourtest.rs`
# and then create `yourtest.expected` by stripping unnecessary details from `yourtest.json`.

import copy
import sys
import json
import types

# Used instead of the string ids when used as references.
# Not used as keys in `index` or `paths`
class ID(str):
    pass

class SubsetException(Exception):
    def __init__(self, msg, trace):
        self.msg = msg
        self.trace = msg
        super().__init__("{}: {}".format(trace, msg))

def check_subset(expected_main, actual_main, base_dir):
    expected_index = expected_main["index"]
    expected_paths = expected_main["paths"]
    actual_index = actual_main["index"]
    actual_paths = actual_main["paths"]
    already_checked = set()

    def _check_subset(expected, actual, trace):
        expected_type = type(expected)
        actual_type = type(actual)

        if actual_type is str:
            actual = actual.replace(base_dir, "$TEST_BASE_DIR")

        if expected_type is not actual_type:
            raise SubsetException(
                "expected type `{}`, got `{}`".format(expected_type, actual_type), trace
            )

        if expected_type in (int, bool, str) and expected != actual:
            raise SubsetException("expected `{}`, got: `{}`".format(expected, actual), trace)
        if expected_type is dict:
            for key in expected:
                if key not in actual:
                    raise SubsetException(
                        "Key `{}` not found in output".format(key), trace
                    )
                new_trace = copy.deepcopy(trace)
                new_trace.append(key)
                _check_subset(expected[key], actual[key], new_trace)
        elif expected_type is list:
            expected_elements = len(expected)
            actual_elements = len(actual)
            if expected_elements != actual_elements:
                raise SubsetException(
                    "Found {} items, expected {}".format(
                        expected_elements, actual_elements
                    ),
                    trace,
                )
            for expected, actual in zip(expected, actual):
                new_trace = copy.deepcopy(trace)
                new_trace.append(expected)
                _check_subset(expected, actual, new_trace)
        elif expected_type is ID and expected not in already_checked:
            already_checked.add(expected)
            _check_subset(
                expected_index.get(expected, {}), actual_index.get(actual, {}), trace
            )
            _check_subset(
                expected_paths.get(expected, {}), actual_paths.get(actual, {}), trace
            )

    _check_subset(expected_main["root"], actual_main["root"], [])

def rustdoc_object_hook(obj):
    # No need to convert paths, index and external_crates keys to ids, since
    # they are the target of resolution, and never a source itself.
    if "id" in obj and obj["id"]:
        obj["id"] = ID(obj["id"])
    if "root" in obj:
        obj["root"] = ID(obj["root"])
    if "items" in obj:
        obj["items"] = [ID(id) for id in obj["items"]]
    if "variants" in obj:
        obj["variants"] = [ID(id) for id in obj["variants"]]
    if "fields" in obj:
        obj["fields"] = [ID(id) for id in obj["fields"]]
    if "impls" in obj:
        obj["impls"] = [ID(id) for id in obj["impls"]]
    if "implementors" in obj:
        obj["implementors"] = [ID(id) for id in obj["implementors"]]
    if "links" in obj:
        obj["links"] = {s: ID(id) for s, id in obj["links"]}
    if "variant_kind" in obj and obj["variant_kind"] == "struct":
        obj["variant_inner"] = [ID(id) for id in obj["variant_inner"]]
    return obj

def main(expected_fpath, actual_fpath, base_dir):
    print(
        "checking that {} is a logical subset of {}".format(
            expected_fpath, actual_fpath
        )
    )
    with open(expected_fpath) as expected_file:
        expected_main = json.load(expected_file, object_hook=rustdoc_object_hook)
    with open(actual_fpath) as actual_file:
        actual_main = json.load(actual_file, object_hook=rustdoc_object_hook)
    check_subset(expected_main, actual_main, base_dir)
    print("all checks passed")

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: `compare.py expected.json actual.json test-dir`")
    else:
        main(sys.argv[1], sys.argv[2], sys.argv[3])
<commit_msg>Normalize windows path seperators.<commit_after>#!/usr/bin/env python

# This script can check that an expected json blob is a subset of what actually gets produced.
# The comparison is independent of the value of IDs (which are unstable) and instead uses their
# relative ordering to check them against eachother by looking them up in their respective blob's
# `index` or `paths` mappings. To add a new test run `rustdoc --output-format json -o . yourtest.rs`
# and then create `yourtest.expected` by stripping unnecessary details from `yourtest.json`. If
# you're on windows, replace `\` with `/`.

import copy
import sys
import json
import types

# Used instead of the string ids when used as references.
# Not used as keys in `index` or `paths`
class ID(str):
    pass

class SubsetException(Exception):
    def __init__(self, msg, trace):
        self.msg = msg
        self.trace = msg
        super().__init__("{}: {}".format(trace, msg))

def check_subset(expected_main, actual_main, base_dir):
    expected_index = expected_main["index"]
    expected_paths = expected_main["paths"]
    actual_index = actual_main["index"]
    actual_paths = actual_main["paths"]
    already_checked = set()

    def _check_subset(expected, actual, trace):
        expected_type = type(expected)
        actual_type = type(actual)

        if actual_type is str:
            actual = normalize(actual).replace(base_dir, "$TEST_BASE_DIR")

        if expected_type is not actual_type:
            raise SubsetException(
                "expected type `{}`, got `{}`".format(expected_type, actual_type), trace
            )

        if expected_type in (int, bool, str) and expected != actual:
            raise SubsetException("expected `{}`, got: `{}`".format(expected, actual), trace)
        if expected_type is dict:
            for key in expected:
                if key not in actual:
                    raise SubsetException(
                        "Key `{}` not found in output".format(key), trace
                    )
                new_trace = copy.deepcopy(trace)
                new_trace.append(key)
                _check_subset(expected[key], actual[key], new_trace)
        elif expected_type is list:
            expected_elements = len(expected)
            actual_elements = len(actual)
            if expected_elements != actual_elements:
                raise SubsetException(
                    "Found {} items, expected {}".format(
                        expected_elements, actual_elements
                    ),
                    trace,
                )
            for expected, actual in zip(expected, actual):
                new_trace = copy.deepcopy(trace)
                new_trace.append(expected)
                _check_subset(expected, actual, new_trace)
        elif expected_type is ID and expected not in already_checked:
            already_checked.add(expected)
            _check_subset(
                expected_index.get(expected, {}), actual_index.get(actual, {}), trace
            )
            _check_subset(
                expected_paths.get(expected, {}), actual_paths.get(actual, {}), trace
            )

    _check_subset(expected_main["root"], actual_main["root"], [])

def rustdoc_object_hook(obj):
    # No need to convert paths, index and external_crates keys to ids, since
    # they are the target of resolution, and never a source itself.
    if "id" in obj and obj["id"]:
        obj["id"] = ID(obj["id"])
    if "root" in obj:
        obj["root"] = ID(obj["root"])
    if "items" in obj:
        obj["items"] = [ID(id) for id in obj["items"]]
    if "variants" in obj:
        obj["variants"] = [ID(id) for id in obj["variants"]]
    if "fields" in obj:
        obj["fields"] = [ID(id) for id in obj["fields"]]
    if "impls" in obj:
        obj["impls"] = [ID(id) for id in obj["impls"]]
    if "implementors" in obj:
        obj["implementors"] = [ID(id) for id in obj["implementors"]]
    if "links" in obj:
        obj["links"] = {s: ID(id) for s, id in obj["links"]}
    if "variant_kind" in obj and obj["variant_kind"] == "struct":
        obj["variant_inner"] = [ID(id) for id in obj["variant_inner"]]
    return obj

def main(expected_fpath, actual_fpath, base_dir):
    print(
        "checking that {} is a logical subset of {}".format(
            expected_fpath, actual_fpath
        )
    )
    with open(expected_fpath) as expected_file:
        expected_main = json.load(expected_file, object_hook=rustdoc_object_hook)
    with open(actual_fpath) as actual_file:
        actual_main = json.load(actual_file, object_hook=rustdoc_object_hook)
    check_subset(expected_main, actual_main, base_dir)
    print("all checks passed")

def normalize(s):
    return s.replace('\\', '/')

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: `compare.py expected.json actual.json test-dir`")
    else:
        main(sys.argv[1], sys.argv[2], normalize(sys.argv[3]))
<|endoftext|>


<commit_before>from django.core.exceptions import ImproperlyConfigured
from django.contrib.auth.signals import user_logged_in

import larper
from larper import UserSession

def handle_login(sender, **kwargs):
    request = kwargs['request']
    larper.store_password(request, request.POST.get('password', ''))

user_logged_in.connect(handle_login)

class LarperMiddleware(object):
    """
    Responsible for populating the request.user object
    with the following attributes:
    * unique_id

    This complements the dn and password management from larper.dn and
    larper.password
    """
    def process_request(self, request):
        if not hasattr(request, 'user'):
            msg = "django.contrib.auth.middleware.AuthenticationMiddleware "
            "is missing from your settings.py"
            raise ImproperlyConfigured(msg)

        user = request.user

        if not hasattr(request, 'session'):
            msg = "django.contrib.sessions.middleware.SessionMiddleware "
            "is missing from your settings.py"
            raise ImproperlyConfigured(msg)

        session = request.session

        if '/en-US/media' not in request.path and\
           request.user.is_authenticated():
            _populate(user, session)

    def process_response(self, request, response):
        UserSession.disconnect(request)
        return response

def _populate(user, session):
    if 'unique_id' in session:
        user.unique_id = session['unique_id']
    else:
        unique_id = user.ldap_user.attrs['uniqueIdentifier'][0]
        user.unique_id = session['unique_id'] = unique_id
<commit_msg>Storing vouchedness on the user<commit_after>from django.core.exceptions import ImproperlyConfigured
from django.contrib.auth.signals import user_logged_in

import larper
from larper import UserSession

def handle_login(sender, **kwargs):
    request = kwargs['request']
    larper.store_password(request, request.POST.get('password', ''))

user_logged_in.connect(handle_login)

class LarperMiddleware(object):
    """
    Responsible for populating the request.user object
    with the following attributes:
    * unique_id

    This complements the dn and password management from larper.dn and
    larper.password
    """
    def process_request(self, request):
        if not hasattr(request, 'user'):
            msg = ('django.contrib.auth.middleware.AuthenticationMiddleware '
                   'is missing from your settings.py')
            raise ImproperlyConfigured(msg)

        if not hasattr(request, 'session'):
            msg = ('django.contrib.sessions.middleware.SessionMiddleware '
                   'is missing from your settings.py')
            raise ImproperlyConfigured(msg)

        if request.user.is_authenticated():
            _populate(request)

    def process_response(self, request, response):
        UserSession.disconnect(request)
        return response

def is_vouched(request):
    user = request.user

    def f():
        if not hasattr(user, 'person'):
            directory = UserSession.connect(request)
            user.person = directory.get_by_unique_id(user.unique_id)
        return bool(user.person.voucher_unique_id)
    return f

def _populate(request):
    user = request.user
    session = request.session

    if 'unique_id' in session:
        user.unique_id = session['unique_id']
    else:
        unique_id = user.ldap_user.attrs['uniqueIdentifier'][0]
        user.unique_id = session['unique_id'] = unique_id

    user.is_vouched = is_vouched(request)
<|endoftext|>


<commit_before># http://www.raspberrypi-spy.co.uk/2012/07/16x2-lcd-module-control-using-python/
# Hitachi HD44780 LCD controller
#!/usr/bin/python
#--------------------------------------
#    ___  ___  _ ____
#   / _ \/ _ \(_) __/__  __ __
#  / , _/ ___/ /\ \/ _ \/ // /
# /_/|_/_/  /_/___/ .__/\_, /
#                /_/   /___/
#
#  lcd_16x2.py
#  16x2 LCD Test Script
#
# Author : Matt Hawkins
# Date   : 06/04/2015
#
# http://www.raspberrypi-spy.co.uk/
#
#--------------------------------------

# The wiring for the LCD is as follows:
# 1 : GND
# 2 : 5V
# 3 : Contrast (0-5V)*
# 4 : RS (Register Select)
# 5 : R/W (Read Write)       - GROUND THIS PIN
# 6 : Enable or Strobe
# 7 : Data Bit 0             - NOT USED
# 8 : Data Bit 1             - NOT USED
# 9 : Data Bit 2             - NOT USED
# 10: Data Bit 3             - NOT USED
# 11: Data Bit 4
# 12: Data Bit 5
# 13: Data Bit 6
# 14: Data Bit 7
# 15: LCD Backlight +5V**
# 16: LCD Backlight GND

#import
import RPi.GPIO as GPIO
import time

# Define GPIO to LCD mapping
LCD_RS = 7
LCD_E  = 8
LCD_D4 = 25
LCD_D5 = 24
LCD_D6 = 23
LCD_D7 = 18

# Define some device constants
LCD_WIDTH = 16    # Maximum characters per line
LCD_CHR = True
LCD_CMD = False

LCD_LINE_1 = 0x80 # LCD RAM address for the 1st line
LCD_LINE_2 = 0xC0 # LCD RAM address for the 2nd line

# Timing constants
E_PULSE = 0.0005
E_DELAY = 0.0005

def main():
  # Main program block

  GPIO.setmode(GPIO.BCM)       # Use BCM GPIO numbers
  GPIO.setup(LCD_E, GPIO.OUT)  # E
  GPIO.setup(LCD_RS, GPIO.OUT) # RS
  GPIO.setup(LCD_D4, GPIO.OUT) # DB4
  GPIO.setup(LCD_D5, GPIO.OUT) # DB5
  GPIO.setup(LCD_D6, GPIO.OUT) # DB6
  GPIO.setup(LCD_D7, GPIO.OUT) # DB7

  # Initialise display
  lcd_init()

  while True:

    # Send some test
    lcd_string("Rasbperry Pi",LCD_LINE_1)
    lcd_string("16x2 LCD Test",LCD_LINE_2)

    time.sleep(3) # 3 second delay

    # Send some text
    lcd_string("1234567890123456",LCD_LINE_1)
    lcd_string("abcdefghijklmnop",LCD_LINE_2)

    time.sleep(3) # 3 second delay

    # Send some text
    lcd_string("RaspberryPi-spy",LCD_LINE_1)
    lcd_string(".co.uk",LCD_LINE_2)

    time.sleep(3)

    # Send some text
    lcd_string("Follow me on",LCD_LINE_1)
    lcd_string("Twitter @RPiSpy",LCD_LINE_2)

    time.sleep(3)

def lcd_init():
  # Initialise display
  lcd_byte(0x33,LCD_CMD) # 110011 Initialise
  lcd_byte(0x32,LCD_CMD) # 110010 Initialise
  lcd_byte(0x06,LCD_CMD) # 000110 Cursor move direction
  lcd_byte(0x0C,LCD_CMD) # 001100 Display On,Cursor Off, Blink Off
  lcd_byte(0x28,LCD_CMD) # 101000 Data length, number of lines, font size
  lcd_byte(0x01,LCD_CMD) # 000001 Clear display
  time.sleep(E_DELAY)

def lcd_byte(bits, mode):
  # Send byte to data pins
  # bits = data
  # mode = True  for character
  #        False for command

  GPIO.output(LCD_RS, mode) # RS

  # High bits
  GPIO.output(LCD_D4, False)
  GPIO.output(LCD_D5, False)
  GPIO.output(LCD_D6, False)
  GPIO.output(LCD_D7, False)
  if bits&0x10==0x10:
    GPIO.output(LCD_D4, True)
  if bits&0x20==0x20:
    GPIO.output(LCD_D5, True)
  if bits&0x40==0x40:
    GPIO.output(LCD_D6, True)
  if bits&0x80==0x80:
    GPIO.output(LCD_D7, True)

  # Toggle 'Enable' pin
  lcd_toggle_enable()

  # Low bits
  GPIO.output(LCD_D4, False)
  GPIO.output(LCD_D5, False)
  GPIO.output(LCD_D6, False)
  GPIO.output(LCD_D7, False)
  if bits&0x01==0x01:
    GPIO.output(LCD_D4, True)
  if bits&0x02==0x02:
    GPIO.output(LCD_D5, True)
  if bits&0x04==0x04:
    GPIO.output(LCD_D6, True)
  if bits&0x08==0x08:
    GPIO.output(LCD_D7, True)

  # Toggle 'Enable' pin
  lcd_toggle_enable()

def lcd_toggle_enable():
  # Toggle enable
  time.sleep(E_DELAY)
  GPIO.output(LCD_E, True)
  time.sleep(E_PULSE)
  GPIO.output(LCD_E, False)
  time.sleep(E_DELAY)

def lcd_string(message,line):
  # Send string to display

  message = message.ljust(LCD_WIDTH," ")

  lcd_byte(line, LCD_CMD)

  for i in range(LCD_WIDTH):
    lcd_byte(ord(message[i]),LCD_CHR)

if __name__ == '__main__':

  try:
    main()
  except KeyboardInterrupt:
    pass
  finally:
    lcd_byte(0x01, LCD_CMD)
    lcd_string("Goodbye!",LCD_LINE_1)
    GPIO.cleanup()
<commit_msg>clean files<commit_after><|endoftext|>


<commit_before>"""
RESTful API controller.

Endpoint for queries on taxonomic occurrences in time and space.
"""

# import connexion
# from swagger_server.models.error_model import ErrorModel
# from swagger_server.models.occurrence import Occurrence
# from datetime import date, datetime
# from typing import List, Dict
# from six import iteritems
# from ..util import deserialize_date, deserialize_datetime
from flask import request, jsonify
import requests
import time
from statistics import mean

def occ(bbox=None, minage=None, maxage=None, agescale=None, timerule=None,
        taxon=None, includelower=None, limit=None, offset=None, show=None):
    """
    Return occurrence identifiers from Neotoma and PBDB.

    With show=all, return expanded geography and age data for each occurrence.
    """
    # Initialization and parameter checks

    if request.args == {}:
        return jsonify(status_code=400, error='No parameters provided.')

    desc_obj = dict()
    occ_return = list()
    age_units = 'ma'
    geog_units = 'dec_deg_modern'
    full_return = False

    if show and show.lower() == 'all':
        full_return = True

    # Query the Neotoma Database (Occurrences)

    t0 = time.time()
    neotoma_base = 'http://apidev.neotomadb.org/v1/data/occurrences'
    payload = dict()

    if bbox:
        payload.update(bbox=bbox)

    if agescale and agescale.lower() == 'yr':
        age_scaler = 1
        age_units = 'yr'
        if minage:
            payload.update(ageyoung=int(minage))
        if maxage:
            payload.update(ageold=int(maxage))
    elif agescale and agescale.lower() == 'ka':
        age_scaler = 1e-03
        age_units = 'ka'
        if minage:
            payload.update(ageyoung=int(minage/age_scaler))
        if maxage:
            payload.update(ageold=int(maxage/age_scaler))
    else:
        age_scaler = 1e-06
        age_units = 'ma'
        if minage:
            payload.update(ageyoung=int(minage/age_scaler))
        if maxage:
            payload.update(ageold=int(maxage/age_scaler))

    if timerule and timerule.lower() == 'major':
        payload.update(agedocontain=0)
    elif timerule and timerule.lower() == 'overlap':
        payload.update(agedocontain=1)

    if includelower and includelower.lower() == 'false':
        payload.update(nametype='tax', taxonname=taxon)
    else:
        payload.update(nametype='base', taxonname=taxon)

    if limit:
        payload.update(limit=limit)
    else:
        payload.update(limit='999999')

    if offset:
        payload.update(offset=offset)

    neotoma_res = requests.get(neotoma_base, params=payload, timeout=None)

    if neotoma_res.status_code == 200:
        neotoma_json = neotoma_res.json()
        if 'data' in neotoma_json:
            for occ in neotoma_json['data']:
                occ_obj = dict()
                occ_id = 'neot:occ:' + str(occ['OccurID'])
                occ_obj.update(occ_id=occ_id, taxon=occ['TaxonName'])

                if full_return:
                    taxon_id = 'neot:txn:' + str(occ['TaxonID'])
                    if occ['AgeOlder'] and occ['AgeYounger']:
                        max_age = occ['AgeOlder'] * age_scaler
                        min_age = occ['AgeYounger'] * age_scaler
                    else:
                        max_age = None
                        min_age = None
                    lat = mean([occ['LatitudeNorth'],
                                occ['LatitudeSouth']])
                    # lat = round(lat, 4)
                    lon = mean([occ['LongitudeEast'],
                                occ['LongitudeWest']])
                    # lon = round(lon, 4)
                    occ_obj.update(taxon_id=taxon_id, max_age=max_age,
                                   min_age=min_age, age_units=age_units,
                                   lat=lat, lon=lon, geog_units=geog_units)

                occ_return.append(occ_obj)

            t1 = round(time.time()-t0, 5)
            desc_obj.update(neot_time=t1)
            desc_obj.update(neot_url=neotoma_res.url)
            desc_obj.update(neot_occs=len(neotoma_json['data']))

    # Query the Paleobiology Database (Occurrences)

    t0 = time.time()
    pbdb_base = 'http://paleobiodb.org/data1.2/occs/list.json'
    payload = dict()
    if full_return:
        payload.update(show='loc,coords,coll')
    payload.update(vocab='pbdb')

    if bbox:
        bbox_list = bbox.split(',')
        payload.update(lngmin=bbox_list[0], latmin=bbox_list[1],
                       lngmax=bbox_list[2], latmax=bbox_list[3])

    if agescale and agescale.lower() == 'yr':
        age_scaler = 1e06
        age_units = 'yr'
        if minage:
            payload.update(min_ma=float(minage)/age_scaler)
        if maxage:
            payload.update(max_ma=float(maxage)/age_scaler)
    elif agescale and agescale.lower() == 'ka':
        age_scaler = 1e03
        age_units = 'ka'
        if minage:
            payload.update(min_ma=float(minage)/age_scaler)
        if maxage:
            payload.update(max_ma=float(maxage)/age_scaler)
    else:
        age_scaler = 1
        age_units = 'ma'
        if minage:
            payload.update(min_ma=minage)
        if maxage:
            payload.update(max_ma=maxage)

    if timerule:
        payload.update(timerule=timerule)

    if includelower and includelower.lower() == 'false':
        payload.update(taxon_name=taxon)
    else:
        payload.update(base_name=taxon)

    if limit:
        payload.update(limit=limit)
    else:
        payload.update(limit='999999')

    if offset:
        payload.update(offset=offset)

    pbdb_res = requests.get(pbdb_base, params=payload, timeout=None)

    if pbdb_res.status_code == 200:
        pbdb_json = pbdb_res.json()
        if 'records' in pbdb_json:
            for occ in pbdb_json['records']:
                occ_obj = dict()
                occ_id = 'pbdb:occ:' + str(occ['occurrence_no'])
                occ_obj.update(occ_id=occ_id, taxon=occ['accepted_name'])

                if full_return:
                    taxon_id = 'pbdb:txn:' + str(occ['accepted_no'])
                    if occ['max_ma'] and occ['min_ma']:
                        max_age = float(occ['max_ma']) * age_scaler
                        min_age = float(occ['min_ma']) * age_scaler
                    else:
                        max_age = None
                        min_age = None
                    lat = float(occ['lat'])
                    lat = round(lat, 4)
                    lon = float(occ['lng'])
                    lon = round(lon, 4)
                    occ_obj.update(taxon_id=taxon_id, max_age=max_age,
                                   min_age=min_age, age_units=age_units,
                                   lat=lat, lon=lon, geog_units=geog_units)

                occ_return.append(occ_obj)

            t1 = round(time.time()-t0, 5)
            desc_obj.update(pbdb_time=t1)
            desc_obj.update(pbdb_url=pbdb_res.url)
            desc_obj.update(pbdb_occs=len(pbdb_json['records']))

    # Composite response

    return jsonify(description=desc_obj, data=occ_return)
<commit_msg>Change JSON object list to "records" standard<commit_after>"""
RESTful API controller.

Endpoint for queries on taxonomic occurrences in time and space.
"""

# import connexion
# from swagger_server.models.error_model import ErrorModel
# from swagger_server.models.occurrence import Occurrence
# from datetime import date, datetime
# from typing import List, Dict
# from six import iteritems
# from ..util import deserialize_date, deserialize_datetime
from flask import request, jsonify
import requests
import time
from statistics import mean

def occ(bbox=None, minage=None, maxage=None, agescale=None, timerule=None,
        taxon=None, includelower=None, limit=None, offset=None, show=None):
    """
    Return occurrence identifiers from Neotoma and PBDB.

    With show=all, return expanded geography and age data for each occurrence.
    """
    # Initialization and parameter checks

    if request.args == {}:
        return jsonify(status_code=400, error='No parameters provided.')

    desc_obj = dict()
    occ_return = list()
    age_units = 'ma'
    geog_units = 'dec_deg_modern'
    full_return = False

    if show and show.lower() == 'all':
        full_return = True

    # Query the Neotoma Database (Occurrences)

    t0 = time.time()
    neotoma_base = 'http://apidev.neotomadb.org/v1/data/occurrences'
    payload = dict()

    if bbox:
        payload.update(bbox=bbox)

    if agescale and agescale.lower() == 'yr':
        age_scaler = 1
        age_units = 'yr'
        if minage:
            payload.update(ageyoung=int(minage))
        if maxage:
            payload.update(ageold=int(maxage))
    elif agescale and agescale.lower() == 'ka':
        age_scaler = 1e-03
        age_units = 'ka'
        if minage:
            payload.update(ageyoung=int(minage/age_scaler))
        if maxage:
            payload.update(ageold=int(maxage/age_scaler))
    else:
        age_scaler = 1e-06
        age_units = 'ma'
        if minage:
            payload.update(ageyoung=int(minage/age_scaler))
        if maxage:
            payload.update(ageold=int(maxage/age_scaler))

    if timerule and timerule.lower() == 'major':
        payload.update(agedocontain=0)
    elif timerule and timerule.lower() == 'overlap':
        payload.update(agedocontain=1)

    if includelower and includelower.lower() == 'false':
        payload.update(nametype='tax', taxonname=taxon)
    else:
        payload.update(nametype='base', taxonname=taxon)

    if limit:
        payload.update(limit=limit)
    else:
        payload.update(limit='999999')

    if offset:
        payload.update(offset=offset)

    neotoma_res = requests.get(neotoma_base, params=payload, timeout=None)

    if neotoma_res.status_code == 200:
        neotoma_json = neotoma_res.json()
        if 'data' in neotoma_json:
            for occ in neotoma_json['data']:
                occ_obj = dict()
                occ_id = 'neot:occ:' + str(occ['OccurID'])
                occ_obj.update(occ_id=occ_id, taxon=occ['TaxonName'])

                if full_return:
                    taxon_id = 'neot:txn:' + str(occ['TaxonID'])
                    if occ['AgeOlder'] and occ['AgeYounger']:
                        max_age = occ['AgeOlder'] * age_scaler
                        min_age = occ['AgeYounger'] * age_scaler
                    else:
                        max_age = None
                        min_age = None
                    lat = mean([occ['LatitudeNorth'],
                                occ['LatitudeSouth']])
                    # lat = round(lat, 4)
                    lon = mean([occ['LongitudeEast'],
                                occ['LongitudeWest']])
                    # lon = round(lon, 4)
                    occ_obj.update(taxon_id=taxon_id, max_age=max_age,
                                   min_age=min_age, age_units=age_units,
                                   lat=lat, lon=lon, geog_units=geog_units)

                occ_return.append(occ_obj)

            t1 = round(time.time()-t0, 5)
            desc_obj.update(neot_time=t1)
            desc_obj.update(neot_url=neotoma_res.url)
            desc_obj.update(neot_occs=len(neotoma_json['data']))

    # Query the Paleobiology Database (Occurrences)

    t0 = time.time()
    pbdb_base = 'http://paleobiodb.org/data1.2/occs/list.json'
    payload = dict()
    if full_return:
        payload.update(show='loc,coords,coll')
    payload.update(vocab='pbdb')

    if bbox:
        bbox_list = bbox.split(',')
        payload.update(lngmin=bbox_list[0], latmin=bbox_list[1],
                       lngmax=bbox_list[2], latmax=bbox_list[3])

    if agescale and agescale.lower() == 'yr':
        age_scaler = 1e06
        age_units = 'yr'
        if minage:
            payload.update(min_ma=float(minage)/age_scaler)
        if maxage:
            payload.update(max_ma=float(maxage)/age_scaler)
    elif agescale and agescale.lower() == 'ka':
        age_scaler = 1e03
        age_units = 'ka'
        if minage:
            payload.update(min_ma=float(minage)/age_scaler)
        if maxage:
            payload.update(max_ma=float(maxage)/age_scaler)
    else:
        age_scaler = 1
        age_units = 'ma'
        if minage:
            payload.update(min_ma=minage)
        if maxage:
            payload.update(max_ma=maxage)

    if timerule:
        payload.update(timerule=timerule)

    if includelower and includelower.lower() == 'false':
        payload.update(taxon_name=taxon)
    else:
        payload.update(base_name=taxon)

    if limit:
        payload.update(limit=limit)
    else:
        payload.update(limit='999999')

    if offset:
        payload.update(offset=offset)

    pbdb_res = requests.get(pbdb_base, params=payload, timeout=None)

    if pbdb_res.status_code == 200:
        pbdb_json = pbdb_res.json()
        if 'records' in pbdb_json:
            for occ in pbdb_json['records']:
                occ_obj = dict()
                occ_id = 'pbdb:occ:' + str(occ['occurrence_no'])
                occ_obj.update(occ_id=occ_id, taxon=occ['accepted_name'])

                if full_return:
                    taxon_id = 'pbdb:txn:' + str(occ['accepted_no'])
                    if occ['max_ma'] and occ['min_ma']:
                        max_age = float(occ['max_ma']) * age_scaler
                        min_age = float(occ['min_ma']) * age_scaler
                    else:
                        max_age = None
                        min_age = None
                    lat = float(occ['lat'])
                    lat = round(lat, 4)
                    lon = float(occ['lng'])
                    lon = round(lon, 4)
                    occ_obj.update(taxon_id=taxon_id, max_age=max_age,
                                   min_age=min_age, age_units=age_units,
                                   lat=lat, lon=lon, geog_units=geog_units)

                occ_return.append(occ_obj)

            t1 = round(time.time()-t0, 5)
            desc_obj.update(pbdb_time=t1)
            desc_obj.update(pbdb_url=pbdb_res.url)
            desc_obj.update(pbdb_occs=len(pbdb_json['records']))

    # Composite response

    return jsonify(description=desc_obj, records=occ_return)
<|endoftext|>


<commit_before>import unittest

class TestScriptDB(unittest.TestCase):
    def test___init__(self):
        # script_d_b = ScriptDB(*args, **kwargs)
        assert True # TODO: implement your test here

    def test_at_typeclass_error(self):
        # script_d_b = ScriptDB(*args, **kwargs)
        # self.assertEqual(expected, script_d_b.at_typeclass_error())
        assert True # TODO: implement your test here

    def test_delete(self):
        # script_d_b = ScriptDB(*args, **kwargs)
        # self.assertEqual(expected, script_d_b.delete())
        assert True # TODO: implement your test here

if __name__ == '__main__':
    unittest.main()
<commit_msg>Some simple and proposed tests for the ScriptDB object. These pass OK with ./manage.py test, at least after running the server at least once.<commit_after>try:
    # this is an optimized version only available in later Django versions
    from django.utils.unittest import TestCase
except ImportError:
    # if the first fails, we use the old version
    from django.test import TestCase

from src.scripts.models import ScriptDB, ObjectDoesNotExist
from src.utils.create import create_script
from src.scripts import DoNothing
import unittest
from django.conf import settings

class TestScriptDB(TestCase):
    "Check the singleton/static ScriptDB object works correctly"
    def setUp(self):
        self.scr = create_script(DoNothing)

    def tearDown(self):
        self.scr.delete()
        del self.scr

    def test_delete(self):
        "Check the script is removed from the database"
        self.scr.delete()
        self.assertFalse(self.scr in ScriptDB.objects.get_all_scripts())

    def test_double_delete(self):
        "What should happen? Isn't it already deleted?"
        self.scr.delete()
        self.scr.delete()

    @unittest.skip("not implemented")
    def test___init__fails(self):  # Users should be told not to do this
        with self.assertRaises(Exception):
            ScriptDB()

    @unittest.skip("not implemented")
    def test_deleted_script_fails_start(self):
        "Would it ever be necessary to start a deleted script?"
        self.scr.delete()
        with self.assertRaises(ObjectDoesNotExist):  # See issue #509
            self.scr.start()
        # Check the script is not recreated as a side-effect
        self.assertFalse(self.scr in ScriptDB.objects.get_all_scripts())
        self.scr = create_script(DoNothing)  # for tearDown()

    @unittest.skip("not implemented")
    def test_deleted_script_is_invalid(self):
        "Can deleted scripts be said to be valid?"
        self.scr.delete()
        self.assertFalse(self.scr.is_valid())  # assertRaises? See issue #509
        self.scr = create_script(DoNothing)  # for tearDown()

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


<commit_before>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Various names for properties, transforms, etc."""

# All constants are for internal use only; no backwards-compatibility
# guarantees.

from __future__ import absolute_import

# Standard file names used for staging files.
from builtins import object

DATAFLOW_SDK_TARBALL_FILE = 'dataflow_python_sdk.tar'

# String constants related to sources framework
SOURCE_FORMAT = 'custom_source'
SOURCE_TYPE = 'CustomSourcesType'
SERIALIZED_SOURCE_KEY = 'serialized_source'

# In a released SDK, container tags are selected based on the SDK version.
# Unreleased versions use container versions based on values of
# BEAM_CONTAINER_VERSION and BEAM_FNAPI_CONTAINER_VERSION (see below).

# Update this version to the next version whenever there is a change that will
# require changes to legacy Dataflow worker execution environment.
BEAM_CONTAINER_VERSION = 'beam-master-20190208'
# Update this version to the next version whenever there is a change that
# requires changes to SDK harness container or SDK harness launcher.
BEAM_FNAPI_CONTAINER_VERSION = 'beam-master-20190213'

# TODO(BEAM-5939): Remove these shared names once Dataflow worker is updated.
PICKLED_MAIN_SESSION_FILE = 'pickled_main_session'
STAGED_PIPELINE_FILENAME = "pipeline.pb"
STAGED_PIPELINE_URL_METADATA_FIELD = "pipeline_url"

# Package names for different distributions
BEAM_PACKAGE_NAME = 'apache-beam'

# SDK identifiers for different distributions
BEAM_SDK_NAME = 'Apache Beam SDK for Python'
# TODO(BEAM-5393): End duplicated constants (see above).

DATAFLOW_CONTAINER_IMAGE_REPOSITORY = 'gcr.io/cloud-dataflow/v1beta3'

class TransformNames(object):
  """For internal use only; no backwards-compatibility guarantees.

  Transform strings as they are expected in the CloudWorkflow protos."""
  COLLECTION_TO_SINGLETON = 'CollectionToSingleton'
  COMBINE = 'CombineValues'
  CREATE_PCOLLECTION = 'CreateCollection'
  DO = 'ParallelDo'
  FLATTEN = 'Flatten'
  GROUP = 'GroupByKey'
  READ = 'ParallelRead'
  WRITE = 'ParallelWrite'

class PropertyNames(object):
  """For internal use only; no backwards-compatibility guarantees.

  Property strings as they are expected in the CloudWorkflow protos."""
  BIGQUERY_CREATE_DISPOSITION = 'create_disposition'
  BIGQUERY_DATASET = 'dataset'
  BIGQUERY_EXPORT_FORMAT = 'bigquery_export_format'
  BIGQUERY_FLATTEN_RESULTS = 'bigquery_flatten_results'
  BIGQUERY_KMS_KEY = 'bigquery_kms_key'
  BIGQUERY_PROJECT = 'project'
  BIGQUERY_QUERY = 'bigquery_query'
  BIGQUERY_SCHEMA = 'schema'
  BIGQUERY_TABLE = 'table'
  BIGQUERY_USE_LEGACY_SQL = 'bigquery_use_legacy_sql'
  BIGQUERY_WRITE_DISPOSITION = 'write_disposition'
  DISPLAY_DATA = 'display_data'
  ELEMENT = 'element'
  ELEMENTS = 'elements'
  ENCODING = 'encoding'
  FILE_PATTERN = 'filepattern'
  FILE_NAME_PREFIX = 'filename_prefix'
  FILE_NAME_SUFFIX = 'filename_suffix'
  FORMAT = 'format'
  INPUTS = 'inputs'
  IMPULSE_ELEMENT = 'impulse_element'
  NON_PARALLEL_INPUTS = 'non_parallel_inputs'
  NUM_SHARDS = 'num_shards'
  OUT = 'out'
  OUTPUT = 'output'
  OUTPUT_INFO = 'output_info'
  OUTPUT_NAME = 'output_name'
  PARALLEL_INPUT = 'parallel_input'
  PUBSUB_ID_LABEL = 'pubsub_id_label'
  PUBSUB_SERIALIZED_ATTRIBUTES_FN = 'pubsub_serialized_attributes_fn'
  PUBSUB_SUBSCRIPTION = 'pubsub_subscription'
  PUBSUB_TIMESTAMP_ATTRIBUTE = 'pubsub_timestamp_label'
  PUBSUB_TOPIC = 'pubsub_topic'
  SERIALIZED_FN = 'serialized_fn'
  SHARD_NAME_TEMPLATE = 'shard_template'
  SOURCE_STEP_INPUT = 'custom_source_step_input'
  STEP_NAME = 'step_name'
  USER_FN = 'user_fn'
  USER_NAME = 'user_name'
  VALIDATE_SINK = 'validate_sink'
  VALIDATE_SOURCE = 'validate_source'
  VALUE = 'value'
  WINDOWING_STRATEGY = 'windowing_strategy'
<commit_msg>Update Dataflow container version to 20190213<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Various names for properties, transforms, etc."""

# All constants are for internal use only; no backwards-compatibility
# guarantees.

from __future__ import absolute_import

# Standard file names used for staging files.
from builtins import object

DATAFLOW_SDK_TARBALL_FILE = 'dataflow_python_sdk.tar'

# String constants related to sources framework
SOURCE_FORMAT = 'custom_source'
SOURCE_TYPE = 'CustomSourcesType'
SERIALIZED_SOURCE_KEY = 'serialized_source'

# In a released SDK, container tags are selected based on the SDK version.
# Unreleased versions use container versions based on values of
# BEAM_CONTAINER_VERSION and BEAM_FNAPI_CONTAINER_VERSION (see below).

# Update this version to the next version whenever there is a change that will
# require changes to legacy Dataflow worker execution environment.
BEAM_CONTAINER_VERSION = 'beam-master-20190213'
# Update this version to the next version whenever there is a change that
# requires changes to SDK harness container or SDK harness launcher.
BEAM_FNAPI_CONTAINER_VERSION = 'beam-master-20190213'

# TODO(BEAM-5939): Remove these shared names once Dataflow worker is updated.
PICKLED_MAIN_SESSION_FILE = 'pickled_main_session'
STAGED_PIPELINE_FILENAME = "pipeline.pb"
STAGED_PIPELINE_URL_METADATA_FIELD = "pipeline_url"

# Package names for different distributions
BEAM_PACKAGE_NAME = 'apache-beam'

# SDK identifiers for different distributions
BEAM_SDK_NAME = 'Apache Beam SDK for Python'
# TODO(BEAM-5393): End duplicated constants (see above).

DATAFLOW_CONTAINER_IMAGE_REPOSITORY = 'gcr.io/cloud-dataflow/v1beta3'

class TransformNames(object):
  """For internal use only; no backwards-compatibility guarantees.

  Transform strings as they are expected in the CloudWorkflow protos."""
  COLLECTION_TO_SINGLETON = 'CollectionToSingleton'
  COMBINE = 'CombineValues'
  CREATE_PCOLLECTION = 'CreateCollection'
  DO = 'ParallelDo'
  FLATTEN = 'Flatten'
  GROUP = 'GroupByKey'
  READ = 'ParallelRead'
  WRITE = 'ParallelWrite'

class PropertyNames(object):
  """For internal use only; no backwards-compatibility guarantees.

  Property strings as they are expected in the CloudWorkflow protos."""
  BIGQUERY_CREATE_DISPOSITION = 'create_disposition'
  BIGQUERY_DATASET = 'dataset'
  BIGQUERY_EXPORT_FORMAT = 'bigquery_export_format'
  BIGQUERY_FLATTEN_RESULTS = 'bigquery_flatten_results'
  BIGQUERY_KMS_KEY = 'bigquery_kms_key'
  BIGQUERY_PROJECT = 'project'
  BIGQUERY_QUERY = 'bigquery_query'
  BIGQUERY_SCHEMA = 'schema'
  BIGQUERY_TABLE = 'table'
  BIGQUERY_USE_LEGACY_SQL = 'bigquery_use_legacy_sql'
  BIGQUERY_WRITE_DISPOSITION = 'write_disposition'
  DISPLAY_DATA = 'display_data'
  ELEMENT = 'element'
  ELEMENTS = 'elements'
  ENCODING = 'encoding'
  FILE_PATTERN = 'filepattern'
  FILE_NAME_PREFIX = 'filename_prefix'
  FILE_NAME_SUFFIX = 'filename_suffix'
  FORMAT = 'format'
  INPUTS = 'inputs'
  IMPULSE_ELEMENT = 'impulse_element'
  NON_PARALLEL_INPUTS = 'non_parallel_inputs'
  NUM_SHARDS = 'num_shards'
  OUT = 'out'
  OUTPUT = 'output'
  OUTPUT_INFO = 'output_info'
  OUTPUT_NAME = 'output_name'
  PARALLEL_INPUT = 'parallel_input'
  PUBSUB_ID_LABEL = 'pubsub_id_label'
  PUBSUB_SERIALIZED_ATTRIBUTES_FN = 'pubsub_serialized_attributes_fn'
  PUBSUB_SUBSCRIPTION = 'pubsub_subscription'
  PUBSUB_TIMESTAMP_ATTRIBUTE = 'pubsub_timestamp_label'
  PUBSUB_TOPIC = 'pubsub_topic'
  SERIALIZED_FN = 'serialized_fn'
  SHARD_NAME_TEMPLATE = 'shard_template'
  SOURCE_STEP_INPUT = 'custom_source_step_input'
  STEP_NAME = 'step_name'
  USER_FN = 'user_fn'
  USER_NAME = 'user_name'
  VALIDATE_SINK = 'validate_sink'
  VALIDATE_SOURCE = 'validate_source'
  VALUE = 'value'
  WINDOWING_STRATEGY = 'windowing_strategy'
<|endoftext|>


<commit_before>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""The TestStream implementation for the DirectRunner

The DirectRunner implements TestStream as the _TestStream class which is used
to store the events in memory, the _WatermarkController which is used to set the
watermark and emit events, and the multiplexer which sends events to the correct
tagged PCollection.
"""

# pytype: skip-file

from __future__ import absolute_import
from __future__ import print_function

import itertools
import logging
from queue import Empty as EmptyException
from queue import Queue
from threading import Thread

import grpc

from apache_beam import ParDo
from apache_beam import coders
from apache_beam import pvalue
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.portability.api import beam_runner_api_pb2_grpc
from apache_beam.testing.test_stream import ElementEvent
from apache_beam.testing.test_stream import ProcessingTimeEvent
from apache_beam.testing.test_stream import WatermarkEvent
from apache_beam.transforms import PTransform
from apache_beam.transforms import core
from apache_beam.transforms import window
from apache_beam.transforms.window import TimestampedValue
from apache_beam.utils import timestamp
from apache_beam.utils.timestamp import Duration
from apache_beam.utils.timestamp import Timestamp

_LOGGER = logging.getLogger(__name__)

class _EndOfStream:
  pass

class _WatermarkController(PTransform):
  """A runner-overridable PTransform Primitive to control the watermark.

  Expected implementation behavior:
   - If the instance recieves a WatermarkEvent, it sets its output watermark to
     the specified value then drops the event.
   - If the instance receives an ElementEvent, it emits all specified elements
     to the Global Window with the event time set to the element's timestamp.
  """
  def __init__(self, output_tag):
    self.output_tag = output_tag

  def get_windowing(self, _):
    return core.Windowing(window.GlobalWindows())

  def expand(self, pcoll):
    ret = pvalue.PCollection.from_(pcoll)
    ret.tag = self.output_tag
    return ret

class _ExpandableTestStream(PTransform):
  def __init__(self, test_stream):
    self.test_stream = test_stream

  def expand(self, pbegin):
    """Expands the TestStream into the DirectRunner implementation.

    Takes the TestStream transform and creates a _TestStream -> multiplexer ->
    _WatermarkController.
    """

    assert isinstance(pbegin, pvalue.PBegin)

    # If there is only one tag there is no need to add the multiplexer.
    if len(self.test_stream.output_tags) == 1:
      return (
          pbegin
          | _TestStream(
              self.test_stream.output_tags,
              events=self.test_stream._events,
              coder=self.test_stream.coder,
              endpoint=self.test_stream._endpoint)
          | _WatermarkController(list(self.test_stream.output_tags)[0]))

    # Multiplex to the correct PCollection based upon the event tag.
    def mux(event):
      if event.tag:
        yield pvalue.TaggedOutput(event.tag, event)
      else:
        yield event

    mux_output = (
        pbegin
        | _TestStream(
            self.test_stream.output_tags,
            events=self.test_stream._events,
            coder=self.test_stream.coder,
            endpoint=self.test_stream._endpoint)
        | 'TestStream Multiplexer' >> ParDo(mux).with_outputs())

    # Apply a way to control the watermark per output. It is necessary to
    # have an individual _WatermarkController per PCollection because the
    # calculation of the input watermark of a transform is based on the event
    # timestamp of the elements flowing through it. Meaning, it is impossible
    # to control the output watermarks of the individual PCollections solely
    # on the event timestamps.
    outputs = {}
    for tag in self.test_stream.output_tags:
      label = '_WatermarkController[{}]'.format(tag)
      outputs[tag] = (mux_output[tag] | label >> _WatermarkController(tag))

    return outputs

class _TestStream(PTransform):
  """Test stream that generates events on an unbounded PCollection of elements.

  Each event emits elements, advances the watermark or advances the processing
  time.  After all of the specified elements are emitted, ceases to produce
  output.

  Expected implementation behavior:
   - If the instance receives a WatermarkEvent with the WATERMARK_CONTROL_TAG
     then the instance sets its own watermark hold at the specified value and
     drops the event.
   - If the instance receives any other WatermarkEvent or ElementEvent, it
     passes it to the consumer.
  """

  # This tag is used on WatermarkEvents to control the watermark at the root
  # TestStream.
  WATERMARK_CONTROL_TAG = '_TestStream_Watermark'

  def __init__(
      self,
      output_tags,
      coder=coders.FastPrimitivesCoder(),
      events=None,
      endpoint=None):
    assert coder is not None
    self.coder = coder
    self._raw_events = events
    self._events = self._add_watermark_advancements(output_tags, events)
    self.output_tags = output_tags
    self.endpoint = endpoint

  def _watermark_starts(self, output_tags):
    """Sentinel values to hold the watermark of outputs to -inf.

    The output watermarks of the output PCollections (fake unbounded sources) in
    a TestStream are controlled by watermark holds. This sets the hold of each
    output PCollection so that the individual holds can be controlled by the
    given events.
    """
    return [WatermarkEvent(timestamp.MIN_TIMESTAMP, tag) for tag in output_tags]

  def _watermark_stops(self, output_tags):
    """Sentinel values to close the watermark of outputs."""
    return [WatermarkEvent(timestamp.MAX_TIMESTAMP, tag) for tag in output_tags]

  def _test_stream_start(self):
    """Sentinel value to move the watermark hold of the TestStream to +inf.

    This sets a hold to +inf such that the individual holds of the output
    PCollections are allowed to modify their individial output watermarks with
    their holds. This is because the calculation of the output watermark is a
    min over all input watermarks.
    """
    return [
        WatermarkEvent(
            timestamp.MAX_TIMESTAMP - timestamp.TIME_GRANULARITY,
            _TestStream.WATERMARK_CONTROL_TAG)
    ]

  def _test_stream_stop(self):
    """Sentinel value to close the watermark of the TestStream."""
    return [
        WatermarkEvent(
            timestamp.MAX_TIMESTAMP, _TestStream.WATERMARK_CONTROL_TAG)
    ]

  def _test_stream_init(self):
    """Sentinel value to hold the watermark of the TestStream to -inf.

    This sets a hold to ensure that the output watermarks of the output
    PCollections do not advance to +inf before their watermark holds are set.
    """
    return [
        WatermarkEvent(
            timestamp.MIN_TIMESTAMP, _TestStream.WATERMARK_CONTROL_TAG)
    ]

  def _set_up(self, output_tags):
    return (
        self._test_stream_init() + self._watermark_starts(output_tags) +
        self._test_stream_start())

  def _tear_down(self, output_tags):
    return self._watermark_stops(output_tags) + self._test_stream_stop()

  def _add_watermark_advancements(self, output_tags, events):
    """Adds watermark advancements to the given events.

    The following watermark advancements can be done on the runner side.
    However, it makes the logic on the runner side much more complicated than
    it needs to be.

    In order for watermarks to be properly advanced in a TestStream, a specific
    sequence of watermark holds must be sent:

    1. Hold the root watermark at -inf (this prevents the pipeline from
       immediately returning).
    2. Hold the watermarks at the WatermarkControllerss at -inf (this prevents
       the pipeline from immediately returning).
    3. Advance the root watermark to +inf - 1 (this allows the downstream
       WatermarkControllers to control their watermarks via holds).
    4. Advance watermarks as normal.
    5. Advance WatermarkController watermarks to +inf
    6. Advance root watermark to +inf.
    """
    if not events:
      return []

    return self._set_up(output_tags) + events + self._tear_down(output_tags)

  def get_windowing(self, unused_inputs):
    return core.Windowing(window.GlobalWindows())

  def expand(self, pcoll):
    return pvalue.PCollection(pcoll.pipeline, is_bounded=False)

  def _infer_output_coder(self, input_type=None, input_coder=None):
    return self.coder

  @staticmethod
  def events_from_script(events):
    """Yields the in-memory events.
    """
    return itertools.chain(events)

  @staticmethod
  def _stream_events_from_rpc(endpoint, output_tags, coder, channel, is_alive):
    """Yields the events received from the given endpoint.

    This is the producer thread that reads events from the TestStreamService and
    puts them onto the shared queue. At the end of the stream, an _EndOfStream
    is placed on the channel to signify a successful end.
    """
    stub_channel = grpc.insecure_channel(endpoint)
    stub = beam_runner_api_pb2_grpc.TestStreamServiceStub(stub_channel)

    # Request the PCollections that we are looking for from the service.
    event_request = beam_runner_api_pb2.EventsRequest(
        output_ids=[str(tag) for tag in output_tags])

    event_stream = stub.Events(event_request)
    try:
      for e in event_stream:
        channel.put(_TestStream.test_stream_payload_to_events(e, coder))
        if not is_alive():
          return
    except grpc.RpcError as e:
      # This happens when the Python interpreter shuts down or whn in a
      # notebook environment when the kernel is interrupted.
      if e.code() == grpc.StatusCode.UNAVAILABLE:
        return
      raise e
    finally:
      # Gracefully stop the job if there is an exception.
      channel.put(_EndOfStream())

  @staticmethod
  def events_from_rpc(endpoint, output_tags, coder, evaluation_context):
    """Yields the events received from the given endpoint.

    This method starts a new thread that reads from the TestStreamService and
    puts the events onto a shared queue. This method then yields all elements
    from the queue. Unfortunately, this is necessary because the GRPC API does
    not allow for non-blocking calls when utilizing a streaming RPC. It is
    officially suggested from the docs to use a producer/consumer pattern to
    handle streaming RPCs. By doing so, this gives this method control over when
    to cancel reading from the RPC if the server takes too long to respond.
    """
    # Shared variable with the producer queue. This shuts down the producer if
    # the consumer exits early.
    shutdown_requested = False

    def is_alive():
      return not (shutdown_requested or evaluation_context.shutdown_requested)

    # The shared queue that allows the producer and consumer to communicate.
    channel = Queue()  # type: Queue[Union[test_stream.Event, _EndOfStream]]
    event_stream = Thread(
        target=_TestStream._stream_events_from_rpc,
        args=(endpoint, output_tags, coder, channel, is_alive))
    event_stream.setDaemon(True)
    event_stream.start()

    # This pumps the shared queue for events until the _EndOfStream sentinel is
    # reached. If the TestStreamService takes longer than expected, the queue
    # will timeout and an EmptyException will be raised. This also sets the
    # shared is_alive sentinel to shut down the producer.
    while True:
      try:
        # Raise an EmptyException if there are no events during the last timeout
        # period.
        event = channel.get(timeout=30)
        if isinstance(event, _EndOfStream):
          break
        yield event
      except EmptyException as e:
        _LOGGER.warning(
            'TestStream timed out waiting for new events from service.'
            ' Stopping pipeline.')
        shutdown_requested = True
        raise e

  @staticmethod
  def test_stream_payload_to_events(payload, coder):
    """Returns a TestStream Python event object from a TestStream event Proto.
    """
    if payload.HasField('element_event'):
      element_event = payload.element_event
      elements = [
          TimestampedValue(
              coder.decode(e.encoded_element), Timestamp(micros=e.timestamp))
          for e in element_event.elements
      ]
      return ElementEvent(timestamped_values=elements, tag=element_event.tag)

    if payload.HasField('watermark_event'):
      watermark_event = payload.watermark_event
      return WatermarkEvent(
          Timestamp(micros=watermark_event.new_watermark),
          tag=watermark_event.tag)

    if payload.HasField('processing_time_event'):
      processing_time_event = payload.processing_time_event
      return ProcessingTimeEvent(
          Duration(micros=processing_time_event.advance_duration))

    raise RuntimeError(
        'Received a proto without the specified fields: {}'.format(payload))
<commit_msg>Add CANCELLED to non error codes for test stream events from grpc<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""The TestStream implementation for the DirectRunner

The DirectRunner implements TestStream as the _TestStream class which is used
to store the events in memory, the _WatermarkController which is used to set the
watermark and emit events, and the multiplexer which sends events to the correct
tagged PCollection.
"""

# pytype: skip-file

from __future__ import absolute_import
from __future__ import print_function

import itertools
import logging
from queue import Empty as EmptyException
from queue import Queue
from threading import Thread

import grpc

from apache_beam import ParDo
from apache_beam import coders
from apache_beam import pvalue
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.portability.api import beam_runner_api_pb2_grpc
from apache_beam.testing.test_stream import ElementEvent
from apache_beam.testing.test_stream import ProcessingTimeEvent
from apache_beam.testing.test_stream import WatermarkEvent
from apache_beam.transforms import PTransform
from apache_beam.transforms import core
from apache_beam.transforms import window
from apache_beam.transforms.window import TimestampedValue
from apache_beam.utils import timestamp
from apache_beam.utils.timestamp import Duration
from apache_beam.utils.timestamp import Timestamp

_LOGGER = logging.getLogger(__name__)

class _EndOfStream:
  pass

class _WatermarkController(PTransform):
  """A runner-overridable PTransform Primitive to control the watermark.

  Expected implementation behavior:
   - If the instance recieves a WatermarkEvent, it sets its output watermark to
     the specified value then drops the event.
   - If the instance receives an ElementEvent, it emits all specified elements
     to the Global Window with the event time set to the element's timestamp.
  """
  def __init__(self, output_tag):
    self.output_tag = output_tag

  def get_windowing(self, _):
    return core.Windowing(window.GlobalWindows())

  def expand(self, pcoll):
    ret = pvalue.PCollection.from_(pcoll)
    ret.tag = self.output_tag
    return ret

class _ExpandableTestStream(PTransform):
  def __init__(self, test_stream):
    self.test_stream = test_stream

  def expand(self, pbegin):
    """Expands the TestStream into the DirectRunner implementation.

    Takes the TestStream transform and creates a _TestStream -> multiplexer ->
    _WatermarkController.
    """

    assert isinstance(pbegin, pvalue.PBegin)

    # If there is only one tag there is no need to add the multiplexer.
    if len(self.test_stream.output_tags) == 1:
      return (
          pbegin
          | _TestStream(
              self.test_stream.output_tags,
              events=self.test_stream._events,
              coder=self.test_stream.coder,
              endpoint=self.test_stream._endpoint)
          | _WatermarkController(list(self.test_stream.output_tags)[0]))

    # Multiplex to the correct PCollection based upon the event tag.
    def mux(event):
      if event.tag:
        yield pvalue.TaggedOutput(event.tag, event)
      else:
        yield event

    mux_output = (
        pbegin
        | _TestStream(
            self.test_stream.output_tags,
            events=self.test_stream._events,
            coder=self.test_stream.coder,
            endpoint=self.test_stream._endpoint)
        | 'TestStream Multiplexer' >> ParDo(mux).with_outputs())

    # Apply a way to control the watermark per output. It is necessary to
    # have an individual _WatermarkController per PCollection because the
    # calculation of the input watermark of a transform is based on the event
    # timestamp of the elements flowing through it. Meaning, it is impossible
    # to control the output watermarks of the individual PCollections solely
    # on the event timestamps.
    outputs = {}
    for tag in self.test_stream.output_tags:
      label = '_WatermarkController[{}]'.format(tag)
      outputs[tag] = (mux_output[tag] | label >> _WatermarkController(tag))

    return outputs

class _TestStream(PTransform):
  """Test stream that generates events on an unbounded PCollection of elements.

  Each event emits elements, advances the watermark or advances the processing
  time.  After all of the specified elements are emitted, ceases to produce
  output.

  Expected implementation behavior:
   - If the instance receives a WatermarkEvent with the WATERMARK_CONTROL_TAG
     then the instance sets its own watermark hold at the specified value and
     drops the event.
   - If the instance receives any other WatermarkEvent or ElementEvent, it
     passes it to the consumer.
  """

  # This tag is used on WatermarkEvents to control the watermark at the root
  # TestStream.
  WATERMARK_CONTROL_TAG = '_TestStream_Watermark'

  def __init__(
      self,
      output_tags,
      coder=coders.FastPrimitivesCoder(),
      events=None,
      endpoint=None):
    assert coder is not None
    self.coder = coder
    self._raw_events = events
    self._events = self._add_watermark_advancements(output_tags, events)
    self.output_tags = output_tags
    self.endpoint = endpoint

  def _watermark_starts(self, output_tags):
    """Sentinel values to hold the watermark of outputs to -inf.

    The output watermarks of the output PCollections (fake unbounded sources) in
    a TestStream are controlled by watermark holds. This sets the hold of each
    output PCollection so that the individual holds can be controlled by the
    given events.
    """
    return [WatermarkEvent(timestamp.MIN_TIMESTAMP, tag) for tag in output_tags]

  def _watermark_stops(self, output_tags):
    """Sentinel values to close the watermark of outputs."""
    return [WatermarkEvent(timestamp.MAX_TIMESTAMP, tag) for tag in output_tags]

  def _test_stream_start(self):
    """Sentinel value to move the watermark hold of the TestStream to +inf.

    This sets a hold to +inf such that the individual holds of the output
    PCollections are allowed to modify their individial output watermarks with
    their holds. This is because the calculation of the output watermark is a
    min over all input watermarks.
    """
    return [
        WatermarkEvent(
            timestamp.MAX_TIMESTAMP - timestamp.TIME_GRANULARITY,
            _TestStream.WATERMARK_CONTROL_TAG)
    ]

  def _test_stream_stop(self):
    """Sentinel value to close the watermark of the TestStream."""
    return [
        WatermarkEvent(
            timestamp.MAX_TIMESTAMP, _TestStream.WATERMARK_CONTROL_TAG)
    ]

  def _test_stream_init(self):
    """Sentinel value to hold the watermark of the TestStream to -inf.

    This sets a hold to ensure that the output watermarks of the output
    PCollections do not advance to +inf before their watermark holds are set.
    """
    return [
        WatermarkEvent(
            timestamp.MIN_TIMESTAMP, _TestStream.WATERMARK_CONTROL_TAG)
    ]

  def _set_up(self, output_tags):
    return (
        self._test_stream_init() + self._watermark_starts(output_tags) +
        self._test_stream_start())

  def _tear_down(self, output_tags):
    return self._watermark_stops(output_tags) + self._test_stream_stop()

  def _add_watermark_advancements(self, output_tags, events):
    """Adds watermark advancements to the given events.

    The following watermark advancements can be done on the runner side.
    However, it makes the logic on the runner side much more complicated than
    it needs to be.

    In order for watermarks to be properly advanced in a TestStream, a specific
    sequence of watermark holds must be sent:

    1. Hold the root watermark at -inf (this prevents the pipeline from
       immediately returning).
    2. Hold the watermarks at the WatermarkControllerss at -inf (this prevents
       the pipeline from immediately returning).
    3. Advance the root watermark to +inf - 1 (this allows the downstream
       WatermarkControllers to control their watermarks via holds).
    4. Advance watermarks as normal.
    5. Advance WatermarkController watermarks to +inf
    6. Advance root watermark to +inf.
    """
    if not events:
      return []

    return self._set_up(output_tags) + events + self._tear_down(output_tags)

  def get_windowing(self, unused_inputs):
    return core.Windowing(window.GlobalWindows())

  def expand(self, pcoll):
    return pvalue.PCollection(pcoll.pipeline, is_bounded=False)

  def _infer_output_coder(self, input_type=None, input_coder=None):
    return self.coder

  @staticmethod
  def events_from_script(events):
    """Yields the in-memory events.
    """
    return itertools.chain(events)

  @staticmethod
  def _stream_events_from_rpc(endpoint, output_tags, coder, channel, is_alive):
    """Yields the events received from the given endpoint.

    This is the producer thread that reads events from the TestStreamService and
    puts them onto the shared queue. At the end of the stream, an _EndOfStream
    is placed on the channel to signify a successful end.
    """
    stub_channel = grpc.insecure_channel(endpoint)
    stub = beam_runner_api_pb2_grpc.TestStreamServiceStub(stub_channel)

    # Request the PCollections that we are looking for from the service.
    event_request = beam_runner_api_pb2.EventsRequest(
        output_ids=[str(tag) for tag in output_tags])

    event_stream = stub.Events(event_request)
    try:
      for e in event_stream:
        channel.put(_TestStream.test_stream_payload_to_events(e, coder))
        if not is_alive():
          return
    except grpc.RpcError as e:
      # Do not raise an exception in the non-error status codes. These can occur
      # when the Python interpreter shuts down or when in a notebook environment
      # when the kernel is interrupted.
      if e.code() in (grpc.StatusCode.CANCELLED, grpc.StatusCode.UNAVAILABLE):
        return
      raise e
    finally:
      # Gracefully stop the job if there is an exception.
      channel.put(_EndOfStream())

  @staticmethod
  def events_from_rpc(endpoint, output_tags, coder, evaluation_context):
    """Yields the events received from the given endpoint.

    This method starts a new thread that reads from the TestStreamService and
    puts the events onto a shared queue. This method then yields all elements
    from the queue. Unfortunately, this is necessary because the GRPC API does
    not allow for non-blocking calls when utilizing a streaming RPC. It is
    officially suggested from the docs to use a producer/consumer pattern to
    handle streaming RPCs. By doing so, this gives this method control over when
    to cancel reading from the RPC if the server takes too long to respond.
    """
    # Shared variable with the producer queue. This shuts down the producer if
    # the consumer exits early.
    shutdown_requested = False

    def is_alive():
      return not (shutdown_requested or evaluation_context.shutdown_requested)

    # The shared queue that allows the producer and consumer to communicate.
    channel = Queue()  # type: Queue[Union[test_stream.Event, _EndOfStream]]
    event_stream = Thread(
        target=_TestStream._stream_events_from_rpc,
        args=(endpoint, output_tags, coder, channel, is_alive))
    event_stream.setDaemon(True)
    event_stream.start()

    # This pumps the shared queue for events until the _EndOfStream sentinel is
    # reached. If the TestStreamService takes longer than expected, the queue
    # will timeout and an EmptyException will be raised. This also sets the
    # shared is_alive sentinel to shut down the producer.
    while True:
      try:
        # Raise an EmptyException if there are no events during the last timeout
        # period.
        event = channel.get(timeout=30)
        if isinstance(event, _EndOfStream):
          break
        yield event
      except EmptyException as e:
        _LOGGER.warning(
            'TestStream timed out waiting for new events from service.'
            ' Stopping pipeline.')
        shutdown_requested = True
        raise e

  @staticmethod
  def test_stream_payload_to_events(payload, coder):
    """Returns a TestStream Python event object from a TestStream event Proto.
    """
    if payload.HasField('element_event'):
      element_event = payload.element_event
      elements = [
          TimestampedValue(
              coder.decode(e.encoded_element), Timestamp(micros=e.timestamp))
          for e in element_event.elements
      ]
      return ElementEvent(timestamped_values=elements, tag=element_event.tag)

    if payload.HasField('watermark_event'):
      watermark_event = payload.watermark_event
      return WatermarkEvent(
          Timestamp(micros=watermark_event.new_watermark),
          tag=watermark_event.tag)

    if payload.HasField('processing_time_event'):
      processing_time_event = payload.processing_time_event
      return ProcessingTimeEvent(
          Duration(micros=processing_time_event.advance_duration))

    raise RuntimeError(
        'Received a proto without the specified fields: {}'.format(payload))
<|endoftext|>