<commit_before>import bpy, bpy.path, os.path, struct

def get_texture_path(image):
    return os.path.normpath(bpy.path.abspath(image.filepath))

SHADER_TYPES = {
    'RetroShader': b'RSHD',
    'RetroDynamicShader': b'RDYN',
    'RetroDynamicAlphaShader': b'RDAL',
    'RetroDynamicCharacterShader': b'RCHR',
}

PASS_TYPE = {
    'Lightmap': b'LMAP',
    'Diffuse': b'DIFF',
    'Emissive': b'EMIS',
    'Specular': b'SPEC',
    'ExtendedSpecular': b'ESPC',
    'Reflection': b'REFL',
    'IndirectTex': b'INDR',
    'Alpha': b'ALPH',
}

def write_chunks(writebuf, mat_obj, mesh_obj):

    if not mat_obj.use_nodes:
        raise RuntimeError("HMDL *requires* that shader nodes are used; '{0}' does not".format(mat_obj.name))

    if 'Output' not in mat_obj.node_tree.nodes or mat_obj.node_tree.nodes['Output'].type != 'GROUP':
        raise RuntimeError("HMDL *requires* that an group shader node named 'Output' is present")

    # Root (output) node
    output_node = mat_obj.node_tree.nodes['Output']
    if output_node.node_tree.name not in SHADER_TYPES:
        raise RuntimeError("HMDL *requires* one of the RetroShader group nodes for the 'Output' node")
    writebuf(SHADER_TYPES[output_node.node_tree.name])

    # Count sockets
    chunk_count = 0
    for inp in output_node.inputs:
        if inp.name in PASS_TYPE:
            if inp.is_linked:
                chunk_count += 1
            else:
                # Color pass
                color_set = False
                if inp.type == 'VALUE':
                    color_set = bool(inp.default_value)
                else:
                    for comp in inp.default_value:
                        color_set |= bool(comp)
                if color_set:
                    chunk_count += 1

    writebuf(struct.pack('I', chunk_count))

    # Enumerate sockets
    for inp in output_node.inputs:
        if inp.name in PASS_TYPE:
            pass_fourcc = PASS_TYPE[inp.name]
            if inp.is_linked:
                socket = inp.links[0].from_socket
                node = socket.node
                if node.type != 'TEX_IMAGE':
                    raise RuntimeError("HMDL requires all group node inputs connect to Image Texture nodes")

                if not node.image:
                    raise RuntimeError("HMDL texture nodes must specify an image object")

                if not node.inputs['Vector'].is_linked:
                    raise RuntimeError("HMDL texture nodes must have a 'Texture Coordinate', 'UV Map' or 'Group' UV modifier node linked")

                # Determine matrix generator type
                tex_coord_source = 0xff
                uv_anim_type = 0xff
                uv_anim_args = []
                soc_from = node.inputs['Vector'].links[0].from_socket

                if soc_from.node.type == 'GROUP':
                    if soc_from.node.node_tree.name.startswith('RetroUVMode'):
                        uv_anim_type = int(soc_from.node.node_tree.name[11:12])
                        if len(soc_from.node.inputs)-1:
                            for s in range(len(soc_from.node.inputs)-1):
                                soc = soc_from.node.inputs[s+1]
                                if len(soc.links):
                                    raise RuntimeError("UV Modifier nodes may not have parameter links (default values only)")
                                if soc.type == 'VALUE':
                                    uv_anim_args.append(soc.default_value)
                                else:
                                    uv_anim_args.append(soc.default_value[0])
                                    uv_anim_args.append(soc.default_value[1])
                    soc_from = soc_from.node.inputs[0].links[0].from_socket

                elif soc_from.node.type == 'UVMAP' or soc_from.node.type == 'TEX_COORD':
                    pass

                else:
                    raise RuntimeError("HMDL texture nodes must have a 'Texture Coordinate', 'UV Map' or 'Group' UV modifier node linked")

                if soc_from.node.type != 'UVMAP' and soc_from.node.type != 'TEX_COORD':
                    raise RuntimeError("Matrix animator nodes must connect to 'Texture Coordinate' or 'UV Map' node")

                # Resolve map and matrix index
                node_label = soc_from.node.label
                matrix_idx = None
                if node_label.startswith('MTX_'):
                    matrix_idx = int(node_label[4:])

                if soc_from.name == 'UV':
                    if hasattr(soc_from.node, 'uv_map'):
                        uv_name = soc_from.node.uv_map
                        uv_idx = mesh_obj.data.uv_layers.find(uv_name)
                        if uv_idx == -1:
                            raise RuntimeError('UV Layer "%s" doesn\'t exist' % uv_name)
                        tex_coord_source = uv_idx + 2
                    else:
                        tex_coord_source = 2

                elif soc_from.name == 'Normal':
                    tex_coord_source = 1

                elif soc_from.name == 'Window':
                    tex_coord_source = 0

                else:
                    raise RuntimeError("Only the 'UV', 'Normal' and 'Window' sockets may be used from 'Texture Coordinate' nodes")

                alpha = False
                if socket.name == 'Alpha':
                    alpha = True

                writebuf(b'PASS')
                writebuf(pass_fourcc)
                path = get_texture_path(node.image)
                writebuf(struct.pack('I', len(path)))
                writebuf(path.encode())
                writebuf(struct.pack('B', tex_coord_source))
                writebuf(struct.pack('B', uv_anim_type))
                writebuf(struct.pack('I', len(uv_anim_args)))
                for arg in uv_anim_args:
                    writebuf(struct.pack('f', arg))
                writebuf(struct.pack('B', alpha))

            else:
                # Color pass
                color_set = False
                if inp.type == 'VALUE':
                    color_set = bool(inp.default_value)
                else:
                    for comp in inp.default_value:
                        color_set |= bool(comp)

                if color_set:
                    writebuf(b'CLR ')
                    writebuf(pass_fourcc)
                    if inp.type == 'VALUE':
                        writebuf(struct.pack('ffff', inp.default_value, inp.default_value,
                                             inp.default_value, inp.default_value))
                    else:
                        writebuf(struct.pack('ffff', inp.default_value[0], inp.default_value[1],
                                             inp.default_value[2], inp.default_value[3]))
<commit_msg>Emit material chunk if socket value differs from default<commit_after>import bpy, bpy.path, os.path, struct

def get_texture_path(image):
    return os.path.normpath(bpy.path.abspath(image.filepath))

SHADER_TYPES = {
    'RetroShader': b'RSHD',
    'RetroDynamicShader': b'RDYN',
    'RetroDynamicAlphaShader': b'RDAL',
    'RetroDynamicCharacterShader': b'RCHR',
}

PASS_TYPE = {
    'Lightmap': b'LMAP',
    'Diffuse': b'DIFF',
    'DiffuseMod': b'DIFM',
    'Emissive': b'EMIS',
    'Specular': b'SPEC',
    'ExtendedSpecular': b'ESPC',
    'Reflection': b'REFL',
    'IndirectTex': b'INDR',
    'Alpha': b'ALPH',
    'AlphaMod': b'ALPM'
}

def write_chunks(writebuf, mat_obj, mesh_obj):

    if not mat_obj.use_nodes:
        raise RuntimeError("HMDL *requires* that shader nodes are used; '{0}' does not".format(mat_obj.name))

    if 'Output' not in mat_obj.node_tree.nodes or mat_obj.node_tree.nodes['Output'].type != 'GROUP':
        raise RuntimeError("HMDL *requires* that an group shader node named 'Output' is present")

    # Root (output) node
    output_node = mat_obj.node_tree.nodes['Output']
    if output_node.node_tree.name not in SHADER_TYPES:
        raise RuntimeError("HMDL *requires* one of the RetroShader group nodes for the 'Output' node")
    writebuf(SHADER_TYPES[output_node.node_tree.name])

    # Count sockets
    chunk_count = 0
    for inp, def_inp in zip(output_node.inputs, output_node.node_tree.inputs):
        if inp.name in PASS_TYPE:
            if inp.is_linked:
                chunk_count += 1
            else:
                # Color pass
                color_set = False
                if inp.type == 'VALUE':
                    color_set = inp.default_value != def_inp.default_value
                else:
                    for comp, def_comp in zip(inp.default_value, def_inp.default_value):
                        color_set |= comp != def_comp
                if color_set:
                    chunk_count += 1

    writebuf(struct.pack('I', chunk_count))

    # Enumerate sockets
    for inp, def_inp in zip(output_node.inputs, output_node.node_tree.inputs):
        if inp.name in PASS_TYPE:
            pass_fourcc = PASS_TYPE[inp.name]
            if inp.is_linked:
                socket = inp.links[0].from_socket
                node = socket.node
                if node.type != 'TEX_IMAGE':
                    raise RuntimeError("HMDL requires all group node inputs connect to Image Texture nodes")

                if not node.image:
                    raise RuntimeError("HMDL texture nodes must specify an image object")

                if not node.inputs['Vector'].is_linked:
                    raise RuntimeError("HMDL texture nodes must have a 'Texture Coordinate', 'UV Map' or 'Group' UV modifier node linked")

                # Determine matrix generator type
                tex_coord_source = 0xff
                uv_anim_type = 0xff
                uv_anim_args = []
                soc_from = node.inputs['Vector'].links[0].from_socket

                if soc_from.node.type == 'GROUP':
                    if soc_from.node.node_tree.name.startswith('RetroUVMode'):
                        uv_anim_type = int(soc_from.node.node_tree.name[11:12])
                        if len(soc_from.node.inputs)-1:
                            for s in range(len(soc_from.node.inputs)-1):
                                soc = soc_from.node.inputs[s+1]
                                if len(soc.links):
                                    raise RuntimeError("UV Modifier nodes may not have parameter links (default values only)")
                                if soc.type == 'VALUE':
                                    uv_anim_args.append(soc.default_value)
                                else:
                                    uv_anim_args.append(soc.default_value[0])
                                    uv_anim_args.append(soc.default_value[1])
                    soc_from = soc_from.node.inputs[0].links[0].from_socket

                elif soc_from.node.type == 'UVMAP' or soc_from.node.type == 'TEX_COORD':
                    pass

                else:
                    raise RuntimeError("HMDL texture nodes must have a 'Texture Coordinate', 'UV Map' or 'Group' UV modifier node linked")

                if soc_from.node.type != 'UVMAP' and soc_from.node.type != 'TEX_COORD':
                    raise RuntimeError("Matrix animator nodes must connect to 'Texture Coordinate' or 'UV Map' node")

                # Resolve map and matrix index
                node_label = soc_from.node.label
                matrix_idx = None
                if node_label.startswith('MTX_'):
                    matrix_idx = int(node_label[4:])

                if soc_from.name == 'UV':
                    if hasattr(soc_from.node, 'uv_map'):
                        uv_name = soc_from.node.uv_map
                        uv_idx = mesh_obj.data.uv_layers.find(uv_name)
                        if uv_idx == -1:
                            raise RuntimeError('UV Layer "%s" doesn\'t exist' % uv_name)
                        tex_coord_source = uv_idx + 2
                    else:
                        tex_coord_source = 2

                elif soc_from.name == 'Normal':
                    tex_coord_source = 1

                elif soc_from.name == 'Window':
                    tex_coord_source = 0

                else:
                    raise RuntimeError("Only the 'UV', 'Normal' and 'Window' sockets may be used from 'Texture Coordinate' nodes")

                alpha = False
                if socket.name == 'Alpha':
                    alpha = True

                writebuf(b'PASS')
                writebuf(pass_fourcc)
                path = get_texture_path(node.image)
                writebuf(struct.pack('I', len(path)))
                writebuf(path.encode())
                writebuf(struct.pack('B', tex_coord_source))
                writebuf(struct.pack('B', uv_anim_type))
                writebuf(struct.pack('I', len(uv_anim_args)))
                for arg in uv_anim_args:
                    writebuf(struct.pack('f', arg))
                writebuf(struct.pack('B', alpha))

            else:
                # Color pass
                color_set = False
                if inp.type == 'VALUE':
                    color_set = inp.default_value != def_inp.default_value
                else:
                    for comp, def_comp in zip(inp.default_value, def_inp.default_value):
                        color_set |= comp != def_comp

                if color_set:
                    writebuf(b'CLR ')
                    writebuf(pass_fourcc)
                    if inp.type == 'VALUE':
                        writebuf(struct.pack('ffff', inp.default_value, inp.default_value,
                                             inp.default_value, inp.default_value))
                    else:
                        writebuf(struct.pack('ffff', inp.default_value[0], inp.default_value[1],
                                             inp.default_value[2], inp.default_value[3]))
<|endoftext|>


<commit_before>
from blimpy import Waterfall
import numpy as np
from scipy.optimize import curve_fit
from fluxcal import foldcal

def get_stokes(cross_dat, **kwargs):
    '''Output stokes parameters (I,Q,U,V) for a rawspec
    cross polarization data array'''
    #Load observation data
    #obs = Waterfall(, max_load=150, **kwargs)
    dat = cross_dat

    #Obtain cross polarization data arrays
    XX = dat[:,0,:]
    YY = dat[:,1,:]
    Re_XY = dat[:,2,:]
    Im_XY = dat[:,3,:]

    #Add a dimension to each to match blimpy file format
    XX = np.expand_dims(XX,axis=1)
    YY = np.expand_dims(YY,axis=1)
    Re_XY = np.expand_dims(Re_XY,axis=1)
    Im_XY = np.expand_dims(Im_XY,axis=1)

    #Compute Stokes Parameters
    I = XX+YY
    Q = XX-YY
    U = 2*Re_XY
    V = -2*Im_XY

    #Compute linear polarization
    #L=np.sqrt(np.square(Q)+np.square(U))

    return I,Q,U,V

def convert_to_coarse(data,chan_per_coarse):
    '''
    Converts a data array with length n_chans to an array of length n_coarse_chans
    by averaging over the coarse channels
    '''
    #find number of coarse channels and reshape array
    num_coarse = data.size/chan_per_coarse
    data_shaped = np.array(np.reshape(data,(num_coarse,chan_per_coarse)))

    #Return the average over each coarse channel
    return np.mean(data_shaped[:,2:-1],axis=1)

def phase_offsets(Udat,Vdat,tsamp,chan_per_coarse,fit=False,**kwargs):
    '''
    Calculates phase difference between X and Y feeds given U and V
    data from a noise diode measurement on the target
    '''
    #Fold noise diode data and calculate ON OFF diferences for U and V
    U_OFF,U_ON = foldcal(Udat,tsamp,**kwargs)
    V_OFF,V_ON = foldcal(Vdat,tsamp,**kwargs)
    Udiff = U_ON-U_OFF
    Vdiff = V_ON-V_OFF

    '''The following is used if using a fit to calculate the phase offsets.'''
    #Fit U data with a cosine and V data with a sine
    #Ufit,Upars = fit_sinusoid(freqs,Udiff,kind='cos')
    #Vfit,Vpars = fit_sinusoid(freqs,Vdiff,kind='sin')

    #Calculate phase offsets for each frequency of the best fits
    #U_phase = Upars[1]*(freqs-Upars[2])
    #V_phase = Vpars[1]*(freqs-Vpars[2])

    #If the fit calculated a negative sine or cosine curve, add a phase of pi
    #if Upars[0]<0:
    #    U_phase += np.pi
    #if Vpars[0]<0:
    #    V_phase += np.pi

    #Return phase offset for each coarse channel
    #if fit==False:
    #    return convert_to_coarse(V_phase,chan_per_coarse)
    #else:
    #    return convert_to_coarse(V_phase,chan_per_coarse),Ufit,Vfit

    # Return the arctan of V/U for each coarse channel
    return convert_to_coarse(np.arctan2(Vdiff,Udiff),chan_per_coarse)

def gain_offsets(Idat,Qdat,tsamp,chan_per_coarse,**kwargs):
    '''
    Determines relative gain error in the X and Y feeds for an
    observation given I and Q noise diode data.
    '''
    #Fold noise diode data and calculate ON OFF differences for I and Q
    I_OFF,I_ON = foldcal(Idat,tsamp,**kwargs)
    Q_OFF,Q_ON = foldcal(Qdat,tsamp,**kwargs)

    #Calculate power in each feed for noise diode ON and OFF
    XX_ON = (I_ON+Q_ON)/2
    XX_OFF = (I_OFF+Q_OFF)/2
    YY_ON = (I_ON-Q_ON)/2
    YY_OFF = (I_OFF-Q_OFF)/2

    #Calculate gain offset (divided by 2) as defined in Heiles (2001)
    G = (XX_OFF-YY_OFF)/(XX_OFF+YY_OFF)

    return convert_to_coarse(G,chan_per_coarse)
def apply_Mueller(I,Q,U,V, gain_offsets, phase_offsets, chan_per_coarse):
    '''
    Returns calibrated Stokes parameters for an observation given an array
    of differential gains and phase differences. Use 'hyptrig' to use the Mueller
    matrix in Liao et al. (2018)
    '''

    #Find shape of data arrays and calculate number of coarse channels
    shape = I.shape
    ax0 = I.shape[0]
    ax1 = I.shape[1]
    nchans = I.shape[2]
    ncoarse = nchans/chan_per_coarse

    #Reshape data arrays to separate coarse channels
    I = np.reshape(I,(ax0,ax1,ncoarse,chan_per_coarse))
    Q = np.reshape(Q,(ax0,ax1,ncoarse,chan_per_coarse))
    U = np.reshape(U,(ax0,ax1,ncoarse,chan_per_coarse))
    V = np.reshape(V,(ax0,ax1,ncoarse,chan_per_coarse))

    #Swap axes 2 and 3 to in order for broadcasting to work correctly
    I = np.swapaxes(I,2,3)
    Q = np.swapaxes(Q,2,3)
    U = np.swapaxes(U,2,3)
    V = np.swapaxes(V,2,3)

    #Apply top left corner of electronics chain inverse Mueller matrix
    a = 1/(1-gain_offsets**2)
    Icorr = a*(I-gain_offsets*Q)
    Qcorr = a*(-1*gain_offsets*I+Q)

    #Apply bottom right corner of electronics chain inverse Mueller matrix
    Ucorr = U*np.cos(phase_offsets)+V*np.sin(phase_offsets)
    Vcorr = -1*U*np.sin(phase_offsets)+V*np.cos(phase_offsets)

    I = Icorr
    Q = Qcorr
    U = Ucorr
    V = Vcorr

    #Reshape arrays to original shape
    I = np.reshape(np.swapaxes(I,2,3),shape)
    Q = np.reshape(np.swapaxes(Q,2,3),shape)
    U = np.reshape(np.swapaxes(U,2,3),shape)
    V = np.reshape(np.swapaxes(V,2,3),shape)

    #Return corrected data arrays
    return I,Q,U,V

def calibrate_pols(cross_pols,diode_cross,obsI=None,onefile=True,**kwargs):
    '''
    Write four calibrated Stokes filterbank files for a given observation
    with a calibrator noise diode measurement on the source
    '''
    #Obtain time sample length, frequencies, and noise diode data
    obs = Waterfall(diode_cross,max_load=150)
    cross_dat = obs.data
    tsamp = obs.header['tsamp']

    #Calculate number of coarse channels in the noise diode measurement (usually 8)
    dio_ncoarse = obs.calc_n_coarse_chan()
    dio_nchans = obs.header['nchans']
    dio_chan_per_coarse = dio_nchans/dio_ncoarse
    obs = None
    Idat,Qdat,Udat,Vdat = get_stokes(cross_dat)
    cross_dat = None
    #Calculate differential gain and phase from noise diode measurements
    gams = gain_offsets(Idat,Qdat,tsamp,dio_chan_per_coarse,**kwargs)
    psis = phase_offsets(Udat,Vdat,tsamp,dio_chan_per_coarse,**kwargs)

    #Clear data arrays for space
    Idat = None
    Qdat = None
    Udat = None
    Vdat = None

    #Get corrected Stokes parameters
    cross_obs = Waterfall(cross_pols,max_load=150)
    obs_ncoarse = cross_obs.calc_n_coarse_chan()
    obs_nchans = cross_obs.header['nchans']
    obs_chan_per_coarse = obs_nchans/obs_ncoarse
    cross_dat = cross_obs.data
    cross_obs = None
    I,Q,U,V = get_stokes(cross_dat)

    I,Q,U,V = apply_Mueller(I,Q,U,V,gams,psis,obs_chan_per_coarse)

    #Use onefile (default) to produce one filterbank file containing all Stokes information
    if onefile==True:
        cross_obs.data[:,0,:] = np.squeeze(I)
        cross_obs.data[:,1,:] = np.squeeze(Q)
        cross_obs.data[:,2,:] = np.squeeze(U)
        cross_obs.data[:,3,:] = np.squeeze(V)
        cross_obs.write_to_fil(cross_pols[:-15]+'.SIQUV.polcal.fil')
        print 'Calibrated Stokes parameters written to '+cross_pols[:-15]+'.SIQUV.polcal.fil'
        return

    #Write corrected Stokes parameters to four filterbank files if onefile==False
    obs = Waterfall(obs_I,max_load=150)
    obs.data = I
    obs.write_to_fil(cross_pols[:-15]+'.SI.polcal.fil')   #assuming file is named *.cross_pols.fil
    print 'Calibrated Stokes I written to '+cross_pols[:-15]+'.SI.polcal.fil'

    obs.data = Q
    obs.write_to_fil(cross_pols[:-15]+'.Q.polcal.fil')   #assuming file is named *.cross_pols.fil
    print 'Calibrated Stokes Q written to '+cross_pols[:-15]+'.Q.polcal.fil'

    obs.data = U
    obs.write_to_fil(cross_pols[:-15]+'.U.polcal.fil')   #assuming file is named *.cross_pols.fil
    print 'Calibrated Stokes U written to '+cross_pols[:-15]+'.U.polcal.fil'

    obs.data = V
    obs.write_to_fil(cross_pols[:-15]+'.V.polcal.fil')   #assuming file is named *.cross_pols.fil
    print 'Calibrated Stokes V written to '+cross_pols[:-15]+'.V.polcal.fil'

def fracpols(str, **kwargs):
    '''Output fractional linear and circular polarizations for a
    rawspec cross polarization .fil file'''

    I,Q,U,V,L=get_stokes(str, **kwargs)
    return L/I,V/I

def write_stokefils(str, str_I, Ifil=False, Qfil=False, Ufil=False, Vfil=False, Lfil=False, **kwargs):
    '''Writes up to 5 new filterbank files corresponding to each Stokes
    parameter (and total linear polarization L) for a given cross polarization .fil file'''

    I,Q,U,V,L=get_stokes(str, **kwargs)
    obs = Waterfall(str_I, max_load=150) #Load filterbank file to write stokes data to
    if Ifil:
        obs.data = I
        obs.write_to_fil(str[:-15]+'.I.fil')   #assuming file is named *.cross_pols.fil

    if Qfil:
        obs.data = Q
        obs.write_to_fil(str[:-15]+'.Q.fil')   #assuming file is named *.cross_pols.fil

    if Ufil:
        obs.data = U
        obs.write_to_fil(str[:-15]+'.U.fil')   #assuming file is named *.cross_pols.fil

    if Vfil:
        obs.data = V
        obs.write_to_fil(str[:-15]+'.V.fil')   #assuming file is named *.cross_pols.fil

    if Lfil:
        obs.data = L
        obs.write_to_fil(str[:-15]+'.L.fil')   #assuming file is named *.cross_pols.fil

def write_polfils(str, str_I, **kwargs):
    '''Writes two new filterbank files containing fractional linear and
    circular polarization data'''

    lin,circ=fracpols(str, **kwargs)
    obs = Waterfall(str_I, max_load=150)

    obs.data = lin
    obs.write_to_fil(str[:-15]+'.linpol.fil')   #assuming file is named *.cross_pols.fil

    obs.data = circ
    obs.write_to_fil(str[:-15]+'.circpol.fil')   #assuming file is named *.cross_pols.fil

def fit_sinusoid(freqs,ary,kind,chan_per_core=8):
    '''
    Calculates a sinusoidal (sin or cos) fit given a range of frequencies and data
    '''

    freqs_clean=np.array(freqs)
    ary_clean=np.array(ary)

    freqs_clean[::chan_per_core]=np.NaN
    freqs_clean[1::chan_per_core]=np.NaN
    freqs_clean[(chan_per_core-1)::chan_per_core]=np.NaN
    ary_clean[::chan_per_core]=np.NaN
    ary_clean[1::chan_per_core]=np.NaN
    ary_clean[(chan_per_core-1)::chan_per_core]=np.NaN

    freqs_clean = [x for x in freqs_clean if str(x) != 'nan']
    ary_clean = [x for x in ary_clean if str(x) != 'nan']

    f_est = np.pi/75
    if kind=='sin':
        popt,pcov = curve_fit(sine,freqs_clean,ary_clean,p0=[10.0,f_est,0.0,0.0])
        return sine(freqs,*popt),popt
    else:
        popt,pcov = curve_fit(cosine,freqs_clean,ary_clean,p0=[10.0,f_est,0.0,0.0])
        return cosine(freqs,*popt),popt

def sine(x,a,b,h,c):
    return(a*np.sin(b*(x-h))+c)

def cosine(x,a,b,h,c):
    return(a*np.cos(b*(x-h))+c)

#end module

<commit_msg>Memory allocation adjustments<commit_after>
from blimpy import Waterfall
import numpy as np
from scipy.optimize import curve_fit
from fluxcal import foldcal

def get_stokes(cross_dat, **kwargs):
    '''Output stokes parameters (I,Q,U,V) for a rawspec
    cross polarization data array'''
    #Load observation data
    #obs = Waterfall(, max_load=150, **kwargs)
    dat = cross_dat

    #Obtain cross polarization data arrays
    XX = dat[:,0,:]
    YY = dat[:,1,:]
    Re_XY = dat[:,2,:]
    Im_XY = dat[:,3,:]

    #Add a dimension to each to match blimpy file format
    XX = np.expand_dims(XX,axis=1)
    YY = np.expand_dims(YY,axis=1)
    Re_XY = np.expand_dims(Re_XY,axis=1)
    Im_XY = np.expand_dims(Im_XY,axis=1)

    #Compute Stokes Parameters
    I = XX+YY
    Q = XX-YY
    U = 2*Re_XY
    V = -2*Im_XY

    #Compute linear polarization
    #L=np.sqrt(np.square(Q)+np.square(U))

    return I,Q,U,V

def convert_to_coarse(data,chan_per_coarse):
    '''
    Converts a data array with length n_chans to an array of length n_coarse_chans
    by averaging over the coarse channels
    '''
    #find number of coarse channels and reshape array
    num_coarse = data.size/chan_per_coarse
    data_shaped = np.array(np.reshape(data,(num_coarse,chan_per_coarse)))

    #Return the average over each coarse channel
    return np.mean(data_shaped[:,2:-1],axis=1)

def phase_offsets(Udat,Vdat,tsamp,chan_per_coarse,fit=False,**kwargs):
    '''
    Calculates phase difference between X and Y feeds given U and V
    data from a noise diode measurement on the target
    '''
    #Fold noise diode data and calculate ON OFF diferences for U and V
    U_OFF,U_ON = foldcal(Udat,tsamp,**kwargs)
    V_OFF,V_ON = foldcal(Vdat,tsamp,**kwargs)
    Udiff = U_ON-U_OFF
    Vdiff = V_ON-V_OFF

    '''The following is used if using a fit to calculate the phase offsets.'''
    #Fit U data with a cosine and V data with a sine
    #Ufit,Upars = fit_sinusoid(freqs,Udiff,kind='cos')
    #Vfit,Vpars = fit_sinusoid(freqs,Vdiff,kind='sin')

    #Calculate phase offsets for each frequency of the best fits
    #U_phase = Upars[1]*(freqs-Upars[2])
    #V_phase = Vpars[1]*(freqs-Vpars[2])

    #If the fit calculated a negative sine or cosine curve, add a phase of pi
    #if Upars[0]<0:
    #    U_phase += np.pi
    #if Vpars[0]<0:
    #    V_phase += np.pi

    #Return phase offset for each coarse channel
    #if fit==False:
    #    return convert_to_coarse(V_phase,chan_per_coarse)
    #else:
    #    return convert_to_coarse(V_phase,chan_per_coarse),Ufit,Vfit

    # Return the arctan of V/U for each coarse channel
    return convert_to_coarse(np.arctan2(Vdiff,Udiff),chan_per_coarse)

def gain_offsets(Idat,Qdat,tsamp,chan_per_coarse,**kwargs):
    '''
    Determines relative gain error in the X and Y feeds for an
    observation given I and Q noise diode data.
    '''
    #Fold noise diode data and calculate ON OFF differences for I and Q
    I_OFF,I_ON = foldcal(Idat,tsamp,**kwargs)
    Q_OFF,Q_ON = foldcal(Qdat,tsamp,**kwargs)

    #Calculate power in each feed for noise diode ON and OFF
    XX_ON = (I_ON+Q_ON)/2
    XX_OFF = (I_OFF+Q_OFF)/2
    YY_ON = (I_ON-Q_ON)/2
    YY_OFF = (I_OFF-Q_OFF)/2

    #Calculate gain offset (divided by 2) as defined in Heiles (2001)
    G = (XX_OFF-YY_OFF)/(XX_OFF+YY_OFF)

    return convert_to_coarse(G,chan_per_coarse)
def apply_Mueller(I,Q,U,V, gain_offsets, phase_offsets, chan_per_coarse):
    '''
    Returns calibrated Stokes parameters for an observation given an array
    of differential gains and phase differences. Use 'hyptrig' to use the Mueller
    matrix in Liao et al. (2018)
    '''

    #Find shape of data arrays and calculate number of coarse channels
    shape = I.shape
    ax0 = I.shape[0]
    ax1 = I.shape[1]
    nchans = I.shape[2]
    ncoarse = nchans/chan_per_coarse

    #Reshape data arrays to separate coarse channels
    I = np.reshape(I,(ax0,ax1,ncoarse,chan_per_coarse))
    Q = np.reshape(Q,(ax0,ax1,ncoarse,chan_per_coarse))
    U = np.reshape(U,(ax0,ax1,ncoarse,chan_per_coarse))
    V = np.reshape(V,(ax0,ax1,ncoarse,chan_per_coarse))

    #Swap axes 2 and 3 to in order for broadcasting to work correctly
    I = np.swapaxes(I,2,3)
    Q = np.swapaxes(Q,2,3)
    U = np.swapaxes(U,2,3)
    V = np.swapaxes(V,2,3)

    #Apply top left corner of electronics chain inverse Mueller matrix
    a = 1/(1-gain_offsets**2)
    Icorr = a*(I-gain_offsets*Q)
    Qcorr = a*(-1*gain_offsets*I+Q)

    #Clear uncalibrated I and Q
    I = None
    Q = None

    #Apply bottom right corner of electronics chain inverse Mueller matrix
    Ucorr = U*np.cos(phase_offsets)+V*np.sin(phase_offsets)
    Vcorr = -1*U*np.sin(phase_offsets)+V*np.cos(phase_offsets)

    #Clear uncalibrated U and V
    U = None
    V = None

    #Reshape arrays to original shape
    Icorr = np.reshape(np.swapaxes(Icorr,2,3),shape)
    Qcorr = np.reshape(np.swapaxes(Qcorr,2,3),shape)
    Ucorr = np.reshape(np.swapaxes(Ucorr,2,3),shape)
    Vcorr = np.reshape(np.swapaxes(Vcorr,2,3),shape)

    #Return corrected data arrays
    return Icorr,Qcorr,Ucorr,Vcorr

def calibrate_pols(cross_pols,diode_cross,obsI=None,onefile=True,**kwargs):
    '''
    Write four calibrated Stokes filterbank files for a given observation
    with a calibrator noise diode measurement on the source
    '''
    #Obtain time sample length, frequencies, and noise diode data
    obs = Waterfall(diode_cross,max_load=150)
    cross_dat = obs.data
    tsamp = obs.header['tsamp']

    #Calculate number of coarse channels in the noise diode measurement (usually 8)
    dio_ncoarse = obs.calc_n_coarse_chan()
    dio_nchans = obs.header['nchans']
    dio_chan_per_coarse = dio_nchans/dio_ncoarse
    obs = None
    Idat,Qdat,Udat,Vdat = get_stokes(cross_dat)
    cross_dat = None
    #Calculate differential gain and phase from noise diode measurements
    gams = gain_offsets(Idat,Qdat,tsamp,dio_chan_per_coarse,**kwargs)
    psis = phase_offsets(Udat,Vdat,tsamp,dio_chan_per_coarse,**kwargs)

    #Clear data arrays for space
    Idat = None
    Qdat = None
    Udat = None
    Vdat = None

    #Get corrected Stokes parameters
    cross_obs = Waterfall(cross_pols,max_load=150)
    obs_ncoarse = cross_obs.calc_n_coarse_chan()
    obs_nchans = cross_obs.header['nchans']
    obs_chan_per_coarse = obs_nchans/obs_ncoarse
    cross_dat = cross_obs.data
    cross_obs = None
    I,Q,U,V = get_stokes(cross_dat)

    I,Q,U,V = apply_Mueller(I,Q,U,V,gams,psis,obs_chan_per_coarse)

    #Use onefile (default) to produce one filterbank file containing all Stokes information
    if onefile==True:
        cross_obs.data[:,0,:] = np.squeeze(I)
        cross_obs.data[:,1,:] = np.squeeze(Q)
        cross_obs.data[:,2,:] = np.squeeze(U)
        cross_obs.data[:,3,:] = np.squeeze(V)
        cross_obs.write_to_fil(cross_pols[:-15]+'.SIQUV.polcal.fil')
        print 'Calibrated Stokes parameters written to '+cross_pols[:-15]+'.SIQUV.polcal.fil'
        return

    #Write corrected Stokes parameters to four filterbank files if onefile==False
    obs = Waterfall(obs_I,max_load=150)
    obs.data = I
    obs.write_to_fil(cross_pols[:-15]+'.SI.polcal.fil')   #assuming file is named *.cross_pols.fil
    print 'Calibrated Stokes I written to '+cross_pols[:-15]+'.SI.polcal.fil'

    obs.data = Q
    obs.write_to_fil(cross_pols[:-15]+'.Q.polcal.fil')   #assuming file is named *.cross_pols.fil
    print 'Calibrated Stokes Q written to '+cross_pols[:-15]+'.Q.polcal.fil'

    obs.data = U
    obs.write_to_fil(cross_pols[:-15]+'.U.polcal.fil')   #assuming file is named *.cross_pols.fil
    print 'Calibrated Stokes U written to '+cross_pols[:-15]+'.U.polcal.fil'

    obs.data = V
    obs.write_to_fil(cross_pols[:-15]+'.V.polcal.fil')   #assuming file is named *.cross_pols.fil
    print 'Calibrated Stokes V written to '+cross_pols[:-15]+'.V.polcal.fil'

def fracpols(str, **kwargs):
    '''Output fractional linear and circular polarizations for a
    rawspec cross polarization .fil file'''

    I,Q,U,V,L=get_stokes(str, **kwargs)
    return L/I,V/I

def write_stokefils(str, str_I, Ifil=False, Qfil=False, Ufil=False, Vfil=False, Lfil=False, **kwargs):
    '''Writes up to 5 new filterbank files corresponding to each Stokes
    parameter (and total linear polarization L) for a given cross polarization .fil file'''

    I,Q,U,V,L=get_stokes(str, **kwargs)
    obs = Waterfall(str_I, max_load=150) #Load filterbank file to write stokes data to
    if Ifil:
        obs.data = I
        obs.write_to_fil(str[:-15]+'.I.fil')   #assuming file is named *.cross_pols.fil

    if Qfil:
        obs.data = Q
        obs.write_to_fil(str[:-15]+'.Q.fil')   #assuming file is named *.cross_pols.fil

    if Ufil:
        obs.data = U
        obs.write_to_fil(str[:-15]+'.U.fil')   #assuming file is named *.cross_pols.fil

    if Vfil:
        obs.data = V
        obs.write_to_fil(str[:-15]+'.V.fil')   #assuming file is named *.cross_pols.fil

    if Lfil:
        obs.data = L
        obs.write_to_fil(str[:-15]+'.L.fil')   #assuming file is named *.cross_pols.fil

def write_polfils(str, str_I, **kwargs):
    '''Writes two new filterbank files containing fractional linear and
    circular polarization data'''

    lin,circ=fracpols(str, **kwargs)
    obs = Waterfall(str_I, max_load=150)

    obs.data = lin
    obs.write_to_fil(str[:-15]+'.linpol.fil')   #assuming file is named *.cross_pols.fil

    obs.data = circ
    obs.write_to_fil(str[:-15]+'.circpol.fil')   #assuming file is named *.cross_pols.fil

def fit_sinusoid(freqs,ary,kind,chan_per_core=8):
    '''
    Calculates a sinusoidal (sin or cos) fit given a range of frequencies and data
    '''

    freqs_clean=np.array(freqs)
    ary_clean=np.array(ary)

    freqs_clean[::chan_per_core]=np.NaN
    freqs_clean[1::chan_per_core]=np.NaN
    freqs_clean[(chan_per_core-1)::chan_per_core]=np.NaN
    ary_clean[::chan_per_core]=np.NaN
    ary_clean[1::chan_per_core]=np.NaN
    ary_clean[(chan_per_core-1)::chan_per_core]=np.NaN

    freqs_clean = [x for x in freqs_clean if str(x) != 'nan']
    ary_clean = [x for x in ary_clean if str(x) != 'nan']

    f_est = np.pi/75
    if kind=='sin':
        popt,pcov = curve_fit(sine,freqs_clean,ary_clean,p0=[10.0,f_est,0.0,0.0])
        return sine(freqs,*popt),popt
    else:
        popt,pcov = curve_fit(cosine,freqs_clean,ary_clean,p0=[10.0,f_est,0.0,0.0])
        return cosine(freqs,*popt),popt

def sine(x,a,b,h,c):
    return(a*np.sin(b*(x-h))+c)

def cosine(x,a,b,h,c):
    return(a*np.cos(b*(x-h))+c)

#end module

<|endoftext|>


<commit_before>import numpy as np
import matplotlib.pyplot as plt
from math import sqrt, cos, sin, pi
from time import clock
from copy import deepcopy

def getVideoFrame(VIDEO_PATH, xdim=800, ydim=600, frame=60):
    FFMPEG_BIN = r"C:\Python33\Lib\ffmpeg-20140713-git-42c1cc3-win64-static\bin\ffmpeg.exe"
    #r"C:\Users\El'endia Starman\Desktop\My works\Jaw Tracking\Demo vids 1\demovid1_left0001-0075.mp4"

    ### get video info
    import subprocess as sp
    command = [FFMPEG_BIN, '-i', VIDEO_PATH, '-']
    pipe = sp.Popen(command, stdout=sp.PIPE, stderr=sp.PIPE)
    pipe.stdout.readline()
    pipe.terminate()
    infos = str(pipe.stderr.read())[2:-1]
##    print(infos)

    ### get the frames per second
    end = infos.index("fps")-1
    beg = end-1
    while infos[beg-1] != " ":
        beg -= 1
    fps = float(infos[beg:end])

    ### pick the desired frame
##    frame = frame
    Hz = fps#round(fps)

    ### calculate timecode from frame and fps
    s = (frame // Hz) % 60
    m = (frame // Hz) // 60 % 60
    h = (frame // Hz) // 3600
    f = (frame %  Hz) * Hz
    timecode = "%02d:%02d:%02d.%03d" % (h,m,s,f)
##    print(timecode)

    ### set up pipe to get single video frame
    command = [ FFMPEG_BIN,
                '-ss', timecode,
                '-i', VIDEO_PATH,
                '-f', 'image2pipe',
                '-pix_fmt', 'rgb24',
                '-vcodec', 'rawvideo', '-']
    pipe = sp.Popen(command, stdout = sp.PIPE, bufsize=10**8)

    ### get and format single video frame
    import numpy as np

    ### dimensions of image
##    xdim = 800
##    ydim = 600

    ### get frame image
    raw_image = pipe.stdout.read(xdim*ydim*3)
    image = np.fromstring(raw_image, dtype='uint8')
    image = image.reshape((ydim,xdim,3))
    pipe.stdout.flush()
    pipe.terminate()

    ### display image
    #fig2 = plt.figure(2)
    #plt.imshow(image)

    return image, xdim,ydim, frame

def colorDistance(target,actual): return sqrt(sum([(a-b)**2 for a,b in zip(target,actual)]))
def colorDistanceMagSquared(target,actual): return sum([(a-b)**2 for a,b in zip(target,actual)])
def colorDistanceSum(target,actual): return sum([abs(a-b) for a,b in zip(target,actual)])
def colorIsExact(target,actual): return target == actual

def colDis(targetColor): return lambda col: colorDistance(targetColor,col)
def colDisMagSq(targetColor): return lambda col: colorDistanceMagSquared(targetColor,col)
def colDisSum(targetColor): return lambda col: colorDistanceSum(targetColor,col)
def colIsExact(targetColor): return lambda col: colorIsExact(targetColor,col)

def gray(z): return [np.uint8(z)]*3
def col(r,g,b): return list(map(np.uint8,[r,g,b]))

def circumcircle(p1,p2,p3): #each point is expected to be [x,y]
    #taken from https://en.wikipedia.org/wiki/Circumscribed_circle#Cartesian_coordinates_from_cross-_and_dot-products
    diff1 = [p1[0]-p2[0], p1[1]-p2[1]]
    diff2 = [p1[0]-p3[0], p1[1]-p3[1]]
    diff3 = [p2[0]-p3[0], p2[1]-p3[1]]

    magSquared1 = diff1[0]**2 + diff1[1]**2
    magSquared2 = diff2[0]**2 + diff2[1]**2
    magSquared3 = diff3[0]**2 + diff3[1]**2

    mag1 = sqrt(magSquared1)
    mag2 = sqrt(magSquared2)
    mag3 = sqrt(magSquared3)

    dot1 = diff1[0]*diff2[0] + diff1[1]*diff2[1]
    dot2 = -diff1[0]*diff3[0] + -diff1[1]*diff3[1]
    dot3 = -diff2[0]*-diff3[0] + -diff2[1]*-diff3[1]

    #working in 2D simplifies the cross product and its magnitude a lot
    crossMag = diff1[0]*diff3[1] - diff1[1]*diff3[0]
    crossMagSquared = crossMag**2

    r = mag1*mag2*mag3/(2*crossMag)

    alpha = magSquared3*dot1/(2*crossMagSquared)
    beta  = magSquared2*dot2/(2*crossMagSquared)
    gamma = magSquared1*dot3/(2*crossMagSquared)

    x = alpha*p1[0] + beta*p2[0] + gamma*p3[0]
    y = alpha*p1[1] + beta*p2[1] + gamma*p3[1]

    return [x,y,r]

def squareSpiral(startX,startY, spacing, limit):
    sx = startX
    sy = startY
    yield sx,sy

    level = 0
    dx = -1
    dy = 0

    while level < limit:
        if dx == -1 and dy == level:
            level += 1
            dx = 0
            dy = level
        else:
            if dy == level and dx < level:
                dx += 1
            elif dx == level and dy > -level:
                dy -= 1
            elif dy == -level and dx > -level:
                dx -= 1
            elif dx == -level and dy < level:
                dy += 1

##        print(dx,dy,level)

        yield sx+dx*spacing, sy+dy*spacing

### SEARCH ###

def search_pixels(blobs, xdim, ydim):
    for centerY in range(ydim):
        for centerX in range(xdim):

            goodEnough = (cD(image[centerY][centerX]) <= maxColDis)

            if goodEnough:
                dupe = 0
                for blob in blobs:
                    if (blob[0]-centerX)**2 + (blob[1]-centerY)**2 <= blob[2]**2*1.5:
                        dupe = 1
                        break

                if not dupe:
                    yield centerX, centerY

def search_spacedPixels(blobs, xdim, ydim, spacing=1):
    for y in range(ydim//spacing):
        for x in range(xdim//spacing):

            centerX = x*spacing
            centerY = y*spacing

            goodEnough = (cD(image[centerY][centerX]) <= maxColDis)

            if goodEnough:
                dupe = 0
                for blob in blobs:
                    if (blob[0]-centerX)**2 + (blob[1]-centerY)**2 <= blob[2]**2*1.5:
                        dupe = 1
                        break

                if not dupe:
                    yield centerX, centerY

def search_squareSpiral(blobs, xdim, ydim, spacing=1):
    startX = xdim//2
    startY = ydim//2

    spiral = squareSpiral(startX,startY, spacing, min(xdim,ydim)//(2*spacing)-1)

    while 1:

        try:
            centerX, centerY = next(spiral)
        except StopIteration:
            return

        goodEnough = (cD(image[centerY][centerX]) <= maxColDis)

        if goodEnough:
            dupe = 0
            for blob in blobs:
                if (blob[0]-centerX)**2 + (blob[1]-centerY)**2 <= blob[2]**2*1.5:
                    dupe = 1
                    break

            if not dupe:
                yield centerX, centerY

def search_rectQuadsection(blobs, xdim, ydim, spacing=1):
    quads = [ [[0,0],[xdim-1,ydim-1]] ]

    while len(quads):
        quad = quads.pop(0)
        centerX = (quad[0][0]+quad[1][0])//2
        centerY = (quad[0][1]+quad[1][1])//2

        if cD(image[centerY][centerX]) <= maxColDis:
            dupe = 0
            for blob in blobs:
                if (blob[0]-centerX)**2 + (blob[1]-centerY)**2 <= blob[2]**2*1.5:
                    dupe = 1
                    break

            if not dupe:
                yield centerX,centerY

        if abs(quad[0][0]-quad[1][0]) > 2*spacing and abs(quad[0][1]-quad[1][1]) > 2*spacing:
            quads.append([[centerX,centerY],[quad[0][0],quad[0][1]]])
            quads.append([[centerX,centerY],[quad[0][0],quad[1][1]]])
            quads.append([[centerX,centerY],[quad[1][0],quad[0][1]]])
            quads.append([[centerX,centerY],[quad[1][0],quad[1][1]]])

### SURVEY ###

def survey_floodFill(sX,sY):
    cX,cY,cR = (None,None,None)

    queue = [[sX,sY]]
    q = 0
    while q < len(queue):
        qx,qy = queue[q]

        for dx,dy in [[1,0],[0,1],[-1,0],[0,-1]]:
            x2, y2 = qx+dx, qy+dy

            if cD(image[y2][x2]) <= maxColDis:
                image[y2][x2] = col(255,0,0)
                queue.append([x2,y2])

        q += 1

    sumX,sumY = list(map(sum, zip(*queue)))
    cX, cY = round(sumX/q), round(sumY/q)
    cR = round(sqrt( q/pi ))

    #image[cY][cX] = col(0,0,255)

    return (cX,cY,cR)

def survey_circumcircle(sX,sY, numIterations=3, showVerts=0): #showVerts is also "minRadius"
    cX,cY,cR = (None,None,None)

    for k in range(numIterations): #refines circumcircle estimate by using prior guess
        vertices = []

        for dx,dy in [[1,0],[0,1],[-1,0],[0,-1]]:
            pX = sX+dx
            pY = sY+dy
            while cD(image[pY][pX]) <= maxColDis:
                pX += dx
                pY += dy
            vertices.append( [pX,pY] )

        v = vertices

        #get four possible circle centers/radii
        c1 = circumcircle(v[0],v[1],v[2])
        c2 = circumcircle(v[0],v[1],v[3])
        c3 = circumcircle(v[0],v[2],v[3])
        c4 = circumcircle(v[1],v[2],v[3])

        #average them
        avgX = (c1[0]+c2[0]+c3[0]+c4[0])/4
        avgY = (c1[1]+c2[1]+c3[1]+c4[1])/4
        avgR = (c1[2]+c2[2]+c3[2]+c4[2])/4

        c5 = [round(avgX),round(avgY),round(avgR)]
        cX,cY,cR = c5

        sX = cX
        sY = cY

    if showVerts and cR >= showVerts:
        image[cY][cX] = col(0,0,255)
        for vert in vertices:
            image[vert[1]][vert[0]] = col(255,0,0)

    return (cX,cY,cR)

### The actual start of the program ###

VIDEO_PATH = r".\Demo vids 1\demovid1_left0001-0075.mp4"

targetColor = [0,255,0]
cD = colDisSum(targetColor)
minRadius = 6 #pixels
maxColDis = 40

imgStart = clock()
origImage, xdim,ydim, frame = getVideoFrame(VIDEO_PATH)
imgEnd = clock() - imgStart
print("Time to get frame: %.3f seconds." % imgEnd)

search_algs = [['pixel by pixel', search_pixels],
               ['pixels spaced apart', search_spacedPixels, minRadius],
               ['square spiral', search_squareSpiral, minRadius],
               ['rectangular quadsection', search_rectQuadsection, minRadius],
               ]

survey_algs = [['flood fill', survey_floodFill],
               ['circumcircle', survey_circumcircle],
               ]

for search_alg in search_algs:
    for survey_alg in survey_algs:

        image = deepcopy(origImage)

        print()
        print("Search algorithm: %s" % search_alg[0])
        print("Survey algorithm: %s" % survey_alg[0])

        startTime = clock()
        times = []

        blobs = []
        search = search_alg[1](blobs, xdim,ydim, *search_alg[2:])

        while 1:
            try:
                sx,sy = next(search)
            except StopIteration:
                break

            surveyStart = clock()
            blob = survey_alg[1](sx,sy)
            surveyEnd = clock()
            times.append(surveyEnd - surveyStart)

            if blob[2] >= minRadius: blobs.append(blob)

        endTime = clock()
        print(" Total time: %.3f seconds" % (endTime-startTime))
        print(" Search time: %.3f seconds" % (endTime-startTime - sum(times)))
        print(" Number of surveys: %d" % len(times))
        print(" Number of blobs: %d" % len(blobs))
        print(" Survey time (total): %.3f seconds" % (sum(times)))
        print(" Survey time (average): %.3f seconds" % (sum(times)/len(times)))

##        for i,t in enumerate(times):
##            print("Blob %d was surveyed in %.3f seconds." % (i,t))

##
##print("Color distance code timing:")
##cdStart = clock()
##num = 0
##for dx in range(xdim//minRadius):
##    for dy in range(ydim//minRadius):
##        num += 1
##        eh = (cD(image[dy*minRadius][dx*minRadius]) <= maxColDis)
####        col = image[dy*minRadius][dx*minRadius]
####        eh = (cD(col) == 1)
##
##cdEnd = clock()
##print("Total time: %.3f seconds for %d samples. Average time per sample: %.6f seconds." % (cdEnd-cdStart, num, (cdEnd-cdStart)/num))

### Note on different color distance methods:
### The difference between Euclidean distance and exact match is on the order
### of 0.02 - 0.03 seconds difference in the total time, which is about 10% of
### the total time. Not that much. I'll go with accuracy over speed here.

#plt.imshow(image)
#plt.show()
<commit_msg>Cleaned up the code some.<commit_after>import numpy as np
import matplotlib.pyplot as plt
from math import sqrt, cos, sin, pi
from time import clock
from copy import deepcopy

def getVideoFrame(VIDEO_PATH, xdim=800, ydim=600, frame=60):
    FFMPEG_BIN = r"C:\Python33\Lib\ffmpeg-20140713-git-42c1cc3-win64-static\bin\ffmpeg.exe"

    ### get video info
    import subprocess as sp
    command = [FFMPEG_BIN, '-i', VIDEO_PATH, '-']
    pipe = sp.Popen(command, stdout=sp.PIPE, stderr=sp.PIPE)
    pipe.stdout.readline()
    pipe.terminate()
    infos = str(pipe.stderr.read())[2:-1]

    ### get the frames per second
    end = infos.index("fps")-1
    beg = end-1
    while infos[beg-1] != " ":
        beg -= 1
    fps = float(infos[beg:end])

    ### pick the desired frame
    Hz = fps

    ### calculate timecode from frame and fps
    s = (frame // Hz) % 60
    m = (frame // Hz) // 60 % 60
    h = (frame // Hz) // 3600
    f = (frame %  Hz) * Hz
    timecode = "%02d:%02d:%02d.%03d" % (h,m,s,f)

    ### set up pipe to get single video frame
    command = [ FFMPEG_BIN,
                '-ss', timecode,
                '-i', VIDEO_PATH,
                '-f', 'image2pipe',
                '-pix_fmt', 'rgb24',
                '-vcodec', 'rawvideo', '-']
    pipe = sp.Popen(command, stdout = sp.PIPE, bufsize=10**8)

    ### get and format single video frame
    import numpy as np

    ### dimensions of image
##    xdim = 800
##    ydim = 600

    ### get frame image
    raw_image = pipe.stdout.read(xdim*ydim*3)
    image = np.fromstring(raw_image, dtype='uint8')
    image = image.reshape((ydim,xdim,3))
    pipe.stdout.flush()
    pipe.terminate()

    return image, xdim,ydim, frame

def colorDistance(target,actual): return sqrt(sum([(a-b)**2 for a,b in zip(target,actual)]))
def colorDistanceMagSquared(target,actual): return sum([(a-b)**2 for a,b in zip(target,actual)])
def colorDistanceSum(target,actual): return sum([abs(a-b) for a,b in zip(target,actual)])
def colorIsExact(target,actual): return target == actual

def colDis(targetColor): return lambda col: colorDistance(targetColor,col)
def colDisMagSq(targetColor): return lambda col: colorDistanceMagSquared(targetColor,col)
def colDisSum(targetColor): return lambda col: colorDistanceSum(targetColor,col)
def colIsExact(targetColor): return lambda col: colorIsExact(targetColor,col)

def gray(z): return [np.uint8(z)]*3
def col(r,g,b): return list(map(np.uint8,[r,g,b]))

def circumcircle(p1,p2,p3): #each point is expected to be [x,y]
    #taken from https://en.wikipedia.org/wiki/Circumscribed_circle#Cartesian_coordinates_from_cross-_and_dot-products
    diff1 = [p1[0]-p2[0], p1[1]-p2[1]]
    diff2 = [p1[0]-p3[0], p1[1]-p3[1]]
    diff3 = [p2[0]-p3[0], p2[1]-p3[1]]

    magSquared1 = diff1[0]**2 + diff1[1]**2
    magSquared2 = diff2[0]**2 + diff2[1]**2
    magSquared3 = diff3[0]**2 + diff3[1]**2

    mag1 = sqrt(magSquared1)
    mag2 = sqrt(magSquared2)
    mag3 = sqrt(magSquared3)

    dot1 = diff1[0]*diff2[0] + diff1[1]*diff2[1]
    dot2 = -diff1[0]*diff3[0] + -diff1[1]*diff3[1]
    dot3 = -diff2[0]*-diff3[0] + -diff2[1]*-diff3[1]

    #working in 2D simplifies the cross product and its magnitude a lot
    crossMag = diff1[0]*diff3[1] - diff1[1]*diff3[0]
    crossMagSquared = crossMag**2

    r = mag1*mag2*mag3/(2*crossMag)

    alpha = magSquared3*dot1/(2*crossMagSquared)
    beta  = magSquared2*dot2/(2*crossMagSquared)
    gamma = magSquared1*dot3/(2*crossMagSquared)

    x = alpha*p1[0] + beta*p2[0] + gamma*p3[0]
    y = alpha*p1[1] + beta*p2[1] + gamma*p3[1]

    return [x,y,r]

def squareSpiral(startX,startY, spacing, limit):
    sx = startX
    sy = startY
    yield sx,sy

    level = 0
    dx = -1
    dy = 0

    while level < limit:
        if dx == -1 and dy == level:
            level += 1
            dx = 0
            dy = level
        else:
            if dy == level and dx < level:
                dx += 1
            elif dx == level and dy > -level:
                dy -= 1
            elif dy == -level and dx > -level:
                dx -= 1
            elif dx == -level and dy < level:
                dy += 1

        yield sx+dx*spacing, sy+dy*spacing

### SEARCH ###

def search_pixels(blobs, xdim, ydim):
    for centerY in range(ydim):
        for centerX in range(xdim):

            goodEnough = (cD(image[centerY][centerX]) <= maxColDis)

            if goodEnough:
                dupe = 0
                for blob in blobs:
                    if (blob[0]-centerX)**2 + (blob[1]-centerY)**2 <= blob[2]**2*1.5:
                        dupe = 1
                        break

                if not dupe:
                    yield centerX, centerY

def search_spacedPixels(blobs, xdim, ydim, spacing=1):
    for y in range(ydim//spacing):
        for x in range(xdim//spacing):

            centerX = x*spacing
            centerY = y*spacing

            goodEnough = (cD(image[centerY][centerX]) <= maxColDis)

            if goodEnough:
                dupe = 0
                for blob in blobs:
                    if (blob[0]-centerX)**2 + (blob[1]-centerY)**2 <= blob[2]**2*1.5:
                        dupe = 1
                        break

                if not dupe:
                    yield centerX, centerY

def search_squareSpiral(blobs, xdim, ydim, spacing=1):
    startX = xdim//2
    startY = ydim//2

    spiral = squareSpiral(startX,startY, spacing, min(xdim,ydim)//(2*spacing)-1)

    while 1:

        try:
            centerX, centerY = next(spiral)
        except StopIteration:
            return

        goodEnough = (cD(image[centerY][centerX]) <= maxColDis)

        if goodEnough:
            dupe = 0
            for blob in blobs:
                if (blob[0]-centerX)**2 + (blob[1]-centerY)**2 <= blob[2]**2*1.5:
                    dupe = 1
                    break

            if not dupe:
                yield centerX, centerY

def search_rectQuadsection(blobs, xdim, ydim, spacing=1):
    quads = [ [[0,0],[xdim-1,ydim-1]] ]

    while len(quads):
        quad = quads.pop(0)
        centerX = (quad[0][0]+quad[1][0])//2
        centerY = (quad[0][1]+quad[1][1])//2

        if cD(image[centerY][centerX]) <= maxColDis:
            dupe = 0
            for blob in blobs:
                if (blob[0]-centerX)**2 + (blob[1]-centerY)**2 <= blob[2]**2*1.5:
                    dupe = 1
                    break

            if not dupe:
                yield centerX,centerY

        if abs(quad[0][0]-quad[1][0]) > 2*spacing and abs(quad[0][1]-quad[1][1]) > 2*spacing:
            quads.append([[centerX,centerY],[quad[0][0],quad[0][1]]])
            quads.append([[centerX,centerY],[quad[0][0],quad[1][1]]])
            quads.append([[centerX,centerY],[quad[1][0],quad[0][1]]])
            quads.append([[centerX,centerY],[quad[1][0],quad[1][1]]])

### SURVEY ###

def survey_floodFill(sX,sY):
    cX,cY,cR = (None,None,None)

    queue = [[sX,sY]]
    q = 0
    while q < len(queue):
        qx,qy = queue[q]

        for dx,dy in [[1,0],[0,1],[-1,0],[0,-1]]:
            x2, y2 = qx+dx, qy+dy

            if cD(image[y2][x2]) <= maxColDis:
                image[y2][x2] = col(255,0,0)
                queue.append([x2,y2])

        q += 1

    sumX,sumY = list(map(sum, zip(*queue)))
    cX, cY = round(sumX/q), round(sumY/q)
    cR = round(sqrt( q/pi ))

    return (cX,cY,cR)

def survey_circumcircle(sX,sY, numIterations=3, showVerts=0): #showVerts is also "minRadius"
    cX,cY,cR = (None,None,None)

    for k in range(numIterations): #refines circumcircle estimate by using prior guess
        vertices = []

        for dx,dy in [[1,0],[0,1],[-1,0],[0,-1]]:
            pX = sX+dx
            pY = sY+dy
            while cD(image[pY][pX]) <= maxColDis:
                pX += dx
                pY += dy
            vertices.append( [pX,pY] )

        v = vertices

        #get four possible circle centers/radii
        c1 = circumcircle(v[0],v[1],v[2])
        c2 = circumcircle(v[0],v[1],v[3])
        c3 = circumcircle(v[0],v[2],v[3])
        c4 = circumcircle(v[1],v[2],v[3])

        #average them
        avgX = (c1[0]+c2[0]+c3[0]+c4[0])/4
        avgY = (c1[1]+c2[1]+c3[1]+c4[1])/4
        avgR = (c1[2]+c2[2]+c3[2]+c4[2])/4

        c5 = [round(avgX),round(avgY),round(avgR)]
        cX,cY,cR = c5

        sX = cX
        sY = cY

    if showVerts and cR >= showVerts:
        image[cY][cX] = col(0,0,255)
        for vert in vertices:
            image[vert[1]][vert[0]] = col(255,0,0)

    return (cX,cY,cR)

### The actual start of the program ###

VIDEO_PATH = r".\Demo vids 1\demovid1_left0001-0075.mp4"

targetColor = [0,255,0]
cD = colDisSum(targetColor)
minRadius = 6 #pixels
maxColDis = 40

imgStart = clock()
origImage, xdim,ydim, frame = getVideoFrame(VIDEO_PATH)
imgEnd = clock() - imgStart
print("Time to get frame: %.3f seconds." % imgEnd)

search_algs = [['pixel by pixel', search_pixels],
               ['pixels spaced apart', search_spacedPixels, minRadius],
               ['square spiral', search_squareSpiral, minRadius],
               ['rectangular quadsection', search_rectQuadsection, minRadius],
               ]

survey_algs = [['flood fill', survey_floodFill],
               ['circumcircle', survey_circumcircle],
               ]

for search_alg in search_algs:
    for survey_alg in survey_algs:

        image = deepcopy(origImage)

        print()
        print("Search algorithm: %s" % search_alg[0])
        print("Survey algorithm: %s" % survey_alg[0])

        startTime = clock()
        times = []

        blobs = []
        search = search_alg[1](blobs, xdim,ydim, *search_alg[2:])

        while 1:
            try:
                sx,sy = next(search)
            except StopIteration:
                break

            surveyStart = clock()
            blob = survey_alg[1](sx,sy)
            surveyEnd = clock()
            times.append(surveyEnd - surveyStart)

            if blob[2] >= minRadius: blobs.append(blob)

        endTime = clock()
        print(" Total time: %.3f seconds" % (endTime-startTime))
        print(" Search time: %.3f seconds" % (endTime-startTime - sum(times)))
        print(" Number of surveys: %d" % len(times))
        print(" Number of blobs: %d" % len(blobs))
        print(" Survey time (total): %.3f seconds" % (sum(times)))
        print(" Survey time (average): %.3f seconds" % (sum(times)/len(times)))

### Note on different color distance methods:
### The difference between Euclidean distance and exact match is on the order
### of 0.02 - 0.03 seconds difference in the total time, which is about 10% of
### the total time. Not that much. I'll go with accuracy over speed here.
<|endoftext|>


<commit_before>import uuid
from django.contrib.contenttypes.models import ContentType
from django.contrib.gis.geos import Point

from django.core.cache import cache
from django.core.exceptions import ObjectDoesNotExist
from django.db import connection
from rest_framework import generics
from rest_framework.exceptions import NotFound
from rest_framework.pagination import PageNumberPagination
from rest_framework.response import Response
from rest_framework_json_api.views import AutoPrefetchMixin

from bluebottle.files.views import ImageContentView
from bluebottle.funding.models import Funding
from bluebottle.activities.models import Activity
from bluebottle.files.models import RelatedImage
from bluebottle.geo.models import Location
from bluebottle.initiatives.filters import InitiativeSearchFilter
from bluebottle.initiatives.models import Initiative
from bluebottle.initiatives.serializers import (
    InitiativeSerializer, InitiativeListSerializer, InitiativeReviewTransitionSerializer,
    InitiativeMapSerializer, InitiativeRedirectSerializer,
    RelatedInitiativeImageSerializer, ThemeSerializer
)
from bluebottle.bb_projects.models import ProjectTheme
from bluebottle.tasks.models import Task
from bluebottle.transitions.views import TransitionList
from bluebottle.utils.permissions import (
    OneOf, ResourcePermission, ResourceOwnerPermission, TenantConditionalOpenClose
)
from bluebottle.utils.views import (
    ListCreateAPIView, RetrieveUpdateAPIView, JsonApiViewMixin,
    CreateAPIView, ListAPIView
)

class InitiativeList(JsonApiViewMixin, AutoPrefetchMixin, ListCreateAPIView):
    queryset = Initiative.objects.prefetch_related(
        'place', 'location', 'owner', 'activity_manager', 'image', 'categories', 'theme'
    )

    def get_serializer_class(self):
        if self.request.method == 'POST' or self.request.GET.get('filter[owner.id]'):
            return InitiativeSerializer
        else:
            return InitiativeListSerializer

    permission_classes = (
        OneOf(ResourcePermission, ResourceOwnerPermission),
    )

    filter_backends = (
        InitiativeSearchFilter,
    )

    filter_fields = {
        'owner__id': ('exact', 'in',),
    }

    prefetch_for_includes = {
        'owner': ['owner'],
        'reviewer': ['reviewer'],
        'promoter': ['promoter'],
        'activity_manager': ['activity_manager'],
        'theme': ['theme'],
        'place': ['place'],
        'location': ['location'],
        'categories': ['categories'],
        'image': ['image'],
        'organization': ['organization'],
        'organization_contact': ['organization_contact'],
        'activities': ['activities'],
    }

    def perform_create(self, serializer):
        self.check_object_permissions(
            self.request,
            serializer.Meta.model(owner=self.request.user, **serializer.validated_data)
        )

        serializer.save(owner=self.request.user)

class TinyProjectPagination(PageNumberPagination):
    page_size = 10000

class InitiativeMapList(generics.ListAPIView):
    queryset = Initiative.objects
    serializer_class = InitiativeMapSerializer

    owner_filter_field = 'owner'

    def get_queryset(self):
        queryset = super(InitiativeMapList, self).get_queryset()
        queryset = queryset.filter(status='approved').all()
        if not Location.objects.count():
            # Skip initiatives without proper location
            queryset = queryset.exclude(place__position=Point(0, 0))
        return queryset

    def list(self, request, *args, **kwargs):
        cache_key = '{}.initiative_map_data'.format(connection.tenant.schema_name)
        data = cache.get(cache_key)
        if not data:
            queryset = self.get_queryset().order_by('created')
            serializer = self.get_serializer(queryset, many=True)
            data = serializer.data
            cache.set(cache_key, data)
        return Response(data)

class InitiativeDetail(JsonApiViewMixin, AutoPrefetchMixin, RetrieveUpdateAPIView):
    queryset = Initiative.objects.all()
    serializer_class = InitiativeSerializer

    permission_classes = (
        OneOf(ResourcePermission, ResourceOwnerPermission),
    )

    prefetch_for_includes = {
        'owner': ['owner'],
        'reviewer': ['reviewer'],
        'promoter': ['promoter'],
        'theme': ['theme'],
        'place': ['place'],
        'location': ['location'],
        'categories': ['categories'],
        'image': ['image'],
        'organization': ['organization'],
        'organization_contact': ['organization_contact'],
        'activities': ['activities'],
    }

class InitiativeImage(ImageContentView):
    queryset = Initiative.objects
    field = 'image'

class RelatedInitiativeImageList(JsonApiViewMixin, AutoPrefetchMixin, CreateAPIView):
    def get_queryset(self):
        return RelatedImage.objects.filter(
            content_type=ContentType.objects.get_for_model(Initiative)
        )

    serializer_class = RelatedInitiativeImageSerializer

    related_permission_classes = {
        'content_object': [ResourceOwnerPermission]
    }

    permission_classes = []

class RelatedInitiativeImageContent(ImageContentView):
    queryset = RelatedImage.objects
    field = 'image'

class InitiativeReviewTransitionList(TransitionList):
    serializer_class = InitiativeReviewTransitionSerializer
    queryset = Initiative.objects.all()

class ThemeList(JsonApiViewMixin, ListAPIView):
    serializer_class = ThemeSerializer
    queryset = ProjectTheme.objects.filter(disabled=False)
    permission_classes = [TenantConditionalOpenClose, ]

from collections import namedtuple
Instance = namedtuple('Instance', 'pk, route, params, target_params, target_route')

class InitiativeRedirectList(JsonApiViewMixin, CreateAPIView):
    serializer_class = InitiativeRedirectSerializer
    permission_classes = ()

    def perform_create(self, serializer):
        data = serializer.validated_data
        data['pk'] = str(uuid.uuid1())

        try:
            if data['route'] == 'project':
                initiative = Initiative.objects.get(slug=data['params']['project_id'])
                try:
                    funding = initiative.activities.instance_of(Funding)[0]
                    data['target_route'] = 'initiatives.activities.details.funding'
                    data['target_params'] = [funding.pk, funding.slug]
                except IndexError:
                    data['target_route'] = 'initiatives.details'
                    data['target_params'] = [initiative.pk, initiative.slug]
            elif data['route'] == 'task':
                task = Task.objects.get(id=data['params']['task_id'])
                activity = Activity.objects.get(pk=task.activity_id)
                data['target_route'] = 'initiatives.activities.details.{}'.format(
                    'event' if task.type == 'event' else 'assignment'
                )
                data['target_params'] = [activity.pk, activity.slug]
            else:
                raise NotFound()

            serializer.instance = Instance(**data)
            return data
        except ObjectDoesNotExist:
            raise NotFound()
<commit_msg>Make sure we return all themes in list endpoint<commit_after>import uuid
from django.contrib.contenttypes.models import ContentType
from django.contrib.gis.geos import Point

from django.core.cache import cache
from django.core.exceptions import ObjectDoesNotExist
from django.db import connection
from rest_framework import generics
from rest_framework.exceptions import NotFound
from rest_framework.pagination import PageNumberPagination
from rest_framework.response import Response
from rest_framework_json_api.views import AutoPrefetchMixin

from bluebottle.files.views import ImageContentView
from bluebottle.funding.models import Funding
from bluebottle.activities.models import Activity
from bluebottle.files.models import RelatedImage
from bluebottle.geo.models import Location
from bluebottle.initiatives.filters import InitiativeSearchFilter
from bluebottle.initiatives.models import Initiative
from bluebottle.initiatives.serializers import (
    InitiativeSerializer, InitiativeListSerializer, InitiativeReviewTransitionSerializer,
    InitiativeMapSerializer, InitiativeRedirectSerializer,
    RelatedInitiativeImageSerializer, ThemeSerializer
)
from bluebottle.bb_projects.models import ProjectTheme
from bluebottle.tasks.models import Task
from bluebottle.transitions.views import TransitionList
from bluebottle.utils.permissions import (
    OneOf, ResourcePermission, ResourceOwnerPermission, TenantConditionalOpenClose
)
from bluebottle.utils.views import (
    ListCreateAPIView, RetrieveUpdateAPIView, JsonApiViewMixin,
    CreateAPIView, ListAPIView
)

class InitiativeList(JsonApiViewMixin, AutoPrefetchMixin, ListCreateAPIView):
    queryset = Initiative.objects.prefetch_related(
        'place', 'location', 'owner', 'activity_manager', 'image', 'categories', 'theme'
    )

    def get_serializer_class(self):
        if self.request.method == 'POST' or self.request.GET.get('filter[owner.id]'):
            return InitiativeSerializer
        else:
            return InitiativeListSerializer

    permission_classes = (
        OneOf(ResourcePermission, ResourceOwnerPermission),
    )

    filter_backends = (
        InitiativeSearchFilter,
    )

    filter_fields = {
        'owner__id': ('exact', 'in',),
    }

    prefetch_for_includes = {
        'owner': ['owner'],
        'reviewer': ['reviewer'],
        'promoter': ['promoter'],
        'activity_manager': ['activity_manager'],
        'theme': ['theme'],
        'place': ['place'],
        'location': ['location'],
        'categories': ['categories'],
        'image': ['image'],
        'organization': ['organization'],
        'organization_contact': ['organization_contact'],
        'activities': ['activities'],
    }

    def perform_create(self, serializer):
        self.check_object_permissions(
            self.request,
            serializer.Meta.model(owner=self.request.user, **serializer.validated_data)
        )

        serializer.save(owner=self.request.user)

class TinyProjectPagination(PageNumberPagination):
    page_size = 10000

class InitiativeMapList(generics.ListAPIView):
    queryset = Initiative.objects
    serializer_class = InitiativeMapSerializer

    owner_filter_field = 'owner'

    def get_queryset(self):
        queryset = super(InitiativeMapList, self).get_queryset()
        queryset = queryset.filter(status='approved').all()
        if not Location.objects.count():
            # Skip initiatives without proper location
            queryset = queryset.exclude(place__position=Point(0, 0))
        return queryset

    def list(self, request, *args, **kwargs):
        cache_key = '{}.initiative_map_data'.format(connection.tenant.schema_name)
        data = cache.get(cache_key)
        if not data:
            queryset = self.get_queryset().order_by('created')
            serializer = self.get_serializer(queryset, many=True)
            data = serializer.data
            cache.set(cache_key, data)
        return Response(data)

class InitiativeDetail(JsonApiViewMixin, AutoPrefetchMixin, RetrieveUpdateAPIView):
    queryset = Initiative.objects.all()
    serializer_class = InitiativeSerializer

    permission_classes = (
        OneOf(ResourcePermission, ResourceOwnerPermission),
    )

    prefetch_for_includes = {
        'owner': ['owner'],
        'reviewer': ['reviewer'],
        'promoter': ['promoter'],
        'theme': ['theme'],
        'place': ['place'],
        'location': ['location'],
        'categories': ['categories'],
        'image': ['image'],
        'organization': ['organization'],
        'organization_contact': ['organization_contact'],
        'activities': ['activities'],
    }

class InitiativeImage(ImageContentView):
    queryset = Initiative.objects
    field = 'image'

class RelatedInitiativeImageList(JsonApiViewMixin, AutoPrefetchMixin, CreateAPIView):
    def get_queryset(self):
        return RelatedImage.objects.filter(
            content_type=ContentType.objects.get_for_model(Initiative)
        )

    serializer_class = RelatedInitiativeImageSerializer

    related_permission_classes = {
        'content_object': [ResourceOwnerPermission]
    }

    permission_classes = []

class RelatedInitiativeImageContent(ImageContentView):
    queryset = RelatedImage.objects
    field = 'image'

class InitiativeReviewTransitionList(TransitionList):
    serializer_class = InitiativeReviewTransitionSerializer
    queryset = Initiative.objects.all()

class ThemePagination(PageNumberPagination):
    page_size = 10000

class ThemeList(JsonApiViewMixin, ListAPIView):
    serializer_class = ThemeSerializer
    queryset = ProjectTheme.objects.filter(disabled=False)
    permission_classes = [TenantConditionalOpenClose, ]
    pagination_class = ThemePagination

from collections import namedtuple
Instance = namedtuple('Instance', 'pk, route, params, target_params, target_route')

class InitiativeRedirectList(JsonApiViewMixin, CreateAPIView):
    serializer_class = InitiativeRedirectSerializer
    permission_classes = ()

    def perform_create(self, serializer):
        data = serializer.validated_data
        data['pk'] = str(uuid.uuid1())

        try:
            if data['route'] == 'project':
                initiative = Initiative.objects.get(slug=data['params']['project_id'])
                try:
                    funding = initiative.activities.instance_of(Funding)[0]
                    data['target_route'] = 'initiatives.activities.details.funding'
                    data['target_params'] = [funding.pk, funding.slug]
                except IndexError:
                    data['target_route'] = 'initiatives.details'
                    data['target_params'] = [initiative.pk, initiative.slug]
            elif data['route'] == 'task':
                task = Task.objects.get(id=data['params']['task_id'])
                activity = Activity.objects.get(pk=task.activity_id)
                data['target_route'] = 'initiatives.activities.details.{}'.format(
                    'event' if task.type == 'event' else 'assignment'
                )
                data['target_params'] = [activity.pk, activity.slug]
            else:
                raise NotFound()

            serializer.instance = Instance(**data)
            return data
        except ObjectDoesNotExist:
            raise NotFound()
<|endoftext|>


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

import bpy

from .. import util

class AppleseedMaterialSlots(bpy.types.Panel):
    bl_space_type = 'PROPERTIES'
    bl_label = "Material"
    bl_region_type = 'WINDOW'
    bl_context = "material"
    # COMPAT_ENGINES must be defined in each subclass, external engines can add themselves here

    @classmethod
    def poll(cls, context):
        engine = context.scene.render.engine
        return (context.material or context.object) and engine == 'APPLESEED_RENDER'

    def draw(self, context):
        layout = self.layout

        mat = context.material
        ob = context.object
        slot = context.material_slot
        space = context.space_data

        if ob:
            is_sortable = (len(ob.material_slots) > 1)

            rows = 1
            if is_sortable:
                rows = 4

            row = layout.row()

            row.template_list("MATERIAL_UL_matslots", "", ob, "material_slots", ob, "active_material_index", rows=rows)

            col = row.column(align=True)
            col.operator("object.material_slot_add", icon='ZOOMIN', text="")
            col.operator("object.material_slot_remove", icon='ZOOMOUT', text="")

            col.menu("MATERIAL_MT_specials", icon='DOWNARROW_HLT', text="")

            if is_sortable:
                col.separator()

                col.operator("object.material_slot_move", icon='TRIA_UP', text="").direction = 'UP'
                col.operator("object.material_slot_move", icon='TRIA_DOWN', text="").direction = 'DOWN'

            if ob.mode == 'EDIT':
                row = layout.row(align=True)
                row.operator("object.material_slot_assign", text="Assign")
                row.operator("object.material_slot_select", text="Select")
                row.operator("object.material_slot_deselect", text="Deselect")

        split = layout.split(percentage=0.65)

        if ob:
            split.template_ID(ob, "active_material", new="appleseed.new_mat")
            row = split.row()
            if slot:
                row.prop(slot, "link", text="")
            else:
                row.label()
        elif mat:
            split.template_ID(space, "pin_id")
            split.separator()

class AppleseedMaterialPreview(bpy.types.Panel):
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'

    bl_context = "material"
    bl_label = "Preview"
    COMPAT_ENGINES = {'APPLESEED_RENDER'}

    @classmethod
    def poll(cls, context):
        return context.scene.render.engine in cls.COMPAT_ENGINES and context.object is not None and context.object.active_material is not None

    def draw(self, context):
        layout = self.layout
        obj = context.object
        material = obj.active_material
        asr_mat = material.appleseed

        layout.template_preview(context.material, show_buttons=False)
        layout.prop(asr_mat, "preview_quality", text="Preview Quality")

class AppleseedMaterialShading(bpy.types.Panel):
    bl_label = 'Shader Model'
    bl_space_type = "PROPERTIES"
    bl_region_type = "WINDOW"
    bl_context = "material"
    COMPAT_ENGINES = {'APPLESEED_RENDER'}

    @classmethod
    def poll(cls, context):
        material = context.object.active_material is not None
        if material:
            renderer = context.scene.render.engine == 'APPLESEED_RENDER'
            obj = context.object is not None
            obj_type = context.object.type == 'MESH'
            return renderer and obj and obj_type
        return False

    def draw(self, context):
        layout = self.layout
        object = context.object
        material = object.active_material
        asr_mat = material.appleseed

        if material.appleseed.osl_node_tree is None:
            layout.operator('appleseed.add_osl_nodetree', text="Add appleseed Material Node", icon='NODETREE')
        else:
            layout.operator('appleseed.view_nodetree', text="View Nodetree")

        layout.prop(asr_mat, "shader_lighting_samples", text="Lighting Samples")

def register():
    util.safe_register_class(AppleseedMaterialSlots)
    bpy.types.MATERIAL_PT_custom_props.COMPAT_ENGINES.add('APPLESEED_RENDER')
    util.safe_register_class(AppleseedMaterialPreview)
    util.safe_register_class(AppleseedMaterialShading)

def unregister():
    util.safe_unregister_class(AppleseedMaterialShading)
    util.safe_unregister_class(AppleseedMaterialPreview)
    bpy.types.MATERIAL_PT_custom_props.COMPAT_ENGINES.remove('APPLESEED_RENDER')
    util.safe_unregister_class(AppleseedMaterialSlots)
<commit_msg>Tweak property names.<commit_after>#
# This source file is part of appleseed.
# Visit https://appleseedhq.net/ for additional information and resources.
#
# This software is released under the MIT license.
#
# Copyright (c) 2014-2018 The appleseedhq Organization
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#

import bpy

from .. import util

class AppleseedMaterialSlots(bpy.types.Panel):
    bl_space_type = 'PROPERTIES'
    bl_label = "Material"
    bl_region_type = 'WINDOW'
    bl_context = "material"
    # COMPAT_ENGINES must be defined in each subclass, external engines can add themselves here

    @classmethod
    def poll(cls, context):
        engine = context.scene.render.engine
        return (context.material or context.object) and engine == 'APPLESEED_RENDER'

    def draw(self, context):
        layout = self.layout

        mat = context.material
        ob = context.object
        slot = context.material_slot
        space = context.space_data

        if ob:
            is_sortable = (len(ob.material_slots) > 1)

            rows = 1
            if is_sortable:
                rows = 4

            row = layout.row()

            row.template_list("MATERIAL_UL_matslots", "", ob, "material_slots", ob, "active_material_index", rows=rows)

            col = row.column(align=True)
            col.operator("object.material_slot_add", icon='ZOOMIN', text="")
            col.operator("object.material_slot_remove", icon='ZOOMOUT', text="")

            col.menu("MATERIAL_MT_specials", icon='DOWNARROW_HLT', text="")

            if is_sortable:
                col.separator()

                col.operator("object.material_slot_move", icon='TRIA_UP', text="").direction = 'UP'
                col.operator("object.material_slot_move", icon='TRIA_DOWN', text="").direction = 'DOWN'

            if ob.mode == 'EDIT':
                row = layout.row(align=True)
                row.operator("object.material_slot_assign", text="Assign")
                row.operator("object.material_slot_select", text="Select")
                row.operator("object.material_slot_deselect", text="Deselect")

        split = layout.split(percentage=0.65)

        if ob:
            split.template_ID(ob, "active_material", new="appleseed.new_mat")
            row = split.row()
            if slot:
                row.prop(slot, "link", text="")
            else:
                row.label()
        elif mat:
            split.template_ID(space, "pin_id")
            split.separator()

class AppleseedMaterialPreview(bpy.types.Panel):
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'

    bl_context = "material"
    bl_label = "Preview"
    COMPAT_ENGINES = {'APPLESEED_RENDER'}

    @classmethod
    def poll(cls, context):
        return context.scene.render.engine in cls.COMPAT_ENGINES and context.object is not None and context.object.active_material is not None

    def draw(self, context):
        layout = self.layout
        obj = context.object
        material = obj.active_material
        asr_mat = material.appleseed

        layout.template_preview(context.material, show_buttons=False)
        layout.prop(asr_mat, "preview_quality", text="Preview Quality")

class AppleseedMaterialShading(bpy.types.Panel):
    bl_label = 'Shader Model'
    bl_space_type = "PROPERTIES"
    bl_region_type = "WINDOW"
    bl_context = "material"
    COMPAT_ENGINES = {'APPLESEED_RENDER'}

    @classmethod
    def poll(cls, context):
        material = context.object.active_material is not None
        if material:
            renderer = context.scene.render.engine == 'APPLESEED_RENDER'
            obj = context.object is not None
            obj_type = context.object.type == 'MESH'
            return renderer and obj and obj_type
        return False

    def draw(self, context):
        layout = self.layout
        obj = context.object
        material = obj.active_material
        asr_mat = material.appleseed

        if material.appleseed.osl_node_tree is None:
            layout.operator('appleseed.add_osl_nodetree', text="Add appleseed Material Node", icon='NODETREE')
        else:
            layout.operator('appleseed.view_nodetree', text="View Nodetree")

        layout.prop(asr_mat, "shader_lighting_samples", text="Lighting Samples")

def register():
    util.safe_register_class(AppleseedMaterialSlots)
    bpy.types.MATERIAL_PT_custom_props.COMPAT_ENGINES.add('APPLESEED_RENDER')
    util.safe_register_class(AppleseedMaterialPreview)
    util.safe_register_class(AppleseedMaterialShading)

def unregister():
    util.safe_unregister_class(AppleseedMaterialShading)
    util.safe_unregister_class(AppleseedMaterialPreview)
    bpy.types.MATERIAL_PT_custom_props.COMPAT_ENGINES.remove('APPLESEED_RENDER')
    util.safe_unregister_class(AppleseedMaterialSlots)
<|endoftext|>


<commit_before># Copyright 2009 - Erik Karulf - MIT License - See LICENSE.txt

"""
A collection of useful decorators. When combined these decorators allow an
administartor to setup a failry precise authentication system.

These Decorators are designed specifically to work with the access control
mechanisms in mod_wsgi.

The following assumptions are made:
 * Wrapped functions' argument lists are of the format: (environ, user, ...)
 * Wrapped functions are assumed to return True, False, or None
"""

import time
import logging
import collections
import random
import sqlite3
import os
from functools import wraps
from hashlib import sha256
from wsgi_googleauth.util import valid_email, parse_email, encode_value, decode_value

class Cache(object):
    """
    Caches the result of the mod_wsgi authorization functions for a given
    amount of time. This uses an internal datastructure instead of an external
    key-value store like memcached or redis.

    >>> import time
    >>>
    >>> @Cache(":memory:", timeout=10.0)
    ... def test(environ, user, password):
    ...     time.sleep(5.0)
    ...     return password == 'password'
    ...
    >>> # cache miss
    ... cache_begin = begin = time.time()
    >>> result = test(None, 'user', 'password')
    >>> assert result
    >>> time.time() - begin < 1.0
    False
    >>>
    >>> # cache hit
    ... begin = time.time()
    >>> result = test(None, 'user', 'password')
    >>> assert result
    >>> time.time() - begin < 1.0
    True
    >>>
    >>> # cache miss
    ... begin = time.time()
    >>> result = test(None, 'user', 'password2')
    >>> assert not result
    >>> time.time() - begin < 1.0
    False
    >>>
    >>> # wait for timeout
    ... while time.time() - cache_begin < 20.0:
    ...     time.sleep(1.0)
    ...
    >>> # cache miss
    ... begin = time.time()
    >>> result = test(None, 'user', 'password')
    >>> assert result
    >>> time.time() - begin < 1.0
    False
    >>>
    """
    SALT_NAME = 'salt'

    def _normalize_filename(self, filename):
        if filename[0] == '~':
            # Relative to policy file
            import policy
            filename = os.path.abspath(os.path.join(os.path.dirname(policy.__file__), '..')) + filename[1:]
        return filename

    def __init__(self, filename, **options):
        # Options
        self.timeout = options.get('timeout', 60*30) # 30 minutes
        self.ignore_environ = options.get('ignore_environ', True) # by default, ignore the environment variables
        self.cache_true_only = options.get('cache_true_only', True)
        self.salt = None

        # Connect to the database
        filename = self._normalize_filename(filename)
        # TODO: Enforce 0600 perms on UNIX?
        self.conn = sqlite3.connect(filename)
        c = self.conn.cursor()
        try:
            c.execute("CREATE TABLE cache (key TEXT UNIQUE, value TEXT, expiration INTEGER);")
        except sqlite3.OperationalError:
            # Cache table exists
            c.execute('DELETE FROM cache WHERE expiration < ?;', (time.time(),))
            c.execute("SELECT value FROM cache WHERE key = ? LIMIT 1", (Cache.SALT_NAME, ))
            for row in c:
                self.salt = decode_value(row[0])
        else:
            # Cache table just created
            self.salt = os.urandom(128)
            c.execute("CREATE INDEX IF NOT EXISTS expiration ON cache (expiration ASC);")
            c.execute('INSERT INTO cache(key, value) VALUES (?, ?);', (Cache.SALT_NAME, encode_value(self.salt)))
        self.conn.commit()
        assert self.salt is not None, "wsgi_googleauth Error: Salt not initialized"

    def __call__(self, f):
        @wraps(f)
        def wrapper(*args):
            # Invariant: cache is sorted by timeout
            # Requirement: time.time is a monotonically increasing function
            c = self.conn.cursor()

            # Generate Key
            # TODO: Use a more secure alternative like scrypt or PBKDF2
            hash_args = args[1:] if self.ignore_environ else args[:]
            hash_args = (self.salt,) + hash_args
            key = sha256(encode_value(hash_args)).hexdigest()

            # Cache Lookup
            found = False
            c.execute("SELECT value FROM cache WHERE key = ? AND expiration > ? LIMIT 1;", (key, time.time()))
            for row in c:
                result = decode_value(row[0])
                found = True

            # Query Function
            if not found:
                result = f(*args)
                if not self.cache_true_only or (self.cache_true_only and result is True):
                    timeout = time.time() + self.timeout
                    try:
                        c.execute("INSERT INTO cache(key, value, expiration) VALUES (?, ?, ?);", (key, encode_value(result), timeout))
                    except sqlite3.IntegrityError:
                        c.execute("UPDATE cache SET expiration = ? WHERE key = ?;", (timeout, key))
                    self.conn.commit()
            return result
        return wrapper

class DefaultDomain(object):
    """
    Fills in a domain name if none is none is specified. If this does not yield
    a valid email address it passes the username through unchanged.

    Example:
    >>> @DefaultDomain('example.com')
    ... def test(environ, name, password):
    ...     print name
    ...
    >>> test(None, 'erik', None)
    erik@example.com
    >>> test(None, 'erik@karulf.com', None)
    erik@karulf.com
    >>> test(None, ';', None)
    ;
    >>>
    """
    def __init__(self, domain):
        self.domain = domain

    def __call__(self, f):
        @wraps(f)
        def wrapper(environ, name, *args):
            if not valid_email(name):
                new_name = "{0}@{1}".format(name, self.domain)
                if valid_email(new_name):
                    logging.info("Defaulting username %s to %s" % (name, new_name))
                    name = new_name
            return f(environ, name, *args)
        return wrapper

class RequireDomain(object):
    """
    Requires both the username to be a valid email and the domain name to match
    a value given at initialization.

    Example:
    >>> @RequireDomain('example.com')
    ... def test(environ, name, password):
    ...     return True
    ...
    >>> test(None, 'erik', None)
    >>> test(None, 'erik@karulf.com', None)
    >>> test(None, 'erik@example.com', None)
    True
    >>>
    """
    def __init__(self, domain):
        self.domain = domain

    def __call__(self, f):
        @wraps(f)
        def wrapper(environ, name, *args):
            match = parse_email(name)
            if match is not None and match[1] == self.domain:
                return f(environ, name, *args)
            else:
                logging.info("Ignoring username %s as it does not match domain %s" % (name, self.domain))
                return None
        return wrapper

class RequireValidEmail(object):
    """
    Requires a valid email address

    Example:
    >>> @RequireValidEmail()
    ... def test(environ, user, password):
    ...     return True
    ...
    >>> test(None, 'erik', 'password')
    >>> test(None, 'erik@example.com', 'password')
    True
    """
    def __init__(self):
        pass

    def __call__(self, f):
        @wraps(f)
        def wrapper(environ, name, *args):
            if valid_email(name):
                return f(environ, name, *args)
            else:
                logging.info("Ignoring username %s as it is not a valid email")
                return None
        return wrapper

<commit_msg>Replacing str.format for Python 2.5 compatibility<commit_after># Copyright 2009 - Erik Karulf - MIT License - See LICENSE.txt

"""
A collection of useful decorators. When combined these decorators allow an
administartor to setup a failry precise authentication system.

These Decorators are designed specifically to work with the access control
mechanisms in mod_wsgi.

The following assumptions are made:
 * Wrapped functions' argument lists are of the format: (environ, user, ...)
 * Wrapped functions are assumed to return True, False, or None
"""

import time
import logging
import collections
import random
import sqlite3
import os
from functools import wraps
from hashlib import sha256
from wsgi_googleauth.util import valid_email, parse_email, encode_value, decode_value

class Cache(object):
    """
    Caches the result of the mod_wsgi authorization functions for a given
    amount of time. This uses an internal datastructure instead of an external
    key-value store like memcached or redis.

    >>> import time
    >>>
    >>> @Cache(":memory:", timeout=10.0)
    ... def test(environ, user, password):
    ...     time.sleep(5.0)
    ...     return password == 'password'
    ...
    >>> # cache miss
    ... cache_begin = begin = time.time()
    >>> result = test(None, 'user', 'password')
    >>> assert result
    >>> time.time() - begin < 1.0
    False
    >>>
    >>> # cache hit
    ... begin = time.time()
    >>> result = test(None, 'user', 'password')
    >>> assert result
    >>> time.time() - begin < 1.0
    True
    >>>
    >>> # cache miss
    ... begin = time.time()
    >>> result = test(None, 'user', 'password2')
    >>> assert not result
    >>> time.time() - begin < 1.0
    False
    >>>
    >>> # wait for timeout
    ... while time.time() - cache_begin < 20.0:
    ...     time.sleep(1.0)
    ...
    >>> # cache miss
    ... begin = time.time()
    >>> result = test(None, 'user', 'password')
    >>> assert result
    >>> time.time() - begin < 1.0
    False
    >>>
    """
    SALT_NAME = 'salt'

    def _normalize_filename(self, filename):
        if filename[0] == '~':
            # Relative to policy file
            import policy
            filename = os.path.abspath(os.path.join(os.path.dirname(policy.__file__), '..')) + filename[1:]
        return filename

    def __init__(self, filename, **options):
        # Options
        self.timeout = options.get('timeout', 60*30) # 30 minutes
        self.ignore_environ = options.get('ignore_environ', True) # by default, ignore the environment variables
        self.cache_true_only = options.get('cache_true_only', True)
        self.salt = None

        # Connect to the database
        filename = self._normalize_filename(filename)
        # TODO: Enforce 0600 perms on UNIX?
        self.conn = sqlite3.connect(filename)
        c = self.conn.cursor()
        try:
            c.execute("CREATE TABLE cache (key TEXT UNIQUE, value TEXT, expiration INTEGER);")
        except sqlite3.OperationalError:
            # Cache table exists
            c.execute('DELETE FROM cache WHERE expiration < ?;', (time.time(),))
            c.execute("SELECT value FROM cache WHERE key = ? LIMIT 1", (Cache.SALT_NAME, ))
            for row in c:
                self.salt = decode_value(row[0])
        else:
            # Cache table just created
            self.salt = os.urandom(128)
            c.execute("CREATE INDEX IF NOT EXISTS expiration ON cache (expiration ASC);")
            c.execute('INSERT INTO cache(key, value) VALUES (?, ?);', (Cache.SALT_NAME, encode_value(self.salt)))
        self.conn.commit()
        assert self.salt is not None, "wsgi_googleauth Error: Salt not initialized"

    def __call__(self, f):
        @wraps(f)
        def wrapper(*args):
            # Invariant: cache is sorted by timeout
            # Requirement: time.time is a monotonically increasing function
            c = self.conn.cursor()

            # Generate Key
            # TODO: Use a more secure alternative like scrypt or PBKDF2
            hash_args = args[1:] if self.ignore_environ else args[:]
            hash_args = (self.salt,) + hash_args
            key = sha256(encode_value(hash_args)).hexdigest()

            # Cache Lookup
            found = False
            c.execute("SELECT value FROM cache WHERE key = ? AND expiration > ? LIMIT 1;", (key, time.time()))
            for row in c:
                result = decode_value(row[0])
                found = True

            # Query Function
            if not found:
                result = f(*args)
                if not self.cache_true_only or (self.cache_true_only and result is True):
                    timeout = time.time() + self.timeout
                    try:
                        c.execute("INSERT INTO cache(key, value, expiration) VALUES (?, ?, ?);", (key, encode_value(result), timeout))
                    except sqlite3.IntegrityError:
                        c.execute("UPDATE cache SET expiration = ? WHERE key = ?;", (timeout, key))
                    self.conn.commit()
            return result
        return wrapper

class DefaultDomain(object):
    """
    Fills in a domain name if none is none is specified. If this does not yield
    a valid email address it passes the username through unchanged.

    Example:
    >>> @DefaultDomain('example.com')
    ... def test(environ, name, password):
    ...     print name
    ...
    >>> test(None, 'erik', None)
    erik@example.com
    >>> test(None, 'erik@karulf.com', None)
    erik@karulf.com
    >>> test(None, ';', None)
    ;
    >>>
    """
    def __init__(self, domain):
        self.domain = domain

    def __call__(self, f):
        @wraps(f)
        def wrapper(environ, name, *args):
            if not valid_email(name):
                new_name = "%s@%s" % (name, self.domain)
                if valid_email(new_name):
                    logging.info("Defaulting username %s to %s" % (name, new_name))
                    name = new_name
            return f(environ, name, *args)
        return wrapper

class RequireDomain(object):
    """
    Requires both the username to be a valid email and the domain name to match
    a value given at initialization.

    Example:
    >>> @RequireDomain('example.com')
    ... def test(environ, name, password):
    ...     return True
    ...
    >>> test(None, 'erik', None)
    >>> test(None, 'erik@karulf.com', None)
    >>> test(None, 'erik@example.com', None)
    True
    >>>
    """
    def __init__(self, domain):
        self.domain = domain

    def __call__(self, f):
        @wraps(f)
        def wrapper(environ, name, *args):
            match = parse_email(name)
            if match is not None and match[1] == self.domain:
                return f(environ, name, *args)
            else:
                logging.info("Ignoring username %s as it does not match domain %s" % (name, self.domain))
                return None
        return wrapper

class RequireValidEmail(object):
    """
    Requires a valid email address

    Example:
    >>> @RequireValidEmail()
    ... def test(environ, user, password):
    ...     return True
    ...
    >>> test(None, 'erik', 'password')
    >>> test(None, 'erik@example.com', 'password')
    True
    """
    def __init__(self):
        pass

    def __call__(self, f):
        @wraps(f)
        def wrapper(environ, name, *args):
            if valid_email(name):
                return f(environ, name, *args)
            else:
                logging.info("Ignoring username %s as it is not a valid email")
                return None
        return wrapper

<|endoftext|>