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

from __future__ import print_function

import sys
import pysam
import os

from slamdunk.utils.misc import replaceExtension, getSampleInfo, SlamSeqInfo, md5, callR, getPlotter  # @UnresolvedImport
from slamdunk.utils.BedReader import BedIterator  # @UnresolvedImport

from slamdunk.utils import SNPtools  # @UnresolvedImport
from slamdunk.slamseq.SlamSeqFile import SlamSeqBamFile, ReadDirection, SlamSeqInterval  # @UnresolvedImport

from slamdunk.version import __version__, __bam_version__, __count_version__  # @UnresolvedImport

def collapse(expandedCSV, collapsedCSV, log):

    tcDict = {}

    outCSV = open(collapsedCSV, 'w')

    readNumber = 0

    with open(expandedCSV, 'r') as f:

        # Skip header
#         next(f)

        for line in f:

            fields = line.rstrip().split('\t')
            if (len(fields) == 14) :

                gene = fields[3]
                length = fields[4]
                Tcontent = fields[8]
                coverageOnTs = fields[9]
                conversionsOnTs = fields[10]
                readCount = fields[11]
                tcReadCount = fields[12]
                multimapCount = fields[13]

                if (gene in tcDict.keys()) :
                    tcDict[gene]['length'] += max(int(length),0)
                    tcDict[gene]['Tcontent'] += max(int(Tcontent),0)
                    tcDict[gene]['coverageOnTs'] += max(int(coverageOnTs),0)
                    tcDict[gene]['conversionsOnTs'] += max(int(conversionsOnTs),0)
                    tcDict[gene]['readCount'] += max(int(readCount),0)
                    tcDict[gene]['tcReadCount'] += max(int(tcReadCount),0)
                    tcDict[gene]['multimapCount'] += max(int(multimapCount),0)
                else :
                    tcDict[gene] = {}
                    tcDict[gene]['length'] = max(int(length),0)
                    tcDict[gene]['Tcontent'] = max(int(Tcontent),0)
                    tcDict[gene]['coverageOnTs'] = max(int(coverageOnTs),0)
                    tcDict[gene]['conversionsOnTs'] = max(int(conversionsOnTs),0)
                    tcDict[gene]['readCount'] = max(int(readCount),0)
                    tcDict[gene]['tcReadCount'] = max(int(tcReadCount),0)
                    tcDict[gene]['multimapCount'] = max(int(multimapCount),0)

                readNumber += int(readCount)

        else :
            print("Error in TC file format - unexpected number of fields (" + str(len(fields)) + ") in the following line:\n" + line, file=log)

    print("gene_name", "length", "readsCPM", "conversionRate", "Tcontent", "coverageOnTs", "conversionsOnTs", "readCount", "tcReadCount", "multimapCount", sep='\t', file=outCSV)

    for gene in sorted(tcDict.keys()) :

        print(gene,end="\t",file=outCSV)
        print(tcDict[gene]['length'],end="\t",file=outCSV)
        print(float(tcDict[gene]['readCount']) / float(readNumber) * 1000000,end="\t",file=outCSV)
        conversionRate = 0
        if (tcDict[gene]['coverageOnTs'] > 0) :
            conversionRate = float(tcDict[gene]['conversionsOnTs']) / tcDict[gene]['coverageOnTs']
        print(conversionRate,end="\t",file=outCSV)
        print(tcDict[gene]['Tcontent'],end="\t",file=outCSV)
        print(tcDict[gene]['coverageOnTs'],end="\t",file=outCSV)
        print(tcDict[gene]['conversionsOnTs'],end="\t",file=outCSV)
        print(tcDict[gene]['readCount'],end="\t",file=outCSV)
        print(tcDict[gene]['tcReadCount'],end="\t",file=outCSV)
        print(tcDict[gene]['multimapCount'],file=outCSV)

    outCSV.close()

def getMean(values, skipZeros=True):
    count = 0.0
    totalSum = 0.0
    for value in values:
        if not skipZeros or value > 0:
            totalSum = totalSum + value
            count += 1
    if(count > 0):
        return totalSum / count
    else:
        return 0.0

def computeTconversions(ref, bed, snpsFile, bam, maxReadLength, minQual, outputCSV, outputBedgraphPlus, outputBedgraphMinus, strictTCs, log):

    referenceFile = pysam.FastaFile(ref)

    sampleInfo = getSampleInfo(bam)

    slamseqInfo = SlamSeqInfo(bam)
    readNumber = slamseqInfo.MappedReads

    bedMD5 = md5(bed)

    fileNameTest = replaceExtension(outputCSV, ".tsv", "_perread")
    fileTest = open(fileNameTest,'w')
    print("#slamdunk v" + __version__, __count_version__, "sample info:", sampleInfo.Name, sampleInfo.ID, sampleInfo.Type, sampleInfo.Time, sep="\t", file=fileTest)
    print("#annotation:", os.path.basename(bed), bedMD5, sep="\t", file=fileTest)
    print("utr", "n", "k", file=fileTest)

    fileCSV = open(outputCSV,'w')
    print("#slamdunk v" + __version__, __count_version__, "sample info:", sampleInfo.Name, sampleInfo.ID, sampleInfo.Type, sampleInfo.Time, sep="\t", file=fileCSV)
    print("#annotation:", os.path.basename(bed), bedMD5, sep="\t", file=fileCSV)
    print(SlamSeqInterval.Header, file=fileCSV)

    snps = SNPtools.SNPDictionary(snpsFile)
    snps.read()

    #Go through one chr after the other
    testFile = SlamSeqBamFile(bam, ref, snps)
    if not testFile.bamVersion == __bam_version__:
        raise RuntimeError("Wrong filtered BAM file version detected (" + testFile.bamVersion + "). Expected version " + __bam_version__ + ". Please rerun slamdunk filter.")

    bedMD5 = md5(bed)
    if slamseqInfo.AnnotationMD5 != bedMD5:
        print("Warning: MD5 checksum of annotation (" + bedMD5 + ") does not matched MD5 in filtered BAM files (" + slamseqInfo.AnnotationMD5 + "). Most probably the annotation filed changed after the filtered BAM files were created.", file=log)

    conversionBedGraph = {}

    for utr in BedIterator(bed):
        Tcontent = 0
        slamSeqUtr = SlamSeqInterval(utr.chromosome, utr.start, utr.stop, utr.strand, utr.name, Tcontent, 0, -1, 0, 0, 0, 0, 0)
        if(not utr.hasStrand()):
            raise RuntimeError("Input BED file does not contain stranded intervals.")

        if utr.start < 0:
            raise RuntimeError("Negativ start coordinate found. Please check the following entry in your BED file: " + utr)
        # Retreive reference sequence
        region = utr.chromosome + ":" + str(utr.start + 1) + "-" + str(utr.stop)

        if(utr.chromosome in list(referenceFile.references)):
            #print(refRegion,file=sys.stderr)
            refSeq = referenceFile.fetch(region=region).upper()

            if (utr.strand == "-") :
                #refSeq = complement(refSeq[::-1])
                Tcontent = refSeq.count("A")
            else :
                Tcontent = refSeq.count("T")

            slamSeqUtr._Tcontent = Tcontent

        readIterator = testFile.readInRegion(utr.chromosome, utr.start, utr.stop, utr.strand, maxReadLength)

        tcCountUtr = [0] * utr.getLength()
        coverageUtr = [0] * utr.getLength()

        countFwd = 0
        tcCountFwd = 0
        countRev = 0
        tCountRev = 0

        multiMapFwd = 0
        multiMapRev = 0

        for read in readIterator:

            # Overwrite any conversions for non-TC reads (reads with < 2 TC conversions)
            if (not read.isTcRead and strictTCs) :
                read.tcCount = 0
                read.mismatches = []
                read.conversionRates = 0.0
                read.tcRate = 0.0

            if(read.direction == ReadDirection.Reverse):
                countRev += 1
                if read.tcCount > 0:
                    tCountRev += 1
                if read.isMultimapper:
                    multiMapRev += 1
            else:
                countFwd += 1
                if read.tcCount > 0:
                    tcCountFwd += 1
                if read.isMultimapper:
                    multiMapFwd += 1

            for mismatch in read.mismatches:
                if(mismatch.isTCMismatch(read.direction == ReadDirection.Reverse) and mismatch.referencePosition >= 0 and mismatch.referencePosition < utr.getLength()):
                    tcCountUtr[mismatch.referencePosition] += 1

            testN = read.getTcount()
            testk = 0
            for mismatch in read.mismatches:
                if(mismatch.referencePosition >= 0 and mismatch.referencePosition < utr.getLength()):
                    if(mismatch.isT(read.direction == ReadDirection.Reverse)):
                        testN += 1
                    if(mismatch.isTCMismatch(read.direction == ReadDirection.Reverse)):
                        testk += 1
            #print(utr.name, read.name, read.direction, testN, testk, read.sequence, sep="\t")
            print(utr.name, testN, testk, sep="\t", file=fileTest)

            for i in xrange(read.startRefPos, read.endRefPos):
                if(i >= 0 and i < utr.getLength()):
                    coverageUtr[i] += 1

        if((utr.strand == "+" and countFwd > 0) or (utr.strand == "-" and countRev > 0)):
            tcRateUtr = [ x * 100.0 / y if y > 0 else 0 for x, y in zip(tcCountUtr, coverageUtr)]

            readCount = countFwd
            tcReadCount = tcCountFwd
            multiMapCount = multiMapFwd

            if(utr.strand == "-"):
                readCount = countRev
                tcReadCount = tCountRev
                multiMapCount = multiMapRev

            if((utr.strand == "-" and countFwd > countRev) or (utr.strand == "+" and countRev > countFwd)):
                print("Warning: " + utr.name + " is located on the " + utr.strand + " strand but read counts are higher for the opposite strand (fwd: " + countFwd + ", rev: " + countRev + ")", file=sys.stderr)

            refSeq = readIterator.getRefSeq()

            # Get number of covered Ts/As in the UTR and compute average conversion rate for all covered Ts/As
            coveredTcount = 0
            avgConversationRate = 0
            coveredPositions = 0
            # Get number of reads on T positions and number of reads with T->C conversions on T positions
            coverageOnTs = 0
            conversionsOnTs = 0

            for position in xrange(0, len(coverageUtr)):

                if(coverageUtr[position] > 0 and ((utr.strand == "+" and refSeq[position] == "T") or (utr.strand == "-" and refSeq[position] == "A"))):
                    coveredTcount += 1
                    avgConversationRate += tcRateUtr[position]

                    coverageOnTs += coverageUtr[position]
                    conversionsOnTs += tcCountUtr[position]
                    conversionBedGraph[utr.chromosome + ":" + str(utr.start + position) + ":" + str(utr.strand)] = tcRateUtr[position]
                if(coverageUtr[position] > 0):
                    coveredPositions += 1

            if(coveredTcount > 0):
                avgConversationRate = avgConversationRate / coveredTcount
            else:
                avgConversationRate = 0

            # reads per million mapped to the UTR
            readsCPM = 0
            if(readNumber > 0):
                readsCPM = readCount  * 1000000.0 / readNumber

            # Convert to SlamSeqInterval and print
            conversionRate = 0
            if (coverageOnTs > 0) :
                conversionRate = float(conversionsOnTs) / float(coverageOnTs)
            slamSeqUtr = SlamSeqInterval(utr.chromosome, utr.start, utr.stop, utr.strand, utr.name, Tcontent, readsCPM, coverageOnTs, conversionsOnTs, conversionRate, readCount, tcReadCount, multiMapCount)

        print(slamSeqUtr, file=fileCSV)

    fileCSV.close()

    fileTest.close()

    fileBedgraphPlus = open(outputBedgraphPlus,'w')
    fileBedgraphMinus = open(outputBedgraphMinus,'w')

    for position in conversionBedGraph:
        positionData = position.split(":")
        if(positionData[2] == "+"):
            print(positionData[0], positionData[1], int(positionData[1]) + 1, conversionBedGraph[position], file=fileBedgraphPlus)
        else:
            print(positionData[0], positionData[1], int(positionData[1]) + 1, conversionBedGraph[position], file=fileBedgraphMinus)

    fileBedgraphPlus.close()
    fileBedgraphMinus.close()

    #fileNameMLE = replaceExtension(outputCSV, ".tsv", "_mle")
    #callR(getPlotter("compute_conversion_rate_mle") +  " -f " + fileNameTest + " -r " + "0.024" + " -o " + fileNameMLE)

<commit_msg>MLE count code removed<commit_after>#!/usr/bin/env python

from __future__ import print_function

import sys
import pysam
import os

from slamdunk.utils.misc import replaceExtension, getSampleInfo, SlamSeqInfo, md5, callR, getPlotter  # @UnresolvedImport
from slamdunk.utils.BedReader import BedIterator  # @UnresolvedImport

from slamdunk.utils import SNPtools  # @UnresolvedImport
from slamdunk.slamseq.SlamSeqFile import SlamSeqBamFile, ReadDirection, SlamSeqInterval  # @UnresolvedImport

from slamdunk.version import __version__, __bam_version__, __count_version__  # @UnresolvedImport

def collapse(expandedCSV, collapsedCSV, log):

    tcDict = {}

    outCSV = open(collapsedCSV, 'w')

    readNumber = 0

    with open(expandedCSV, 'r') as f:

        # Skip header
#         next(f)

        for line in f:

            fields = line.rstrip().split('\t')
            if (len(fields) == 14) :

                gene = fields[3]
                length = fields[4]
                Tcontent = fields[8]
                coverageOnTs = fields[9]
                conversionsOnTs = fields[10]
                readCount = fields[11]
                tcReadCount = fields[12]
                multimapCount = fields[13]

                if (gene in tcDict.keys()) :
                    tcDict[gene]['length'] += max(int(length),0)
                    tcDict[gene]['Tcontent'] += max(int(Tcontent),0)
                    tcDict[gene]['coverageOnTs'] += max(int(coverageOnTs),0)
                    tcDict[gene]['conversionsOnTs'] += max(int(conversionsOnTs),0)
                    tcDict[gene]['readCount'] += max(int(readCount),0)
                    tcDict[gene]['tcReadCount'] += max(int(tcReadCount),0)
                    tcDict[gene]['multimapCount'] += max(int(multimapCount),0)
                else :
                    tcDict[gene] = {}
                    tcDict[gene]['length'] = max(int(length),0)
                    tcDict[gene]['Tcontent'] = max(int(Tcontent),0)
                    tcDict[gene]['coverageOnTs'] = max(int(coverageOnTs),0)
                    tcDict[gene]['conversionsOnTs'] = max(int(conversionsOnTs),0)
                    tcDict[gene]['readCount'] = max(int(readCount),0)
                    tcDict[gene]['tcReadCount'] = max(int(tcReadCount),0)
                    tcDict[gene]['multimapCount'] = max(int(multimapCount),0)

                readNumber += int(readCount)

        else :
            print("Error in TC file format - unexpected number of fields (" + str(len(fields)) + ") in the following line:\n" + line, file=log)

    print("gene_name", "length", "readsCPM", "conversionRate", "Tcontent", "coverageOnTs", "conversionsOnTs", "readCount", "tcReadCount", "multimapCount", sep='\t', file=outCSV)

    for gene in sorted(tcDict.keys()) :

        print(gene,end="\t",file=outCSV)
        print(tcDict[gene]['length'],end="\t",file=outCSV)
        print(float(tcDict[gene]['readCount']) / float(readNumber) * 1000000,end="\t",file=outCSV)
        conversionRate = 0
        if (tcDict[gene]['coverageOnTs'] > 0) :
            conversionRate = float(tcDict[gene]['conversionsOnTs']) / tcDict[gene]['coverageOnTs']
        print(conversionRate,end="\t",file=outCSV)
        print(tcDict[gene]['Tcontent'],end="\t",file=outCSV)
        print(tcDict[gene]['coverageOnTs'],end="\t",file=outCSV)
        print(tcDict[gene]['conversionsOnTs'],end="\t",file=outCSV)
        print(tcDict[gene]['readCount'],end="\t",file=outCSV)
        print(tcDict[gene]['tcReadCount'],end="\t",file=outCSV)
        print(tcDict[gene]['multimapCount'],file=outCSV)

    outCSV.close()

def getMean(values, skipZeros=True):
    count = 0.0
    totalSum = 0.0
    for value in values:
        if not skipZeros or value > 0:
            totalSum = totalSum + value
            count += 1
    if(count > 0):
        return totalSum / count
    else:
        return 0.0

def computeTconversions(ref, bed, snpsFile, bam, maxReadLength, minQual, outputCSV, outputBedgraphPlus, outputBedgraphMinus, strictTCs, log):

    referenceFile = pysam.FastaFile(ref)

    sampleInfo = getSampleInfo(bam)

    slamseqInfo = SlamSeqInfo(bam)
    readNumber = slamseqInfo.MappedReads

    bedMD5 = md5(bed)

    fileCSV = open(outputCSV,'w')
    print("#slamdunk v" + __version__, __count_version__, "sample info:", sampleInfo.Name, sampleInfo.ID, sampleInfo.Type, sampleInfo.Time, sep="\t", file=fileCSV)
    print("#annotation:", os.path.basename(bed), bedMD5, sep="\t", file=fileCSV)
    print(SlamSeqInterval.Header, file=fileCSV)

    snps = SNPtools.SNPDictionary(snpsFile)
    snps.read()

    #Go through one chr after the other
    testFile = SlamSeqBamFile(bam, ref, snps)
    if not testFile.bamVersion == __bam_version__:
        raise RuntimeError("Wrong filtered BAM file version detected (" + testFile.bamVersion + "). Expected version " + __bam_version__ + ". Please rerun slamdunk filter.")

    bedMD5 = md5(bed)
    if slamseqInfo.AnnotationMD5 != bedMD5:
        print("Warning: MD5 checksum of annotation (" + bedMD5 + ") does not matched MD5 in filtered BAM files (" + slamseqInfo.AnnotationMD5 + "). Most probably the annotation filed changed after the filtered BAM files were created.", file=log)

    conversionBedGraph = {}

    for utr in BedIterator(bed):
        Tcontent = 0
        slamSeqUtr = SlamSeqInterval(utr.chromosome, utr.start, utr.stop, utr.strand, utr.name, Tcontent, 0, -1, 0, 0, 0, 0, 0)
        if(not utr.hasStrand()):
            raise RuntimeError("Input BED file does not contain stranded intervals.")

        if utr.start < 0:
            raise RuntimeError("Negativ start coordinate found. Please check the following entry in your BED file: " + utr)
        # Retreive reference sequence
        region = utr.chromosome + ":" + str(utr.start + 1) + "-" + str(utr.stop)

        if(utr.chromosome in list(referenceFile.references)):
            #print(refRegion,file=sys.stderr)
            refSeq = referenceFile.fetch(region=region).upper()

            if (utr.strand == "-") :
                #refSeq = complement(refSeq[::-1])
                Tcontent = refSeq.count("A")
            else :
                Tcontent = refSeq.count("T")

            slamSeqUtr._Tcontent = Tcontent

        readIterator = testFile.readInRegion(utr.chromosome, utr.start, utr.stop, utr.strand, maxReadLength)

        tcCountUtr = [0] * utr.getLength()
        coverageUtr = [0] * utr.getLength()

        countFwd = 0
        tcCountFwd = 0
        countRev = 0
        tCountRev = 0

        multiMapFwd = 0
        multiMapRev = 0

        for read in readIterator:

            # Overwrite any conversions for non-TC reads (reads with < 2 TC conversions)
            if (not read.isTcRead and strictTCs) :
                read.tcCount = 0
                read.mismatches = []
                read.conversionRates = 0.0
                read.tcRate = 0.0

            if(read.direction == ReadDirection.Reverse):
                countRev += 1
                if read.tcCount > 0:
                    tCountRev += 1
                if read.isMultimapper:
                    multiMapRev += 1
            else:
                countFwd += 1
                if read.tcCount > 0:
                    tcCountFwd += 1
                if read.isMultimapper:
                    multiMapFwd += 1

            for mismatch in read.mismatches:
                if(mismatch.isTCMismatch(read.direction == ReadDirection.Reverse) and mismatch.referencePosition >= 0 and mismatch.referencePosition < utr.getLength()):
                    tcCountUtr[mismatch.referencePosition] += 1

            for i in xrange(read.startRefPos, read.endRefPos):
                if(i >= 0 and i < utr.getLength()):
                    coverageUtr[i] += 1

        if((utr.strand == "+" and countFwd > 0) or (utr.strand == "-" and countRev > 0)):
            tcRateUtr = [ x * 100.0 / y if y > 0 else 0 for x, y in zip(tcCountUtr, coverageUtr)]

            readCount = countFwd
            tcReadCount = tcCountFwd
            multiMapCount = multiMapFwd

            if(utr.strand == "-"):
                readCount = countRev
                tcReadCount = tCountRev
                multiMapCount = multiMapRev

            if((utr.strand == "-" and countFwd > countRev) or (utr.strand == "+" and countRev > countFwd)):
                print("Warning: " + utr.name + " is located on the " + utr.strand + " strand but read counts are higher for the opposite strand (fwd: " + countFwd + ", rev: " + countRev + ")", file=sys.stderr)

            refSeq = readIterator.getRefSeq()

            # Get number of covered Ts/As in the UTR and compute average conversion rate for all covered Ts/As
            coveredTcount = 0
            avgConversationRate = 0
            coveredPositions = 0
            # Get number of reads on T positions and number of reads with T->C conversions on T positions
            coverageOnTs = 0
            conversionsOnTs = 0

            for position in xrange(0, len(coverageUtr)):

                if(coverageUtr[position] > 0 and ((utr.strand == "+" and refSeq[position] == "T") or (utr.strand == "-" and refSeq[position] == "A"))):
                    coveredTcount += 1
                    avgConversationRate += tcRateUtr[position]

                    coverageOnTs += coverageUtr[position]
                    conversionsOnTs += tcCountUtr[position]
                    conversionBedGraph[utr.chromosome + ":" + str(utr.start + position) + ":" + str(utr.strand)] = tcRateUtr[position]
                if(coverageUtr[position] > 0):
                    coveredPositions += 1

            if(coveredTcount > 0):
                avgConversationRate = avgConversationRate / coveredTcount
            else:
                avgConversationRate = 0

            # reads per million mapped to the UTR
            readsCPM = 0
            if(readNumber > 0):
                readsCPM = readCount  * 1000000.0 / readNumber

            # Convert to SlamSeqInterval and print
            conversionRate = 0
            if (coverageOnTs > 0) :
                conversionRate = float(conversionsOnTs) / float(coverageOnTs)
            slamSeqUtr = SlamSeqInterval(utr.chromosome, utr.start, utr.stop, utr.strand, utr.name, Tcontent, readsCPM, coverageOnTs, conversionsOnTs, conversionRate, readCount, tcReadCount, multiMapCount)

        print(slamSeqUtr, file=fileCSV)

    fileCSV.close()

    fileBedgraphPlus = open(outputBedgraphPlus,'w')
    fileBedgraphMinus = open(outputBedgraphMinus,'w')

    for position in conversionBedGraph:
        positionData = position.split(":")
        if(positionData[2] == "+"):
            print(positionData[0], positionData[1], int(positionData[1]) + 1, conversionBedGraph[position], file=fileBedgraphPlus)
        else:
            print(positionData[0], positionData[1], int(positionData[1]) + 1, conversionBedGraph[position], file=fileBedgraphMinus)

    fileBedgraphPlus.close()
    fileBedgraphMinus.close()

<|endoftext|>


<commit_before>import numpy as np
import os
import pkg_resources

from ctypes.util import find_library
from cffi import FFI

PKG_ROOT = os.path.dirname(__file__)

__ffi = FFI()
with open(os.path.join(PKG_ROOT, 'rust_lib', 'hydromath.h')) as header:
    __ffi.cdef(header.read())

__lib = __ffi.dlopen(os.path.join(PKG_ROOT, 'hydromath.so'))

mse_c = __lib.mse_c

def mse(obs, sim):
    assert len(obs) == len(sim)

    return mse_c(
        __ffi.cast('double *', obs.ctypes.data),
        __ffi.cast('double *', sim.ctypes.data),
        len(obs)
    )
<commit_msg>Expose RMSE and NSE functions in addition to MSE<commit_after>import numpy as np
import os
import pkg_resources

from ctypes.util import find_library
from cffi import FFI

PKG_ROOT = os.path.dirname(__file__)

__ffi = FFI()
with open(os.path.join(PKG_ROOT, 'rust_lib', 'hydromath.h')) as header:
    __ffi.cdef(header.read())

__lib = __ffi.dlopen(os.path.join(PKG_ROOT, 'hydromath.so'))

mse_c = __lib.mse_c

def __map_func(c_func_name, obs, sim):
    """
        Map to a Rust function via the C interface.

        Since all the skill score functions take two arrays (and a length) we
        can avoid repeating code by following the required pattern in one
        place. This function looks up the c/rust function name based on the
        passed argument then passes the converted data into the function. Later
        versions can also handle ensuring the passed arrays are numpy arrays of
        the correct data type.
    """
    assert len(obs) == len(sim)
    return __lib.__getattr__(c_func_name)(
        __ffi.cast('double *', obs.ctypes.data),
        __ffi.cast('double *', sim.ctypes.data),
        len(obs)
    )

def mse(obs, sim):
    return __map_func('mse_c', obs, sim)

def rmse(obs, sim):
    return __map_func('rmse_c', obs, sim)

def nse(obs, sim):
    return __map_func('nse_c', obs, sim)
<|endoftext|>


<commit_before>import json
from collections import defaultdict
from queue import Queue

from django.contrib import messages
from django.core.paginator import Paginator, PageNotAnInteger, EmptyPage
from django.http import HttpResponseRedirect
from django.shortcuts import render, get_object_or_404

from smartshark.datacollection.pluginmanagementinterface import PluginManagementInterface
from smartshark.filters import JobExecutionFilter
from smartshark.models import PluginExecution, Job, Project, Plugin
from smartshark.mongohandler import handler

def index(request):
    return render(request, 'smartshark/frontend/index.html')

class Item(object):
    def __init__(self, id, name, desc=None, sub_fields=None, parent='#', logical_types=None, reference_to=None):
        self.id = id
        self.name = name
        self.reference_to = reference_to

        if desc is None:
            self.desc = []
        else:
            self.desc = [desc]

        if sub_fields is None:
            self.sub_fields = []

        if logical_types is None:
            self.logical_types = []
        else:
            self.logical_types = logical_types

        self.parent=parent

    def add_field(self, field):
        self.sub_fields.append(field)

    def add_description(self, description):
        self.desc.append(description)

def recursion(item, parent, plugin_name, items):
    if 'fields' not in item:
        return

    for field in item['fields']:
        if isinstance(field['logical_type'], list):
            logical_types = field['logical_type']
        else:
            logical_types = [field['logical_type']]

        reference_to = None
        if 'reference_to' in field:
            reference_to = field['reference_to']

        new_field = Item(
            parent+'_'+field['field_name'], field['field_name'], desc={'desc': field['desc'], 'plugin': ' '.join(plugin_name.split('_'))},
            parent=parent, logical_types=logical_types, reference_to=reference_to)

        if new_field.id in items:
            stored_field = items[new_field.id]
            stored_field.add_description({'desc': field['desc'], 'plugin': ' '.join(plugin_name.split('_'))})
        else:
            items[new_field.id] = new_field

        if 'fields' in field:
            recursion(field, new_field.id, plugin_name, items)

def documentation(request):
    items = {}
    data = []

    active_plugins = []
    for p in Plugin.objects.all().filter(active=True):
        active_plugins.append(p.name + '_' + p.version)

    for schema in handler.get_plugin_schemas():

        plugin_name = schema['plugin']

        # only active plugins as the json gets too big otherwise
        if plugin_name not in active_plugins:
            continue

        for mongo_collection in schema['collections']:

            desc = ''
            if 'desc' in mongo_collection.keys():
                desc = mongo_collection['desc']

            collection_name = ''
            if 'collection_name' in mongo_collection.keys():
                collection_name = mongo_collection['collection_name']

            collection = Item(collection_name, collection_name, desc={'desc': desc, 'plugin': ' '.join(plugin_name.split('_'))})
            if collection.id in items:
                stored_collection = items[collection.id]
                stored_collection.add_description({'desc': desc, 'plugin': ' '.join(plugin_name.split('_'))})
            else:
                items[collection.id] = collection

            recursion(mongo_collection, collection_name, plugin_name, items)

    for item_id, item_data in items.items():
        json_collection = {
            'id': item_id,
            'parent': item_data.parent,
            'text': item_data.name,
            'data': {
                'desc': item_data.desc,
                'logical_types': item_data.logical_types,
                'reference_to': item_data.reference_to
            }
        }
        data.append(json_collection)

    return render(request, 'smartshark/frontend/documentation.html', {
        'data': json.dumps(data),
        'plugins': Plugin.objects.all().filter(active=True).order_by('name'),
    })

def plugin_execution_status(request, id):
    if not request.user.is_authenticated() or not request.user.has_perm('smartshark.plugin_execution_status'):
        messages.error(request, 'You are not authorized to perform this action.')
        return HttpResponseRedirect('/admin/smartshark/project')

    interface = PluginManagementInterface.find_correct_plugin_manager()
    plugin_execution = get_object_or_404(PluginExecution, pk=id)

    # Get all jobs from all plugin_executions which did not terminate yet
    # Update the job stati for these jobs
    jobs = Job.objects.filter(plugin_execution=plugin_execution, status='WAIT').all()
    job_stati = interface.get_job_stati(jobs)
    i = 0
    for job in jobs:
        job.status = job_stati[i]
        job.save()
        i += 1

    job_filter = JobExecutionFilter(request.GET, queryset=Job.objects.all().filter(plugin_execution=plugin_execution))

    rev = [exitjob.revision_hash if exitjob.revision_hash else '' for exitjob in job_filter.qs.filter(status='EXIT')]
    exit_job_revisions = ''
    if rev:
        exit_job_revisions = ','.join(rev)

    # Set up pagination
    paginator = Paginator(job_filter.qs, 10)
    page = request.GET.get('page')
    try:
        jobs = paginator.page(page)
    except PageNotAnInteger:
        # If page is not an integer, deliver first page.
        jobs = paginator.page(1)
    except EmptyPage:
        # If page is out of range (e.g. 9999), deliver last page of results.
        jobs = paginator.page(paginator.num_pages)

    return render(request, 'smartshark/project/plugin_execution_status.html', {
        'plugin_execution': plugin_execution,
        'filter': job_filter,
        'jobs': jobs,
        'overall': len(job_filter.qs),
        'queried_status': job_filter.data.get('status', None),
        'done_jobs': len(job_filter.qs.filter(status='DONE')),
        'exit_jobs': len(rev),
        'waiting_jobs': len(job_filter.qs.filter(status='WAIT')),
        'exit_job_revisions': exit_job_revisions
    })

def plugin_status(request, id):
    if not request.user.is_authenticated() or not request.user.has_perm('smartshark.plugin_status'):
        messages.error(request, 'You are not authorized to perform this action.')
        return HttpResponseRedirect('/admin/smartshark/project')

    project = get_object_or_404(Project, pk=id)
    execution_list = PluginExecution.objects.all().filter(project=project).order_by('-submitted_at')

    paginator = Paginator(execution_list, 10)
    page = request.GET.get('page')
    try:
        executions = paginator.page(page)
    except PageNotAnInteger:
        # If page is not an integer, deliver first page.
        executions = paginator.page(1)
    except EmptyPage:
        # If page is out of range (e.g. 9999), deliver last page of results.
        executions = paginator.page(paginator.num_pages)

    return render(request, 'smartshark/project/plugin_status.html', {
        'executions': executions,
        'project': project
    })

def job_output(request, id, type):
    if not request.user.is_authenticated() or not request.user.has_perm('smartshark.job_output'):
        messages.error(request, 'You are not authorized to perform this action.')
        return HttpResponseRedirect('/admin/smartshark/project')

    job = get_object_or_404(Job, pk=id)
    interface = PluginManagementInterface.find_correct_plugin_manager()

    if type == 'output':
        output = interface.get_output_log(job)
    elif type == 'error':
        output = interface.get_error_log(job)
    elif type == 'arguments':
        return render(request, 'smartshark/job/execution_arguments.html', {
            'exe_arguments': job.plugin_execution.executionhistory_set.all().order_by('execution_argument__position'),
            'cmd': interface.get_sent_bash_command(job),
        })

    return render(request, 'smartshark/job/output.html', {
        'output': '\n'.join(output),
        'job': job,
    })
<commit_msg>we need all plugins not just active ones for documentation<commit_after>import json
from collections import defaultdict
from queue import Queue

from django.contrib import messages
from django.core.paginator import Paginator, PageNotAnInteger, EmptyPage
from django.http import HttpResponseRedirect
from django.shortcuts import render, get_object_or_404

from smartshark.datacollection.pluginmanagementinterface import PluginManagementInterface
from smartshark.filters import JobExecutionFilter
from smartshark.models import PluginExecution, Job, Project, Plugin
from smartshark.mongohandler import handler

def index(request):
    return render(request, 'smartshark/frontend/index.html')

class Item(object):
    def __init__(self, id, name, desc=None, sub_fields=None, parent='#', logical_types=None, reference_to=None):
        self.id = id
        self.name = name
        self.reference_to = reference_to

        if desc is None:
            self.desc = []
        else:
            self.desc = [desc]

        if sub_fields is None:
            self.sub_fields = []

        if logical_types is None:
            self.logical_types = []
        else:
            self.logical_types = logical_types

        self.parent=parent

    def add_field(self, field):
        self.sub_fields.append(field)

    def add_description(self, description):
        self.desc.append(description)

def recursion(item, parent, plugin_name, items):
    if 'fields' not in item:
        return

    for field in item['fields']:
        if isinstance(field['logical_type'], list):
            logical_types = field['logical_type']
        else:
            logical_types = [field['logical_type']]

        reference_to = None
        if 'reference_to' in field:
            reference_to = field['reference_to']

        new_field = Item(
            parent+'_'+field['field_name'], field['field_name'], desc={'desc': field['desc'], 'plugin': ' '.join(plugin_name.split('_'))},
            parent=parent, logical_types=logical_types, reference_to=reference_to)

        if new_field.id in items:
            stored_field = items[new_field.id]
            stored_field.add_description({'desc': field['desc'], 'plugin': ' '.join(plugin_name.split('_'))})
        else:
            items[new_field.id] = new_field

        if 'fields' in field:
            recursion(field, new_field.id, plugin_name, items)

def documentation(request):
    items = {}
    data = []

    for schema in handler.get_plugin_schemas():

        plugin_name = schema['plugin']

        for mongo_collection in schema['collections']:

            desc = ''
            if 'desc' in mongo_collection.keys():
                desc = mongo_collection['desc']

            collection_name = ''
            if 'collection_name' in mongo_collection.keys():
                collection_name = mongo_collection['collection_name']

            collection = Item(collection_name, collection_name, desc={'desc': desc, 'plugin': ' '.join(plugin_name.split('_'))})
            if collection.id in items:
                stored_collection = items[collection.id]
                stored_collection.add_description({'desc': desc, 'plugin': ' '.join(plugin_name.split('_'))})
            else:
                items[collection.id] = collection

            recursion(mongo_collection, collection_name, plugin_name, items)

    for item_id, item_data in items.items():
        json_collection = {
            'id': item_id,
            'parent': item_data.parent,
            'text': item_data.name,
            'data': {
                'desc': item_data.desc,
                'logical_types': item_data.logical_types,
                'reference_to': item_data.reference_to
            }
        }
        data.append(json_collection)

    return render(request, 'smartshark/frontend/documentation.html', {
        'data': json.dumps(data),
        'plugins': Plugin.objects.all().filter(active=True).order_by('name'),
    })

def plugin_execution_status(request, id):
    if not request.user.is_authenticated() or not request.user.has_perm('smartshark.plugin_execution_status'):
        messages.error(request, 'You are not authorized to perform this action.')
        return HttpResponseRedirect('/admin/smartshark/project')

    interface = PluginManagementInterface.find_correct_plugin_manager()
    plugin_execution = get_object_or_404(PluginExecution, pk=id)

    # Get all jobs from all plugin_executions which did not terminate yet
    # Update the job stati for these jobs
    jobs = Job.objects.filter(plugin_execution=plugin_execution, status='WAIT').all()
    job_stati = interface.get_job_stati(jobs)
    i = 0
    for job in jobs:
        job.status = job_stati[i]
        job.save()
        i += 1

    job_filter = JobExecutionFilter(request.GET, queryset=Job.objects.all().filter(plugin_execution=plugin_execution))

    rev = [exitjob.revision_hash if exitjob.revision_hash else '' for exitjob in job_filter.qs.filter(status='EXIT')]
    exit_job_revisions = ''
    if rev:
        exit_job_revisions = ','.join(rev)

    # Set up pagination
    paginator = Paginator(job_filter.qs, 10)
    page = request.GET.get('page')
    try:
        jobs = paginator.page(page)
    except PageNotAnInteger:
        # If page is not an integer, deliver first page.
        jobs = paginator.page(1)
    except EmptyPage:
        # If page is out of range (e.g. 9999), deliver last page of results.
        jobs = paginator.page(paginator.num_pages)

    return render(request, 'smartshark/project/plugin_execution_status.html', {
        'plugin_execution': plugin_execution,
        'filter': job_filter,
        'jobs': jobs,
        'overall': len(job_filter.qs),
        'queried_status': job_filter.data.get('status', None),
        'done_jobs': len(job_filter.qs.filter(status='DONE')),
        'exit_jobs': len(rev),
        'waiting_jobs': len(job_filter.qs.filter(status='WAIT')),
        'exit_job_revisions': exit_job_revisions
    })

def plugin_status(request, id):
    if not request.user.is_authenticated() or not request.user.has_perm('smartshark.plugin_status'):
        messages.error(request, 'You are not authorized to perform this action.')
        return HttpResponseRedirect('/admin/smartshark/project')

    project = get_object_or_404(Project, pk=id)
    execution_list = PluginExecution.objects.all().filter(project=project).order_by('-submitted_at')

    paginator = Paginator(execution_list, 10)
    page = request.GET.get('page')
    try:
        executions = paginator.page(page)
    except PageNotAnInteger:
        # If page is not an integer, deliver first page.
        executions = paginator.page(1)
    except EmptyPage:
        # If page is out of range (e.g. 9999), deliver last page of results.
        executions = paginator.page(paginator.num_pages)

    return render(request, 'smartshark/project/plugin_status.html', {
        'executions': executions,
        'project': project
    })

def job_output(request, id, type):
    if not request.user.is_authenticated() or not request.user.has_perm('smartshark.job_output'):
        messages.error(request, 'You are not authorized to perform this action.')
        return HttpResponseRedirect('/admin/smartshark/project')

    job = get_object_or_404(Job, pk=id)
    interface = PluginManagementInterface.find_correct_plugin_manager()

    if type == 'output':
        output = interface.get_output_log(job)
    elif type == 'error':
        output = interface.get_error_log(job)
    elif type == 'arguments':
        return render(request, 'smartshark/job/execution_arguments.html', {
            'exe_arguments': job.plugin_execution.executionhistory_set.all().order_by('execution_argument__position'),
            'cmd': interface.get_sent_bash_command(job),
        })

    return render(request, 'smartshark/job/output.html', {
        'output': '\n'.join(output),
        'job': job,
    })
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# filename: reply.py
import time
from wenpl.divide import showReply

class Msg(object):
    def __init__(self):
        pass

    def send(self):
        return "success"

class TextMsg(Msg):
    def __init__(self, toUserName, fromUserName, content):
        self.__dict = dict()
        self.__dict['ToUserName'] = toUserName
        self.__dict['FromUserName'] = fromUserName
        self.__dict['CreateTime'] = int(time.time())
        self.__dict['Content'] = showReply(content)

    def send(self):
        XmlForm = """
        <xml>
        <ToUserName><![CDATA[{ToUserName}]]></ToUserName>
        <FromUserName><![CDATA[{FromUserName}]]></FromUserName>
        <CreateTime>{CreateTime}</CreateTime>
        <MsgType><![CDATA[text]]></MsgType>
        <Content><![CDATA[{Content}]]></Content>
        </xml>
        """
        return XmlForm.format(**self.__dict)

class LocationMsg(Msg):
    def __init__(self, toUserName, fromUserName, content):
        self.__dict = dict()
        self.__dict['ToUserName'] = toUserName
        self.__dict['FromUserName'] = fromUserName
        self.__dict['CreateTime'] = int(time.time())
        # self.__dict['CreateTime'] = str(time.time())
        # self.__dict['Content'] = show_reply(content)
        # print showReply(content)

    def send(self):
        print "haha"
        XmlForm = """
        <xml>
        <ToUserName><![CDATA[{ToUserName}]]></ToUserName>
        <FromUserName><![CDATA[{FromUserName}]]></FromUserName>
        <CreateTime>{CreateTime}</CreateTime>
        <MsgType><![CDATA[location]]></MsgType>
        <Location_X>23.134521</Location_X>
        <Location_Y>113.358803</Location_Y>
        <![CDATA[位置信息]]>
        </xml>
        """
        # <Scale>20</Scale>
        # < MsgId > 5832828233808572154 < / MsgId >
        print "haha1", XmlForm
        return XmlForm.format(**self.__dict)

class ImageMsg(Msg):
    def __init__(self, toUserName, fromUserName, mediaId):
        self.__dict = dict()
        self.__dict['ToUserName'] = toUserName
        self.__dict['FromUserName'] = fromUserName
        self.__dict['CreateTime'] = int(time.time())
        self.__dict['MediaId'] = mediaId
    def send(self):
        XmlForm = """
        <xml>
        <ToUserName><![CDATA[{ToUserName}]]></ToUserName>
        <FromUserName><![CDATA[{FromUserName}]]></FromUserName>
        <CreateTime>{CreateTime}</CreateTime>
        <MsgType><![CDATA[image]]></MsgType>
        <Image>
        <MediaId><![CDATA[{MediaId}]]></MediaId>
        </Image>
        </xml>
        """
        return XmlForm.format(**self.__dict)

class VoiceMsg(Msg):
    def __init__(self, toUserName, fromUserName, content):
        self.__dict = dict()
        self.__dict['ToUserName'] = toUserName
        self.__dict['FromUserName'] = fromUserName
        self.__dict['CreateTime'] = int(time.time())
        self.__dict['Content'] = "语音: "+content+'\n' +'\n'+showReply(content)
        # print showReply(content)

    def send(self):
        XmlForm = """
        <xml>
        <ToUserName><![CDATA[{ToUserName}]]></ToUserName>
        <FromUserName><![CDATA[{FromUserName}]]></FromUserName>
        <CreateTime>{CreateTime}</CreateTime>
        <MsgType><![CDATA[text]]></MsgType>
        <Content><![CDATA[{Content}]]></Content>
        </xml>
        """
        return XmlForm.format(**self.__dict)
<commit_msg>last updated, project has been removed and replaced in a private project<commit_after># -*- coding: utf-8 -*-
# filename: reply.py
import time
from wenpl.divide import showReply

class Msg(object):
    def __init__(self):
        pass

    def send(self):
        return "success"

class TextMsg(Msg):
    def __init__(self, toUserName, fromUserName, content):
        self.__dict = dict()
        self.__dict['ToUserName'] = toUserName
        self.__dict['FromUserName'] = fromUserName
        self.__dict['CreateTime'] = int(time.time())
        self.__dict['Content'] = showReply(content)

    def send(self):
        XmlForm = """
        <xml>
        <ToUserName><![CDATA[{ToUserName}]]></ToUserName>
        <FromUserName><![CDATA[{FromUserName}]]></FromUserName>
        <CreateTime>{CreateTime}</CreateTime>
        <MsgType><![CDATA[text]]></MsgType>
        <Content><![CDATA[{Content}]]></Content>
        </xml>
        """
        return XmlForm.format(**self.__dict)

# class LocationMsg(Msg):
#     def __init__(self, toUserName, fromUserName, content):
#         self.__dict = dict()
#         self.__dict['ToUserName'] = toUserName
#         self.__dict['FromUserName'] = fromUserName
#         self.__dict['CreateTime'] = int(time.time())
#         # self.__dict['CreateTime'] = str(time.time())
#         # self.__dict['Content'] = show_reply(content)
#         # print showReply(content)
#
#     def send(self):
#         XmlForm = """
#         <xml>
#         <ToUserName><![CDATA[{ToUserName}]]></ToUserName>
#         <FromUserName><![CDATA[{FromUserName}]]></FromUserName>
#         <CreateTime>{CreateTime}</CreateTime>
#         <MsgType><![CDATA[location]]></MsgType>
#         <Location_X>23.134521</Location_X>
#         <Location_Y>113.358803</Location_Y>
#         <![CDATA[位置信息]]>
#         </xml>
#         """
#         # <Scale>20</Scale>
#         # < MsgId > 5832828233808572154 < / MsgId >
#         return XmlForm.format(**self.__dict)

class ImageMsg(Msg):
    def __init__(self, toUserName, fromUserName, mediaId):
        self.__dict = dict()
        self.__dict['ToUserName'] = toUserName
        self.__dict['FromUserName'] = fromUserName
        self.__dict['CreateTime'] = int(time.time())
        self.__dict['MediaId'] = mediaId
    def send(self):
        XmlForm = """
        <xml>
        <ToUserName><![CDATA[{ToUserName}]]></ToUserName>
        <FromUserName><![CDATA[{FromUserName}]]></FromUserName>
        <CreateTime>{CreateTime}</CreateTime>
        <MsgType><![CDATA[image]]></MsgType>
        <Image>
        <MediaId><![CDATA[{MediaId}]]></MediaId>
        </Image>
        </xml>
        """
        return XmlForm.format(**self.__dict)

class VoiceMsg(Msg):
    def __init__(self, toUserName, fromUserName, content):
        self.__dict = dict()
        self.__dict['ToUserName'] = toUserName
        self.__dict['FromUserName'] = fromUserName
        self.__dict['CreateTime'] = int(time.time())
        self.__dict['Content'] = "语音: "+content+'\n' +'\n'+showReply(content)
        # print showReply(content)

    def send(self):
        XmlForm = """
        <xml>
        <ToUserName><![CDATA[{ToUserName}]]></ToUserName>
        <FromUserName><![CDATA[{FromUserName}]]></FromUserName>
        <CreateTime>{CreateTime}</CreateTime>
        <MsgType><![CDATA[text]]></MsgType>
        <Content><![CDATA[{Content}]]></Content>
        </xml>
        """
        return XmlForm.format(**self.__dict)
<|endoftext|>


<commit_before>import boto3
import argparse
import gnupg

# Assumes you have GPG already installed
# Assumes that the Data Czars already have your public key
# Asumes that .boto is configured with edX Prod account

# Parser
parser = argparse.ArgumentParser(description="Username of Data Czar.")
parser.add_argument('-u', '--user', help='Email of Data Czar', required=True)
parser.add_argument('-f', '--file', help='Public Key file', required=True)
parser.add_argument('--credentials-only', help='Only create new credentials', default=False, action='store_true')
parser.add_argument('-o', '--orgs', nargs='*', help='Name of the org(s) as list, User need to be a member', default=None)
args = parser.parse_args()

# Import Data Czar GPG Key
gpg = gnupg.GPG()
key_data = open(args.file).read()
import_result = gpg.import_keys(key_data)

# Connect to AWS and create account
iam = boto3.client('iam')

if not args.credentials_only:
    user_response = iam.create_user(UserName=args.user)

key_response = iam.create_access_key(UserName=args.user)

# Add user to group edx-s3bucket-course-data-readonly
iam.add_user_to_group(GroupName='edx-s3bucket-course-data-readonly', UserName=args.user)

# Add user to it's respective Org
if args.orgs:
    for org in args.orgs:
        user_org = 'edx-course-data-' + org.lower()
        iam.add_user_to_group(GroupName=user_org, UserName=args.user)

# Create AWS Cred String
key = key_response['AccessKey']
credstring = str(f'AWS_ACCESS_KEY_ID = {key["AccessKeyId"]} \nAWS_SECRET_ACCESS_KEY = {key["SecretAccessKey"]}')

# Encrypt file
encrypted_data = gpg.encrypt(credstring, args.user, always_trust=True)
encrypted_string = str(encrypted_data)
gpgfile = open(args.user + '-credentials.txt.gpg', 'w+')
gpgfile.write(encrypted_string)

print('ok: ', encrypted_data.ok)
print('status: ', encrypted_data.status)
print('stderr: ', encrypted_data.stderr)
<commit_msg>feat: Add user tagging in create_data_czar.py | PSRE-1038 (#6546)<commit_after>import boto3
import argparse
import gnupg

# Assumes you have GPG already installed
# Assumes that the Data Czars already have your public key
# Asumes that .boto is configured with edX Prod account

# Parser
parser = argparse.ArgumentParser(description="Username of Data Czar.")
parser.add_argument('-u', '--user', help='Email of Data Czar', required=True)
parser.add_argument('-f', '--file', help='Public Key file', required=True)
parser.add_argument('--credentials-only', help='Only create new credentials', default=False, action='store_true')
parser.add_argument('-o', '--orgs', nargs='*', help='Name of the org(s) as list, User need to be a member', default=None)
parser.add_argument('-c', '--creator', help='Name of the creator', default=None)
args = parser.parse_args()

# Import Data Czar GPG Key
gpg = gnupg.GPG()
key_data = open(args.file).read()
import_result = gpg.import_keys(key_data)

# Connect to AWS and create account
iam = boto3.client('iam')

if not args.credentials_only:
    user_response = iam.create_user(UserName=args.user)
    if args.creator:
        tag_response = iam.tag_user(UserName=args.user, Tags=[{'Key': 'Creator', 'Value': args.creator}])

key_response = iam.create_access_key(UserName=args.user)

# Add user to group edx-s3bucket-course-data-readonly
iam.add_user_to_group(GroupName='edx-s3bucket-course-data-readonly', UserName=args.user)

# Add user to it's respective Org
if args.orgs:
    for org in args.orgs:
        user_org = 'edx-course-data-' + org.lower()
        iam.add_user_to_group(GroupName=user_org, UserName=args.user)

# Create AWS Cred String
key = key_response['AccessKey']
credstring = str(f'AWS_ACCESS_KEY_ID = {key["AccessKeyId"]} \nAWS_SECRET_ACCESS_KEY = {key["SecretAccessKey"]}')

# Encrypt file
encrypted_data = gpg.encrypt(credstring, args.user, always_trust=True)
encrypted_string = str(encrypted_data)
gpgfile = open(args.user + '-credentials.txt.gpg', 'w+')
gpgfile.write(encrypted_string)

print('ok: ', encrypted_data.ok)
print('status: ', encrypted_data.status)
print('stderr: ', encrypted_data.stderr)
<|endoftext|>


<commit_before>#! /usr/bin/env python2
import os
import unittest

from selenium import webdriver

DIR = os.path.dirname(__file__)

os.environ['PATH'] = os.environ['PATH'] + ":" + DIR

# local testing, remove / comment out before CI
PROXY='127.0.0.1:8080'
POBJ = {
    "httpProxy":PROXY,
    "ftpProxy":PROXY,
    "sslProxy":PROXY,
    "noProxy":[],
    "proxyType":"MANUAL"
}
webdriver.DesiredCapabilities.CHROME['proxy'] = POBJ
webdriver.DesiredCapabilities.FIREFOX['proxy'] = POBJ

class DuckMixin(object):
    def testDuckSafe(self):
        self.browser.get("https://duckduckgo.com")
        (self.browser
         .find_element_by_id("search_form_input_homepage")
         .send_keys("porn" + webdriver.common.keys.Keys.RETURN))
        # click to disable temporarily, then wait, then see if available
        self.browser.implicitly_wait(10)
        self.browser.find_element_by_class_name("js-safe-search-temp").click()
        # assert that browser.find_element_by_class_name("js-safe-search-temp") still exists, else throws exception
        self.browser.find_element_by_class_name("js-safe-search-temp")

class QwantMixin(object):
    def testQwant(self):
        self.browser.get("https://qwant.com")
        self.browser.implicitly_wait(10)
        (self.browser
         .find_element_by_css_selector("input[type=search]")
         .send_keys("porn" + webdriver.common.keys.Keys.RETURN))
        self.browser.find_element_by_class_name("no_result")

class FirefoxTestCase(unittest.TestCase, DuckMixin, QwantMixin):
    def setUp(self):
        profile = webdriver.FirefoxProfile()
        profile.add_extension(extension=os.path.join(DIR, "..", "addon"))
        self.browser = webdriver.Firefox(firefox_profile=profile)

    def tearDown(self):
        self.browser.close()

class ChromiumTestCase(unittest.TestCase, DuckMixin, QwantMixin):
    def setUp(self):
        profile = webdriver.chrome.options.Options()
        profile.add_extension(extension=os.path.join(DIR, "..", "safe.zip"))
        self.browser = webdriver.Chrome(chrome_options=profile)

    def tearDown(self):
        self.browser.close()

if __name__ == "__main__":
    unittest.main()
<commit_msg>test: lacked test for duckduckgo POST variant<commit_after>#! /usr/bin/env python2
import os
import unittest

from selenium import webdriver

DIR = os.path.dirname(__file__)

os.environ['PATH'] = os.environ['PATH'] + ":" + DIR

# local testing, remove / comment out before CI
PROXY='127.0.0.1:8080'
POBJ = {
    "httpProxy":PROXY,
    "ftpProxy":PROXY,
    "sslProxy":PROXY,
    "noProxy":[],
    "proxyType":"MANUAL"
}
webdriver.DesiredCapabilities.CHROME['proxy'] = POBJ
webdriver.DesiredCapabilities.FIREFOX['proxy'] = POBJ

class DuckMixin(object):
    def testDuckSafe(self):
        self.browser.get("https://duckduckgo.com")
        self.search_porn()

    def testDuckPOSTSafe(self):
        self.browser.get("https://duckduckgo.com")
        # failed to change settings by hand,
        # self.browser.get("https://duckduckgo.com/settings")
        # self.browser.find_element_by_id("setting_kg")
        # so just set cookie
        self.browser.add_cookie({"name": "g", "value": "p"})
        self.search_porn()

    def search_porn(self):
        (self.browser
         .find_element_by_id("search_form_input_homepage")
         .send_keys("porn" + webdriver.common.keys.Keys.RETURN))
        # click to disable temporarily, then wait, then see if available
        self.browser.implicitly_wait(10)
        self.browser.find_element_by_class_name("js-safe-search-temp").click()
        # assert that browser.find_element_by_class_name("js-safe-search-temp") still exists, else throws exception
        self.browser.find_element_by_class_name("js-safe-search-temp")

class QwantMixin(object):
    def testQwant(self):
        self.browser.get("https://qwant.com")
        self.browser.implicitly_wait(10)
        (self.browser
         .find_element_by_css_selector("input[type=search]")
         .send_keys("porn" + webdriver.common.keys.Keys.RETURN))
        self.browser.find_element_by_class_name("no_result")

#### run the tests, one case for each browser, with mixins for each site
class FirefoxTestCase(unittest.TestCase, DuckMixin, QwantMixin):
    def setUp(self):
        profile = webdriver.FirefoxProfile()
        profile.add_extension(extension=os.path.join(DIR, "..", "addon"))
        self.browser = webdriver.Firefox(firefox_profile=profile)

    def tearDown(self):
        self.browser.close()

class ChromiumTestCase(unittest.TestCase, DuckMixin, QwantMixin):
    def setUp(self):
        profile = webdriver.chrome.options.Options()
        profile.add_extension(extension=os.path.join(DIR, "..", "safe.zip"))
        self.browser = webdriver.Chrome(chrome_options=profile)

    def tearDown(self):
        self.browser.close()

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


<commit_before>#!/usr/bin/env python
import os
import sys
import traceback

sys.path.insert(0, os.path.abspath('./../'))

if __name__ == "__main__":
    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "test_bootstrap3.settings")

    from django.core.management import execute_from_command_line

    try:
        execute_from_command_line(sys.argv)
    except Exception as exception:
        # since this response is sent to the PSP, catch errors locally
        print('%s while performing request' % (exception.__str__()))
        traceback.print_exc()
        raise exception
<commit_msg>Fixed setting usage<commit_after>#!/usr/bin/env python
import os
import sys
import traceback

sys.path.insert(0, os.path.abspath('./../'))

if __name__ == "__main__":
    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "bootstrap3.settings")

    from django.core.management import execute_from_command_line

    try:
        execute_from_command_line(sys.argv)
    except Exception as exception:
        # since this response is sent to the PSP, catch errors locally
        print('%s while performing request' % (exception.__str__()))
        traceback.print_exc()
        raise exception
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import os
import sys

root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(root + '/python')

import pandas_ta as ta
import pandas as pd
import ccxt

# -----------------------------------------------------------------------------

print('CCXT Version:', ccxt.__version__)

# -----------------------------------------------------------------------------

exchange = ccxt.binance()
symbol = 'BTC/USDT'
timeframe = '1m'
limit = 50
while True:
    try:
        ohlcv = exchange.fetch_ohlcv(symbol, timeframe)
        print('--------------------------------------------------------------')
        if len(ohlcv):
            df = pd.DataFrame(ohlcv, columns=['time', 'open', 'high', 'low', 'close', 'volume'])
            df['time'] = pd.to_datetime(df['time'], unit='ms')
            df = pd.concat([df, df.ta.rsi()], axis=1)
            print(df[-20:])
            print(exchange.iso8601 (exchange.milliseconds()))
    except Exception as e:
        print(type(e).__name__, str(e))
<commit_msg>rsi.py config length #12309<commit_after># -*- coding: utf-8 -*-

import os
import sys

root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(root + '/python')

import pandas_ta as ta
import pandas as pd
import ccxt

# -----------------------------------------------------------------------------

print('CCXT Version:', ccxt.__version__)

# -----------------------------------------------------------------------------

exchange = ccxt.binance()
symbol = 'BTC/USDT'
timeframe = '1m'
limit = 500
rsi_length = 100
while True:
    try:
        ohlcv = exchange.fetch_ohlcv(symbol, timeframe)
        print('--------------------------------------------------------------')
        if len(ohlcv):
            df = pd.DataFrame(ohlcv, columns=['time', 'open', 'high', 'low', 'close', 'volume'])
            df['time'] = pd.to_datetime(df['time'], unit='ms')
            df = pd.concat([df, df.ta.rsi(length=rsi_length)], axis=1)
            print(df[-20:])
            print(exchange.iso8601 (exchange.milliseconds()))
    except Exception as e:
        print(type(e).__name__, str(e))

<|endoftext|>


<commit_before># (C) Datadog, Inc. 2018
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
import json
import re
import shlex
import xml.parsers.expat  # python 2.4 compatible
from collections import defaultdict
from distutils.version import LooseVersion
from os import geteuid

from six import PY3, iteritems
from six.moves import filter

from datadog_checks.checks import AgentCheck
from datadog_checks.utils.subprocess_output import get_subprocess_output

if PY3:
    long = int

class BackendStatus(object):
    HEALTHY = 'healthy'
    SICK = 'sick'
    ALL = (HEALTHY, SICK)

    @classmethod
    def to_check_status(cls, status):
        if status == cls.HEALTHY:
            return AgentCheck.OK
        elif status == cls.SICK:
            return AgentCheck.CRITICAL
        return AgentCheck.UNKNOWN

class Varnish(AgentCheck):
    SERVICE_CHECK_NAME = 'varnish.backend_healthy'
    # Parse metrics from varnishstat.
    VARNISHSTAT_FORMAT_OPTION = {
        "text": "-1",  # version < 3.0.0
        "xml": "-x",  # version >= 3.0.0 < 5.0.0
        "json": "-j",  # version >=5.0.0
    }

    # Output of varnishstat -V : `varnishstat (varnish-4.1.1 revision 66bb824)`
    version_pattern = re.compile(r'(\d+\.\d+\.\d+)')

    # XML parsing bits, a.k.a. Kafka in Code
    def _reset(self):
        self._current_element = ""
        self._current_metric = "varnish"
        self._current_value = 0
        self._current_str = ""
        self._current_type = ""

    def _start_element(self, name, attrs):
        self._current_element = name

    def _end_element(self, name, tags):
        if name == "stat":
            m_name = self.normalize(self._current_metric)
            if self._current_type in ("a", "c"):
                self.rate(m_name, long(self._current_value), tags=tags)
            elif self._current_type in ("i", "g"):
                self.gauge(m_name, long(self._current_value), tags=tags)
                if 'n_purges' in m_name:
                    self.rate('varnish.n_purgesps', long(self._current_value), tags=tags)
            else:
                # Unsupported data type, ignore
                self._reset()
                return  # don't save

            # reset for next stat element
            self._reset()
        elif name in ("ident", "name") or (name == "type" and self._current_str != "MAIN"):
            self._current_metric += "." + self._current_str

    def _char_data(self, data):
        self.log.debug("Data %s [%s]" % (data, self._current_element))
        data = data.strip()
        if len(data) > 0 and self._current_element != "":
            if self._current_element == "value":
                self._current_value = long(data)
            elif self._current_element == "flag":
                self._current_type = data
            else:
                self._current_str = data

    def check(self, instance):
        # Not configured? Not a problem.
        if instance.get("varnishstat", None) is None:
            raise Exception("varnishstat is not configured")
        custom_tags = instance.get('tags', [])
        if custom_tags is None:
            custom_tags = []
        else:
            custom_tags = list(set(custom_tags))
        # Split the varnishstat command so that additional arguments can be passed in
        # In order to support monitoring a Varnish instance which is running as a Docker
        # container we need to wrap commands (varnishstat, varnishadm) with scripts which
        # perform a docker exec on the running container. This works fine when running a
        # single container on the host but breaks down when attempting to use the auto
        # discovery feature. This change allows for passing in additional parameters to
        # the script (i.e. %%host%%) so that the command is properly formatted and the
        # desired container is queried.
        varnishstat_path = shlex.split(instance.get("varnishstat"))
        name = instance.get('name')
        metrics_filter = instance.get("metrics_filter", [])
        if not isinstance(metrics_filter, list):
            raise Exception("The parameter 'metrics_filter' must be a list")

        # Get version and version-specific args from varnishstat -V.
        version, varnishstat_format = self._get_version_info(varnishstat_path)

        cmd = varnishstat_path + [self.VARNISHSTAT_FORMAT_OPTION[varnishstat_format]]
        for metric in metrics_filter:
            cmd.extend(["-f", metric])

        if name is not None:
            cmd.extend(['-n', name])
            tags = custom_tags + [u'varnish_name:%s' % name]
        else:
            tags = custom_tags + [u'varnish_name:default']

        output, _, _ = get_subprocess_output(cmd, self.log)

        self._parse_varnishstat(output, varnishstat_format, tags)

        # Parse service checks from varnishadm.
        if instance.get("varnishadm", None):
            # Split the varnishadm command so that additional arguments can be passed in
            # In order to support monitoring a Varnish instance which is running as a Docker
            # container we need to wrap commands (varnishstat, varnishadm) with scripts which
            # perform a docker exec on the running container. This works fine when running a
            # single container on the host but breaks down when attempting to use the auto
            # discovery feature. This change allows for passing in additional parameters to
            # the script (i.e. %%host%%) so that the command is properly formatted and the
            # desired container is queried.
            varnishadm_path = shlex.split(instance.get('varnishadm'))
            secretfile_path = instance.get('secretfile', '/etc/varnish/secret')

            daemon_host = instance.get('daemon_host', 'localhost')
            daemon_port = instance.get('daemon_port', '6082')

            cmd = []
            if geteuid() != 0:
                cmd.append('sudo')

            if version < LooseVersion('4.1.0'):
                cmd.extend(varnishadm_path + ['-S', secretfile_path, 'debug.health'])
            else:
                cmd.extend(
                    varnishadm_path
                    + ['-T', '{}:{}'.format(daemon_host, daemon_port), '-S', secretfile_path, 'backend.list', '-p']
                )

            try:
                output, err, _ = get_subprocess_output(cmd, self.log)
            except OSError as e:
                self.log.error("There was an error running varnishadm. Make sure 'sudo' is available. %s", e)
                output = None
            if err:
                self.log.error('Error getting service check from varnishadm: %s', err)

            if output:
                self._parse_varnishadm(output, custom_tags)

    def _get_version_info(self, varnishstat_path):
        # Get the varnish version from varnishstat
        output, error, _ = get_subprocess_output(varnishstat_path + ["-V"], self.log, raise_on_empty_output=False)

        # Assumptions regarding varnish's version
        varnishstat_format = "json"
        raw_version = None

        m1 = self.version_pattern.search(output, re.MULTILINE)
        # v2 prints the version on stderr, v3 on stdout
        m2 = self.version_pattern.search(error, re.MULTILINE)

        if m1 is None and m2 is None:
            self.log.warn("Cannot determine the version of varnishstat, assuming 3 or greater")
            self.warning("Cannot determine the version of varnishstat, assuming 3 or greater")
        else:
            if m1 is not None:
                raw_version = m1.group()
            elif m2 is not None:
                raw_version = m2.group()

        self.log.debug("Varnish version: %s", raw_version)

        if raw_version:
            self.set_metadata('version', raw_version)

        if raw_version is None:
            raw_version = '3.0.0'

        version = LooseVersion(raw_version)

        # Location of varnishstat
        if version < LooseVersion('3.0.0'):
            varnishstat_format = "text"
        elif version < LooseVersion('5.0.0'):  # we default to json starting version 5.0.0
            varnishstat_format = "xml"

        return version, varnishstat_format

    def _parse_varnishstat(self, output, varnishstat_format, tags=None):
        """
        The text option (-1) is not reliable enough when counters get large.
        VBE.media_video_prd_services_01(10.93.67.16,,8080).happy18446744073709551615

        2 types of data, "a" for counter ("c" in newer versions of varnish), "i" for gauge ("g")
        https://github.com/varnish/Varnish-Cache/blob/master/include/tbl/vsc_fields.h

        Bitmaps are not supported.
        """
        tags = tags or []
        # FIXME: this check is processing an unbounded amount of data
        # we should explicitly list the metrics we want to get from the check
        if varnishstat_format == "xml":
            p = xml.parsers.expat.ParserCreate()
            p.StartElementHandler = self._start_element
            p.EndElementHandler = lambda name: self._end_element(name, tags)
            p.CharacterDataHandler = self._char_data
            self._reset()
            p.Parse(output, True)
        elif varnishstat_format == "json":
            json_output = json.loads(output)
            for name, metric in iteritems(json_output):
                if not isinstance(metric, dict):  # skip 'timestamp' field
                    continue

                if name.startswith("MAIN."):
                    name = name.split('.', 1)[1]
                value = metric.get("value", 0)

                if metric["flag"] in ("a", "c"):
                    self.rate(self.normalize(name, prefix="varnish"), long(value), tags=tags)
                elif metric["flag"] in ("g", "i"):
                    self.gauge(self.normalize(name, prefix="varnish"), long(value), tags=tags)
                    if 'n_purges' in self.normalize(name, prefix="varnish"):
                        self.rate('varnish.n_purgesps', long(value), tags=tags)
        elif varnishstat_format == "text":
            for line in output.split("\n"):
                self.log.debug("Parsing varnish results: %s" % line)
                fields = line.split()
                if len(fields) < 3:
                    break
                name, gauge_val, rate_val = fields[0], fields[1], fields[2]
                metric_name = self.normalize(name, prefix="varnish")

                # Now figure out which value to pick
                if rate_val.lower() in ("nan", "."):
                    # col 2 matters
                    self.log.debug("Varnish (gauge) %s %d" % (metric_name, int(gauge_val)))
                    self.gauge(metric_name, int(gauge_val), tags=tags)
                    if 'n_purges' in metric_name:
                        self.rate('varnish.n_purgesps', float(gauge_val), tags=tags)
                else:
                    # col 3 has a rate (since restart)
                    self.log.debug("Varnish (rate) %s %d" % (metric_name, int(gauge_val)))
                    self.rate(metric_name, float(gauge_val), tags=tags)

    def _parse_varnishadm(self, output, tags):
        """ Parse out service checks from varnishadm.

        Example output:

            Backend b0 is Sick
            Current states  good:  2 threshold:  3 window:  5
            Average responsetime of good probes: 0.000000
            Oldest                                                    Newest
            ================================================================
            -------------------------------------------------------------444 Good IPv4
            -------------------------------------------------------------XXX Good Xmit
            -------------------------------------------------------------RRR Good Recv
            ----------------------------------------------------------HHH--- Happy
            Backend b1 is Sick
            Current states  good:  2 threshold:  3 window:  5
            Average responsetime of good probes: 0.000000
            Oldest                                                    Newest
            ================================================================
            ----------------------------------------------------------HHH--- Happy

        Example output (new output format):

            Backend name                   Admin      Probe
            boot.default                   probe      Healthy (no probe)
            boot.backend2                  probe      Healthy 4/4
              Current states  good:  4 threshold:  3 window:  4
              Average response time of good probes: 0.002504
              Oldest ================================================== Newest
              --------------------------------------------------------------44 Good IPv4
              --------------------------------------------------------------XX Good Xmit
              --------------------------------------------------------------RR Good Recv
              ------------------------------------------------------------HHHH Happy

        """
        # Process status by backend.
        backends_by_status = defaultdict(list)
        for line in output.split("\n"):
            backend, status, message = None, None, None
            # split string and remove all empty fields
            tokens = filter(None, line.strip().split(' '))
            tokens = [t for t in tokens]
            if len(tokens):
                if tokens == ['Backend', 'name', 'Admin', 'Probe']:
                    # skip the column headers that exist in new output format
                    continue
                # parse new output format
                # the backend name will include the vcl name
                # so split on first . to remove prefix
                elif len(tokens) >= 4 and tokens[1] in ['healthy', 'sick']:
                    # If the backend health was overriden, lets grab the
                    # overriden value instead of the probed health
                    backend = tokens[0].split('.', 1)[-1]
                    status = tokens[1].lower()
                elif len(tokens) >= 4 and tokens[1] == 'probe':
                    backend = tokens[0].split('.', 1)[-1]
                    status = tokens[2].lower()
                # Parse older Varnish backend output
                elif tokens[0] == 'Backend':
                    backend = tokens[1]
                    status = tokens[-1].lower()

                if tokens[0] == 'Current' and backend is not None:
                    try:
                        message = ' '.join(tokens[2:]).strip()
                    except Exception:
                        # If we can't parse a message still send a status.
                        self.log.exception('Error when parsing message from varnishadm')
                        message = ''

                if backend is not None:
                    backends_by_status[status].append((backend, message))

        for status, backends in iteritems(backends_by_status):
            check_status = BackendStatus.to_check_status(status)
            for backend, message in backends:
                service_checks_tags = ['backend:%s' % backend] + tags
                self.service_check(self.SERVICE_CHECK_NAME, check_status, tags=service_checks_tags, message=message)
<commit_msg>Remove shlex (#5065)<commit_after># (C) Datadog, Inc. 2018
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
import json
import re
import xml.parsers.expat  # python 2.4 compatible
from collections import defaultdict
from distutils.version import LooseVersion
from os import geteuid

from six import PY3, iteritems
from six.moves import filter

from datadog_checks.checks import AgentCheck
from datadog_checks.utils.subprocess_output import get_subprocess_output

if PY3:
    long = int

class BackendStatus(object):
    HEALTHY = 'healthy'
    SICK = 'sick'
    ALL = (HEALTHY, SICK)

    @classmethod
    def to_check_status(cls, status):
        if status == cls.HEALTHY:
            return AgentCheck.OK
        elif status == cls.SICK:
            return AgentCheck.CRITICAL
        return AgentCheck.UNKNOWN

class Varnish(AgentCheck):
    SERVICE_CHECK_NAME = 'varnish.backend_healthy'
    # Parse metrics from varnishstat.
    VARNISHSTAT_FORMAT_OPTION = {
        "text": "-1",  # version < 3.0.0
        "xml": "-x",  # version >= 3.0.0 < 5.0.0
        "json": "-j",  # version >=5.0.0
    }

    # Output of varnishstat -V : `varnishstat (varnish-4.1.1 revision 66bb824)`
    version_pattern = re.compile(r'(\d+\.\d+\.\d+)')

    # XML parsing bits, a.k.a. Kafka in Code
    def _reset(self):
        self._current_element = ""
        self._current_metric = "varnish"
        self._current_value = 0
        self._current_str = ""
        self._current_type = ""

    def _start_element(self, name, attrs):
        self._current_element = name

    def _end_element(self, name, tags):
        if name == "stat":
            m_name = self.normalize(self._current_metric)
            if self._current_type in ("a", "c"):
                self.rate(m_name, long(self._current_value), tags=tags)
            elif self._current_type in ("i", "g"):
                self.gauge(m_name, long(self._current_value), tags=tags)
                if 'n_purges' in m_name:
                    self.rate('varnish.n_purgesps', long(self._current_value), tags=tags)
            else:
                # Unsupported data type, ignore
                self._reset()
                return  # don't save

            # reset for next stat element
            self._reset()
        elif name in ("ident", "name") or (name == "type" and self._current_str != "MAIN"):
            self._current_metric += "." + self._current_str

    def _char_data(self, data):
        self.log.debug("Data %s [%s]" % (data, self._current_element))
        data = data.strip()
        if len(data) > 0 and self._current_element != "":
            if self._current_element == "value":
                self._current_value = long(data)
            elif self._current_element == "flag":
                self._current_type = data
            else:
                self._current_str = data

    def check(self, instance):
        # Not configured? Not a problem.
        if instance.get("varnishstat", None) is None:
            raise Exception("varnishstat is not configured")
        custom_tags = instance.get('tags', [])
        if custom_tags is None:
            custom_tags = []
        else:
            custom_tags = list(set(custom_tags))
        # Split the varnishstat command so that additional arguments can be passed in
        # In order to support monitoring a Varnish instance which is running as a Docker
        # container we need to wrap commands (varnishstat, varnishadm) with scripts which
        # perform a docker exec on the running container. This works fine when running a
        # single container on the host but breaks down when attempting to use the auto
        # discovery feature. This change allows for passing in additional parameters to
        # the script (i.e. %%host%%) so that the command is properly formatted and the
        # desired container is queried.
        varnishstat_path = instance.get('varnishstat', '').split()
        name = instance.get('name')
        metrics_filter = instance.get("metrics_filter", [])
        if not isinstance(metrics_filter, list):
            raise Exception("The parameter 'metrics_filter' must be a list")

        # Get version and version-specific args from varnishstat -V.
        version, varnishstat_format = self._get_version_info(varnishstat_path)

        cmd = varnishstat_path + [self.VARNISHSTAT_FORMAT_OPTION[varnishstat_format]]
        for metric in metrics_filter:
            cmd.extend(["-f", metric])

        if name is not None:
            cmd.extend(['-n', name])
            tags = custom_tags + [u'varnish_name:%s' % name]
        else:
            tags = custom_tags + [u'varnish_name:default']

        output, _, _ = get_subprocess_output(cmd, self.log)

        self._parse_varnishstat(output, varnishstat_format, tags)

        # Parse service checks from varnishadm.
        if instance.get("varnishadm", None):
            # Split the varnishadm command so that additional arguments can be passed in
            # In order to support monitoring a Varnish instance which is running as a Docker
            # container we need to wrap commands (varnishstat, varnishadm) with scripts which
            # perform a docker exec on the running container. This works fine when running a
            # single container on the host but breaks down when attempting to use the auto
            # discovery feature. This change allows for passing in additional parameters to
            # the script (i.e. %%host%%) so that the command is properly formatted and the
            # desired container is queried.
            varnishadm_path = instance.get('varnishadm', '').split()
            secretfile_path = instance.get('secretfile', '/etc/varnish/secret')

            daemon_host = instance.get('daemon_host', 'localhost')
            daemon_port = instance.get('daemon_port', '6082')

            cmd = []
            if geteuid() != 0:
                cmd.append('sudo')

            if version < LooseVersion('4.1.0'):
                cmd.extend(varnishadm_path + ['-S', secretfile_path, 'debug.health'])
            else:
                cmd.extend(
                    varnishadm_path
                    + ['-T', '{}:{}'.format(daemon_host, daemon_port), '-S', secretfile_path, 'backend.list', '-p']
                )

            try:
                output, err, _ = get_subprocess_output(cmd, self.log)
            except OSError as e:
                self.log.error("There was an error running varnishadm. Make sure 'sudo' is available. %s", e)
                output = None
            if err:
                self.log.error('Error getting service check from varnishadm: %s', err)

            if output:
                self._parse_varnishadm(output, custom_tags)

    def _get_version_info(self, varnishstat_path):
        # Get the varnish version from varnishstat
        output, error, _ = get_subprocess_output(varnishstat_path + ["-V"], self.log, raise_on_empty_output=False)

        # Assumptions regarding varnish's version
        varnishstat_format = "json"
        raw_version = None

        m1 = self.version_pattern.search(output, re.MULTILINE)
        # v2 prints the version on stderr, v3 on stdout
        m2 = self.version_pattern.search(error, re.MULTILINE)

        if m1 is None and m2 is None:
            self.log.warn("Cannot determine the version of varnishstat, assuming 3 or greater")
            self.warning("Cannot determine the version of varnishstat, assuming 3 or greater")
        else:
            if m1 is not None:
                raw_version = m1.group()
            elif m2 is not None:
                raw_version = m2.group()

        self.log.debug("Varnish version: %s", raw_version)

        if raw_version:
            self.set_metadata('version', raw_version)

        if raw_version is None:
            raw_version = '3.0.0'

        version = LooseVersion(raw_version)

        # Location of varnishstat
        if version < LooseVersion('3.0.0'):
            varnishstat_format = "text"
        elif version < LooseVersion('5.0.0'):  # we default to json starting version 5.0.0
            varnishstat_format = "xml"

        return version, varnishstat_format

    def _parse_varnishstat(self, output, varnishstat_format, tags=None):
        """
        The text option (-1) is not reliable enough when counters get large.
        VBE.media_video_prd_services_01(10.93.67.16,,8080).happy18446744073709551615

        2 types of data, "a" for counter ("c" in newer versions of varnish), "i" for gauge ("g")
        https://github.com/varnish/Varnish-Cache/blob/master/include/tbl/vsc_fields.h

        Bitmaps are not supported.
        """
        tags = tags or []
        # FIXME: this check is processing an unbounded amount of data
        # we should explicitly list the metrics we want to get from the check
        if varnishstat_format == "xml":
            p = xml.parsers.expat.ParserCreate()
            p.StartElementHandler = self._start_element
            p.EndElementHandler = lambda name: self._end_element(name, tags)
            p.CharacterDataHandler = self._char_data
            self._reset()
            p.Parse(output, True)
        elif varnishstat_format == "json":
            json_output = json.loads(output)
            for name, metric in iteritems(json_output):
                if not isinstance(metric, dict):  # skip 'timestamp' field
                    continue

                if name.startswith("MAIN."):
                    name = name.split('.', 1)[1]
                value = metric.get("value", 0)

                if metric["flag"] in ("a", "c"):
                    self.rate(self.normalize(name, prefix="varnish"), long(value), tags=tags)
                elif metric["flag"] in ("g", "i"):
                    self.gauge(self.normalize(name, prefix="varnish"), long(value), tags=tags)
                    if 'n_purges' in self.normalize(name, prefix="varnish"):
                        self.rate('varnish.n_purgesps', long(value), tags=tags)
        elif varnishstat_format == "text":
            for line in output.split("\n"):
                self.log.debug("Parsing varnish results: %s" % line)
                fields = line.split()
                if len(fields) < 3:
                    break
                name, gauge_val, rate_val = fields[0], fields[1], fields[2]
                metric_name = self.normalize(name, prefix="varnish")

                # Now figure out which value to pick
                if rate_val.lower() in ("nan", "."):
                    # col 2 matters
                    self.log.debug("Varnish (gauge) %s %d" % (metric_name, int(gauge_val)))
                    self.gauge(metric_name, int(gauge_val), tags=tags)
                    if 'n_purges' in metric_name:
                        self.rate('varnish.n_purgesps', float(gauge_val), tags=tags)
                else:
                    # col 3 has a rate (since restart)
                    self.log.debug("Varnish (rate) %s %d" % (metric_name, int(gauge_val)))
                    self.rate(metric_name, float(gauge_val), tags=tags)

    def _parse_varnishadm(self, output, tags):
        """ Parse out service checks from varnishadm.

        Example output:

            Backend b0 is Sick
            Current states  good:  2 threshold:  3 window:  5
            Average responsetime of good probes: 0.000000
            Oldest                                                    Newest
            ================================================================
            -------------------------------------------------------------444 Good IPv4
            -------------------------------------------------------------XXX Good Xmit
            -------------------------------------------------------------RRR Good Recv
            ----------------------------------------------------------HHH--- Happy
            Backend b1 is Sick
            Current states  good:  2 threshold:  3 window:  5
            Average responsetime of good probes: 0.000000
            Oldest                                                    Newest
            ================================================================
            ----------------------------------------------------------HHH--- Happy

        Example output (new output format):

            Backend name                   Admin      Probe
            boot.default                   probe      Healthy (no probe)
            boot.backend2                  probe      Healthy 4/4
              Current states  good:  4 threshold:  3 window:  4
              Average response time of good probes: 0.002504
              Oldest ================================================== Newest
              --------------------------------------------------------------44 Good IPv4
              --------------------------------------------------------------XX Good Xmit
              --------------------------------------------------------------RR Good Recv
              ------------------------------------------------------------HHHH Happy

        """
        # Process status by backend.
        backends_by_status = defaultdict(list)
        for line in output.split("\n"):
            backend, status, message = None, None, None
            # split string and remove all empty fields
            tokens = filter(None, line.strip().split(' '))
            tokens = [t for t in tokens]
            if len(tokens):
                if tokens == ['Backend', 'name', 'Admin', 'Probe']:
                    # skip the column headers that exist in new output format
                    continue
                # parse new output format
                # the backend name will include the vcl name
                # so split on first . to remove prefix
                elif len(tokens) >= 4 and tokens[1] in ['healthy', 'sick']:
                    # If the backend health was overriden, lets grab the
                    # overriden value instead of the probed health
                    backend = tokens[0].split('.', 1)[-1]
                    status = tokens[1].lower()
                elif len(tokens) >= 4 and tokens[1] == 'probe':
                    backend = tokens[0].split('.', 1)[-1]
                    status = tokens[2].lower()
                # Parse older Varnish backend output
                elif tokens[0] == 'Backend':
                    backend = tokens[1]
                    status = tokens[-1].lower()

                if tokens[0] == 'Current' and backend is not None:
                    try:
                        message = ' '.join(tokens[2:]).strip()
                    except Exception:
                        # If we can't parse a message still send a status.
                        self.log.exception('Error when parsing message from varnishadm')
                        message = ''

                if backend is not None:
                    backends_by_status[status].append((backend, message))

        for status, backends in iteritems(backends_by_status):
            check_status = BackendStatus.to_check_status(status)
            for backend, message in backends:
                service_checks_tags = ['backend:%s' % backend] + tags
                self.service_check(self.SERVICE_CHECK_NAME, check_status, tags=service_checks_tags, message=message)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright © Spyder Project Contributors
# Licensed under the terms of the MIT License
#

"""Tests for the plugin."""

# Test library imports
import os
import pytest
import os.path
from qtpy.QtCore import Qt
from qtpy.QtWebEngineWidgets import WEBENGINE

# Local imports
from spyder_terminal.terminalplugin import TerminalPlugin

LOCATION = os.path.realpath(os.path.join(os.getcwd(),
                                         os.path.dirname(__file__)))

TERM_UP = 5000

def check_pwd(termwidget):
    """Check if pwd command is executed."""
    if WEBENGINE:
        def callback(data):
            global html
            html = data
        termwidget.body.toHtml(callback)
        try:
            return LOCATION in html
        except NameError:
            return False
    else:
        return LOCATION in termwidget.body.toHtml()

@pytest.fixture
def setup_terminal(qtbot):
    """Set up the Notebook plugin."""
    terminal = TerminalPlugin(None)
    qtbot.addWidget(terminal)
    terminal.create_new_term()
    terminal.show()
    return terminal

def test_new_terminal(qtbot):
    """Test if a new terminal is added."""
    # Setup widget
    terminal = setup_terminal(qtbot)
    term = terminal.get_current_term()
    qtbot.wait(1000)

    # Move to LOCATION
    qtbot.keyClicks(term.view, 'cd {}'.format(LOCATION))
    qtbot.keyPress(term.view, Qt.Key_Return)

    # Run pwd
    qtbot.keyClicks(term.view, 'pwd')
    qtbot.keyPress(term.view, Qt.Key_Return)

    # Assert pwd is LOCATION
    qtbot.waitUntil(lambda: check_pwd(term), timeout=TERM_UP)
    assert len(terminal.terms) == 1

<commit_msg>Testing: Clear terminal before getting pwd<commit_after># -*- coding: utf-8 -*-
#
# Copyright © Spyder Project Contributors
# Licensed under the terms of the MIT License
#

"""Tests for the plugin."""

# Test library imports
import os
import pytest
import os.path
from qtpy.QtCore import Qt
from qtpy.QtWebEngineWidgets import WEBENGINE

# Local imports
from spyder_terminal.terminalplugin import TerminalPlugin

LOCATION = os.path.realpath(os.path.join(os.getcwd(),
                                         os.path.dirname(__file__)))

TERM_UP = 5000

def check_pwd(termwidget):
    """Check if pwd command is executed."""
    if WEBENGINE:
        def callback(data):
            global html
            html = data
        termwidget.body.toHtml(callback)
        try:
            return LOCATION in html
        except NameError:
            return False
    else:
        return LOCATION in termwidget.body.toHtml()

@pytest.fixture
def setup_terminal(qtbot):
    """Set up the Notebook plugin."""
    terminal = TerminalPlugin(None)
    qtbot.addWidget(terminal)
    terminal.create_new_term()
    terminal.show()
    return terminal

def test_new_terminal(qtbot):
    """Test if a new terminal is added."""
    # Setup widget
    terminal = setup_terminal(qtbot)
    term = terminal.get_current_term()
    qtbot.wait(TERM_UP)

    # Move to LOCATION
    qtbot.keyClicks(term.view, 'cd {}'.format(LOCATION))
    qtbot.keyPress(term.view, Qt.Key_Return)

    # Clear
    qtbot.keyClicks(term.view, 'clear')
    qtbot.keyPress(term.view, Qt.Key_Return)

    # Run pwd
    qtbot.keyClicks(term.view, 'pwd')
    qtbot.keyPress(term.view, Qt.Key_Return)

    # Assert pwd is LOCATION
    qtbot.waitUntil(lambda: check_pwd(term), timeout=TERM_UP)
    assert len(terminal.terms) == 1

<|endoftext|>