<commit_before>"""
Tools for high level profiling and event logging
"""
import gc
import os
import inspect

global memHardLimit,memLast,memLog,logFile,logLevel,verbose,procID,logAllProcesses,flushBuffer, preInitBuffer

memHardLimit=0.0
memMax=0.0
memLast=0.0
memLog=False
memList=[]
logFile=None
logLevel=0
logAllProcesses=False
verbose=False
procID=None
flushBuffer=False
preInitBuffer=[]

def memProfOn():
    global memLog
    memLog = True

def memHardLimitOn(value):
    global memHardLimit
    if not memLog:
        memProfOn()
    memHardLimit = value*1024.0

#these are so the optparser can turn these on and off
def memProfOn_callback(option,opt,value,parser):
    memProfOn()

def memHardLimitOn_callback(option,opt,value,parser):
    memHardLimitOn(value)

def verboseOn_callback(option,opt,value,parser):
    global verbose
    verbose=True

def openLog(filename,level):
    global logFile
    global logLevel
    global preInitBuffer
    logFile=open(filename,'w')
    logLevel = level
    for (string,level,data) in preInitBuffer:
        logEvent(string,level,data)

def closeLog():
    global logFile
    try:
        logFile.close()
    except:
        pass

def logEvent(stringIn,level=1,data=None):
    global logLevel,procID,logAllProcesses,flushBuffer,preInitBuffer
    if procID != None:
        if logAllProcesses or procID==0:
            if level < logLevel:
                if logAllProcesses and procID != None and stringIn != None:
                    string = "Proc %d : " % procID
                    string += stringIn
                else:
                    string = stringIn
                if string!=None:
                    if data != None:
                        string += `data`
                    string+='\n'
                    global logFile,verbose
                    logFile.write(string)
                    if flushBuffer:
                        logFile.flush()
                    if verbose:
                        import sys
                        sys.stdout.write(string)
                        if flushBuffer:
                            sys.stdout.flush()
    elif procID==None:
        preInitBuffer.append((stringIn,level,data))

def memory(message=None,className='',memSaved=None):
    global memLog
    global memLast
    global memList
    global memMax
    if memLog:
        gc.collect()
        try:
            import memory_profiler
            mem = memory_usage(-1,interval=0.1,timeout=0.1)[0]
        except:
            try:
                mem = int(os.popen('ps -p %d -o rss|tail -1' % os.getpid()).read())/1024.0
            except:
                try:
                    #CLE doesn't have full ps functionality
                    #try to find first occurence of pid, then get memory usage slot
                    #mwf debug
                    #import pdb
                    #pdb.set_trace()
                    plist = os.popen('ps ').read()
                    pr = plist.split('%s' % os.getpid())
                    #print 'pr = %s ' % pr
                    #print 'pr[1] = %s ' % pr[1]
                    mem = int(pr[1].split()[1])/1024.0
                except:
                    logEvent("memory function doesn't work on this platform\n")
                    mem = 0
        if mem > memMax:
            memMax = mem
        if memSaved != None:
            memInc = mem - memSaved
        else:
            memInc = mem-memLast
        memList.append((message,memInc))
        memLast = mem
        caller = inspect.stack()[1][3]
        line = inspect.stack()[1][2]
        filename = inspect.stack()[1][1]
        if className:
            className += '.'
        if memHardLimit:
            if mem > memHardLimit:
                import sys
                if className == None:
                    className = "UnknownClass"
                if caller == None:
                    caller="UnknownCaller"
                if line == None:
                    line = "Unknown Line"
                if message == None:
                    message = ''
                sys.exit("ERROR: MEMORY HARDLIMIT REACHED, EXIT From "+filename.split("/")[-1]+", "+className+caller+", line "+`line`+": "+message+", %d MB in routine, %d MB in program, %d MB is hard limit" % (memInc,mem,memHardLimit))
        if message:
            return "In "+filename.split("/")[-1]+", "+className+caller+", line "+`line`+": "+message+", %d MB in routine, %d MB in program" % (memInc,mem)

def memorySummary():
    global memLog
    global memList
    global memMax
    if memLog:
        gc.collect()
        try:
            import memory_profiler
            mem = memory_usage(-1,interval=0.1,timeout=0.1)[0]
        except:
            try:
                mem = int(os.popen('ps -p %d -o vsz|tail -1' % os.getpid()).read())/1024.0
            except:
                try:
                    #CLE doesn't have full ps functionality
                    #try to find first occurence of pid, then get memory usage slot
                    #mwf debug
                    #import pdb
                    #pdb.set_trace()
                    plist = os.popen('ps ').read()
                    pr = plist.split('%s' % os.getpid())
                    #print 'pr = %s ' % pr
                    #print 'pr[1] = %s ' % pr[1]
                    mem = int(pr[1].split()[1])/1024.0
                except:
                    logEvent("memory function doesn't work on this platform\n")
                    mem = 0
        if mem > 0:
            for pair in memList:
                logEvent(`pair[0]`+"  %"+`100.0*pair[1]/memMax`)
<commit_msg>comment out old shell-based ps commands in mem profiler and added MPI.COMM_WORLD.Abort(1) so we terminate cleanly<commit_after>"""
Tools for high level profiling and event logging
"""
import gc
import os
import inspect

global memHardLimit,memLast,memLog,logFile,logLevel,verbose,procID,logAllProcesses,flushBuffer, preInitBuffer

memHardLimit=0.0
memMax=0.0
memLast=0.0
memLog=False
memList=[]
logFile=None
logLevel=0
logAllProcesses=False
verbose=False
procID=None
flushBuffer=False
preInitBuffer=[]

def memProfOn():
    global memLog
    memLog = True

def memHardLimitOn(value):
    global memHardLimit
    if not memLog:
        memProfOn()
    memHardLimit = value*1024.0

#these are so the optparser can turn these on and off
def memProfOn_callback(option,opt,value,parser):
    memProfOn()

def memHardLimitOn_callback(option,opt,value,parser):
    memHardLimitOn(value)

def verboseOn_callback(option,opt,value,parser):
    global verbose
    verbose=True

def openLog(filename,level):
    global logFile
    global logLevel
    global preInitBuffer
    logFile=open(filename,'w')
    logLevel = level
    for (string,level,data) in preInitBuffer:
        logEvent(string,level,data)

def closeLog():
    global logFile
    try:
        logFile.close()
    except:
        pass

def logEvent(stringIn,level=1,data=None):
    global logLevel,procID,logAllProcesses,flushBuffer,preInitBuffer
    if procID != None:
        if logAllProcesses or procID==0:
            if level < logLevel:
                if logAllProcesses and procID != None and stringIn != None:
                    string = "Proc %d : " % procID
                    string += stringIn
                else:
                    string = stringIn
                if string!=None:
                    if data != None:
                        string += `data`
                    string+='\n'
                    global logFile,verbose
                    logFile.write(string)
                    if flushBuffer:
                        logFile.flush()
                    if verbose:
                        import sys
                        sys.stdout.write(string)
                        if flushBuffer:
                            sys.stdout.flush()
    elif procID==None:
        preInitBuffer.append((stringIn,level,data))

def memory(message=None,className='',memSaved=None):
    global memLog
    global memLast
    global memList
    global memMax
    if memLog:
        gc.collect()
        try:
            from memory_profiler import memory_usage
            memList = memory_usage(-1)
            if len(memList) < 1:
                memList = memory_usage(-1,timeout=10)
                assert(len(memList) > 0)
            mem = memList[-1]
        except:
            raise
        # except:
#             try:
#                 mem = int(os.popen('ps -p %d -o rss|tail -1' % os.getpid()).read())/1024.0
#             except:
#                 try:
#                     #CLE doesn't have full ps functionality
#                     #try to find first occurence of pid, then get memory usage slot
#                     #mwf debug
#                     #import pdb
#                     #pdb.set_trace()
#                     plist = os.popen('ps ').read()
#                     pr = plist.split('%s' % os.getpid())
#                     #print 'pr = %s ' % pr
#                     #print 'pr[1] = %s ' % pr[1]
#                     mem = int(pr[1].split()[1])/1024.0
#                 except:
#                     logEvent("memory function doesn't work on this platform\n")
#                     mem = 0
        if mem > memMax:
            memMax = mem
        if memSaved != None:
            memInc = mem - memSaved
        else:
            memInc = mem-memLast
        memList.append((message,memInc))
        memLast = mem
        caller = inspect.stack()[1][3]
        line = inspect.stack()[1][2]
        filename = inspect.stack()[1][1]
        if className:
            className += '.'
        if memHardLimit:
            if mem > memHardLimit:
                import sys
                from mpi4py import MPI
                if className == None:
                    className = "UnknownClass"
                if caller == None:
                    caller="UnknownCaller"
                if line == None:
                    line = "Unknown Line"
                if message == None:
                    message = ''
                logEvent("PROTEUS ERROR: MEMORY HARDLIMIT REACHED, EXIT From "+filename.split("/")[-1]+", "+className+caller+", line "+`line`+": "+message+", %d MB in routine, %d MB in program, %d MB is hard limit" % (memInc,mem,memHardLimit))
                MPI.COMM_WORLD.Abort(1)
                sys.exit("MPI.COMM_WORLD.Abort(1); exit(1)")
        if message:
            return "In "+filename.split("/")[-1]+", "+className+caller+", line "+`line`+": "+message+", %d MB in routine, %d MB in program" % (memInc,mem)

def memorySummary():
    global memLog
    global memList
    global memMax
    if memLog:
        gc.collect()
        try:
            from memory_profile import memory_usage
            mem = memory_usage(-1)[0]
        except:
            raise
        # except:
#             try:
#                 mem = int(os.popen('ps -p %d -o vsz|tail -1' % os.getpid()).read())/1024.0
#             except:
#                 try:
#                     #CLE doesn't have full ps functionality
#                     #try to find first occurence of pid, then get memory usage slot
#                     #mwf debug
#                     #import pdb
#                     #pdb.set_trace()
#                     plist = os.popen('ps ').read()
#                     pr = plist.split('%s' % os.getpid())
#                     #print 'pr = %s ' % pr
#                     #print 'pr[1] = %s ' % pr[1]
#                     mem = int(pr[1].split()[1])/1024.0
#                 except:
#                     logEvent("memory function doesn't work on this platform\n")
#                     mem = 0
        if mem > 0:
            for pair in memList:
                logEvent(`pair[0]`+"  %"+`100.0*pair[1]/memMax`)
<|endoftext|>


<commit_before># coding: utf-8

from django.db import models
# from django.db.models import Q
from django.db.models.signals import post_save  # , pre_save, m2m_changed
from django.utils.translation import ugettext_lazy as _

import datetime
# from datetime import date
now = datetime.datetime.now()

from protocolle.core.models import (TipoDocumento, Carater, Natureza,
                                    Status, OPERATION_CHOICES)
from protocolle.auxiliar.models import (Instituicao, Pessoa, Setor)

class Documento(models.Model):
    operacao = models.IntegerField(_(u'Operação'), choices=OPERATION_CHOICES)
    tipo_documento = models.ForeignKey(TipoDocumento,
                                       verbose_name=_(u'Tipo de Documento'),
                                       related_name='DocTiposDocumento')
    protocolo = models.CharField(_(u'Número do Protocolo'), max_length=100,
                                 unique=True, default=0, editable=False)
    numero = models.CharField(_(u'Número do Documento'), max_length=100)
    data_recebimento = models.DateTimeField(_(u'Data de Recebimento'),
                                            auto_now_add=True)
    data_documento = models.DateField(_(u'Data do Documento'))
    data_validade = models.DateField(_(u'Data de Validade'),
                                     null=True, blank=True)
    folhas = models.IntegerField(_(u'Número de Folhas'), null=True, blank=True)
    carater = models.ForeignKey(Carater, verbose_name=_(u'Caráter'),
                                related_name='DocCarateres')
    natureza = models.ForeignKey(Natureza, verbose_name=_(u'Natureza'),
                                 related_name='DocNaturezas')
    origem = models.ForeignKey(Instituicao,
                               verbose_name=_(u'Instituição Origem'),
                               related_name='DocInstituicoesOri')
    destino = models.ForeignKey(Instituicao, editable=False,
                                verbose_name=_(u'Instituição Destino'),
                                related_name='DocInstituicoesDes')
    assunto = models.TextField(_(u'Assunto'))
    interessado = models.ForeignKey(Pessoa, verbose_name=_(u'Interessado'),
                                    related_name='DocPessoasRem')
    status = models.ForeignKey(Status, verbose_name=_(u'Status'),
                               related_name='DocSituacoes', editable=False)
    observacao = models.TextField(_(u'Observação'), null=True, blank=True)

    def save(self, *args, **kwargs):
        """
        Funcao para gerar o número de protocolo
        unico e personalizado
        """
        if not self.id:
            # verifica se existem dados na tabela Documento
            d = Documento.objects.all()
            if not d:
                # se nao: gerar protocolo = 1
                k = 1
            else:
                # se sim: gerar protocolo + 1
                i = Documento.objects.all().order_by('-id')[0]
                k = i.id+1
            # inseri um formato com zeros e ano (Ex: 00000002/2014)
            self.protocolo = '%08d/%d' % (k, now.year)

            s = Status.objects.get(nome='Parado')
            self.status = s
        super(Documento, self).save(*args, **kwargs)

    def __unicode__(self):
        # return self.protocolo
        return "%s | %s | %s | %s | %s" % (
            self.protocolo,
            self.data_recebimento.strftime("%d/%m/%Y"),
            unicode(self.tipo_documento),
            unicode(self.origem),
            unicode(self.interessado))

    class Meta:
        verbose_name = _(u'Documento')
        verbose_name_plural = _(u'Documentos')

class DocumentoAnexo(models.Model):
    documento = models.ForeignKey('Documento', related_name='AneDocumentos')
    arquivo = models.FileField(_(u'Arquivo'), upload_to=u'anexos',
                               help_text='Selecione um arquivo')

    def attach_link(self):
        if self.arquivo:
            return "Baixar" % self.arquivo.url
        else:
            return "Nenhum arquivo encontrado"
    attach_link.allow_tags = True
    attach_link.short_description = _(u'Arquivo')

    def __unicode__(self):
        return self.pk

    class Meta:
        verbose_name = _(u'Anexo')
        verbose_name_plural = _(u'Anexos')

class Tramite_Documento(models.Model):
    tramite = models.ForeignKey('Tramite')
    protocolo = models.ForeignKey('Documento', related_name='doc_fk')
    folhas = models.IntegerField(_(u'Número de Folhas'), null=True, blank=True)

    class Meta:
        verbose_name = _(u'Documento do Tramite')
        verbose_name_plural = _(u'Documentos do Tramite')

class Tramite(models.Model):
    data_tramite = models.DateTimeField(_(u'Data do Tramite'),
                                        auto_now_add=True)
    origem = models.ForeignKey(Instituicao,
                               verbose_name=_(u'Instituição Origem'),
                               related_name='TraInstituicoesOri')
    origem_setor = models.ForeignKey(Setor, verbose_name=_(u'Setor Origem'),
                                     related_name='TraSetoresOri')
    destino = models.ForeignKey(Instituicao,
                                verbose_name=_(u'Instituição Destino'),
                                related_name='TraInstituicoesDes')
    destino_setor = models.ForeignKey(Setor, verbose_name=_(u'Setor Destino'),
                                      related_name='TraSetoresDes')
    motivo = models.TextField(_(u'Motivo'), blank=True, null=True)
    protocolo = models.ManyToManyField(Documento,
                                       verbose_name=_(u'Protocolos'),
                                       related_name='TraDocumentos',
                                       through=Tramite_Documento)
                                       # limit_choices_to={'status':
                                       #                   Documento.objects.exclude(
                                       #                       pk=2)})
# Funcao para filtrar os documentos com status diferente de 'Arquivado'
# Nao funciona com o m2m do grappelli, por conta do raw_id_fields

    def get_documentos(self):
        # return self.documento.all()
        # return '
' .join([k.protocolo for k in self.documento.all()])
        out = []
        for k in self.protocolo.all():
            out.append(
                # Removi essa parte porque ela inseria o link
                # para o documento e agora nem todos os documentos
                # sao editaveis, entao nao fazia sentido
                # '%s
' % (
                #     Documento._meta.verbose_name.lower(),
                #     k.pk, k.protocolo)
                '%s
' % k.protocolo
            )
        return '\n'.join(out)
    get_documentos.allow_tags = True
    get_documentos.short_description = 'Protocolos'

    def get_documentos_guia(self):
        out = []
        for k in self.protocolo.all():
            out.append(
                '
%s	%s	%s	%s	%s	 \
                %s
' % (
                k.protocolo, k.tipo_documento, k.numero, k.assunto,
                k.interessado, k.data_recebimento.strftime('%d/%m/%Y %H:%M')))
            # out.append(
            #     k.protocolo, k.tipo_documento, k.numero,
            #     k.assunto, k.interessado, k.data_recebimento)
        return ''.join(out)
        # return ''.join(out)
    get_documentos_guia.allow_tags = True

    def __unicode__(self):
        return unicode(self.pk)

    class Meta:
        verbose_name = u'Tramite'
        verbose_name_plural = u'Tramites'

def status_changed(sender, instance, **kwargs):
    """
    funcao que envia um sinal para mudar o status
    dos documentos atraves da tabela: Tramite_Documento
    por conta do InLine no admin
    """
    s = Status.objects.get(nome='Tramitando')
    if s.pk:
        instance.protocolo.status_id = s.pk
        instance.protocolo.save()

post_save.connect(status_changed,
                  sender=Tramite_Documento,
                  dispatch_uid='status_changed_tramitando')  # weak=False)

# # funcao que envia um sinal para alterar o status dos documentos atraves da
# tabela: Tramite se ela nao tiver um InLine
# start
# def save_handler(sender, instance, *args, **kwargs):
#     m2m_changed.connect(m2m_handler, sender=sender.protocolo.through,
#                         weak=False)

# def m2m_handler(sender, instance, action, *args, **kwargs):
#     if action =='post_clear':
#         for d in Documento.objects.all():
#             # instance.m2m_field.add(sub)
#             d.status_id = 1
#             d.save()

# post_save.connect(save_handler, sender=Tramite, weak=False)
# # end

# def do_something(sender, **kwargs):
    # the object which is saved can be accessed via kwargs 'instance' key.
    # obj = kwargs['instance']
    # print 'the object is now saved.'
    # ...do something else...

# here we connect a post_save signal for MyModel
# in other terms whenever an instance of MyModel is saved
# the 'do_something' function will be called.
# post_save.connect(do_something, sender=Tramite)
<commit_msg>#4721314 - Trocar Interessado por Remetente<commit_after># coding: utf-8

from django.db import models
# from django.db.models import Q
from django.db.models.signals import post_save  # , pre_save, m2m_changed
from django.utils.translation import ugettext_lazy as _

import datetime
# from datetime import date
now = datetime.datetime.now()

from protocolle.core.models import (TipoDocumento, Carater, Natureza,
                                    Status, OPERATION_CHOICES)
from protocolle.auxiliar.models import (Instituicao, Pessoa, Setor)

class Documento(models.Model):
    operacao = models.IntegerField(_(u'Operação'), choices=OPERATION_CHOICES)
    tipo_documento = models.ForeignKey(TipoDocumento,
                                       verbose_name=_(u'Tipo de Documento'),
                                       related_name='DocTiposDocumento')
    protocolo = models.CharField(_(u'Número do Protocolo'), max_length=100,
                                 unique=True, default=0, editable=False)
    numero = models.CharField(_(u'Número do Documento'), max_length=100)
    data_recebimento = models.DateTimeField(_(u'Data de Recebimento'),
                                            auto_now_add=True)
    data_documento = models.DateField(_(u'Data do Documento'))
    data_validade = models.DateField(_(u'Data de Validade'),
                                     null=True, blank=True)
    folhas = models.IntegerField(_(u'Número de Folhas'), null=True, blank=True)
    carater = models.ForeignKey(Carater, verbose_name=_(u'Caráter'),
                                related_name='DocCarateres')
    natureza = models.ForeignKey(Natureza, verbose_name=_(u'Natureza'),
                                 related_name='DocNaturezas')
    origem = models.ForeignKey(Instituicao,
                               verbose_name=_(u'Instituição Origem'),
                               related_name='DocInstituicoesOri')
    destino = models.ForeignKey(Instituicao, editable=False,
                                verbose_name=_(u'Instituição Destino'),
                                related_name='DocInstituicoesDes')
    assunto = models.TextField(_(u'Assunto'))
    interessado = models.ForeignKey(Pessoa, verbose_name=_(u'Remetente'),
                                    related_name='DocPessoasRem')
    status = models.ForeignKey(Status, verbose_name=_(u'Status'),
                               related_name='DocSituacoes', editable=False)
    observacao = models.TextField(_(u'Observação'), null=True, blank=True)

    def save(self, *args, **kwargs):
        """
        Funcao para gerar o número de protocolo
        unico e personalizado
        """
        if not self.id:
            # verifica se existem dados na tabela Documento
            d = Documento.objects.all()
            if not d:
                # se nao: gerar protocolo = 1
                k = 1
            else:
                # se sim: gerar protocolo + 1
                i = Documento.objects.all().order_by('-id')[0]
                k = i.id+1
            # inseri um formato com zeros e ano (Ex: 00000002/2014)
            self.protocolo = '%08d/%d' % (k, now.year)

            s = Status.objects.get(nome='Parado')
            self.status = s
        super(Documento, self).save(*args, **kwargs)

    def __unicode__(self):
        # return self.protocolo
        return "%s | %s | %s | %s | %s" % (
            self.protocolo,
            self.data_recebimento.strftime("%d/%m/%Y"),
            unicode(self.tipo_documento),
            unicode(self.origem),
            unicode(self.interessado))

    class Meta:
        verbose_name = _(u'Documento')
        verbose_name_plural = _(u'Documentos')

class DocumentoAnexo(models.Model):
    documento = models.ForeignKey('Documento', related_name='AneDocumentos')
    arquivo = models.FileField(_(u'Arquivo'), upload_to=u'anexos',
                               help_text='Selecione um arquivo')

    def attach_link(self):
        if self.arquivo:
            return "Baixar" % self.arquivo.url
        else:
            return "Nenhum arquivo encontrado"
    attach_link.allow_tags = True
    attach_link.short_description = _(u'Arquivo')

    def __unicode__(self):
        return self.pk

    class Meta:
        verbose_name = _(u'Anexo')
        verbose_name_plural = _(u'Anexos')

class Tramite_Documento(models.Model):
    tramite = models.ForeignKey('Tramite')
    protocolo = models.ForeignKey('Documento', related_name='doc_fk')
    folhas = models.IntegerField(_(u'Número de Folhas'), null=True, blank=True)

    class Meta:
        verbose_name = _(u'Documento do Tramite')
        verbose_name_plural = _(u'Documentos do Tramite')

class Tramite(models.Model):
    data_tramite = models.DateTimeField(_(u'Data do Tramite'),
                                        auto_now_add=True)
    origem = models.ForeignKey(Instituicao,
                               verbose_name=_(u'Instituição Origem'),
                               related_name='TraInstituicoesOri')
    origem_setor = models.ForeignKey(Setor, verbose_name=_(u'Setor Origem'),
                                     related_name='TraSetoresOri')
    destino = models.ForeignKey(Instituicao,
                                verbose_name=_(u'Instituição Destino'),
                                related_name='TraInstituicoesDes')
    destino_setor = models.ForeignKey(Setor, verbose_name=_(u'Setor Destino'),
                                      related_name='TraSetoresDes')
    motivo = models.TextField(_(u'Motivo'), blank=True, null=True)
    protocolo = models.ManyToManyField(Documento,
                                       verbose_name=_(u'Protocolos'),
                                       related_name='TraDocumentos',
                                       through=Tramite_Documento)
                                       # limit_choices_to={'status':
                                       #                   Documento.objects.exclude(
                                       #                       pk=2)})
# Funcao para filtrar os documentos com status diferente de 'Arquivado'
# Nao funciona com o m2m do grappelli, por conta do raw_id_fields

    def get_documentos(self):
        # return self.documento.all()
        # return '
' .join([k.protocolo for k in self.documento.all()])
        out = []
        for k in self.protocolo.all():
            out.append(
                # Removi essa parte porque ela inseria o link
                # para o documento e agora nem todos os documentos
                # sao editaveis, entao nao fazia sentido
                # '%s
' % (
                #     Documento._meta.verbose_name.lower(),
                #     k.pk, k.protocolo)
                '%s
' % k.protocolo
            )
        return '\n'.join(out)
    get_documentos.allow_tags = True
    get_documentos.short_description = 'Protocolos'

    def get_documentos_guia(self):
        out = []
        for k in self.protocolo.all():
            out.append(
                '
%s	%s	%s	%s	%s	 \
                %s
' % (
                k.protocolo, k.tipo_documento, k.numero, k.assunto,
                k.interessado, k.data_recebimento.strftime('%d/%m/%Y %H:%M')))
            # out.append(
            #     k.protocolo, k.tipo_documento, k.numero,
            #     k.assunto, k.interessado, k.data_recebimento)
        return ''.join(out)
        # return ''.join(out)
    get_documentos_guia.allow_tags = True

    def __unicode__(self):
        return unicode(self.pk)

    class Meta:
        verbose_name = u'Tramite'
        verbose_name_plural = u'Tramites'

def status_changed(sender, instance, **kwargs):
    """
    funcao que envia um sinal para mudar o status
    dos documentos atraves da tabela: Tramite_Documento
    por conta do InLine no admin
    """
    s = Status.objects.get(nome='Tramitando')
    if s.pk:
        instance.protocolo.status_id = s.pk
        instance.protocolo.save()

post_save.connect(status_changed,
                  sender=Tramite_Documento,
                  dispatch_uid='status_changed_tramitando')  # weak=False)

# # funcao que envia um sinal para alterar o status dos documentos atraves da
# tabela: Tramite se ela nao tiver um InLine
# start
# def save_handler(sender, instance, *args, **kwargs):
#     m2m_changed.connect(m2m_handler, sender=sender.protocolo.through,
#                         weak=False)

# def m2m_handler(sender, instance, action, *args, **kwargs):
#     if action =='post_clear':
#         for d in Documento.objects.all():
#             # instance.m2m_field.add(sub)
#             d.status_id = 1
#             d.save()

# post_save.connect(save_handler, sender=Tramite, weak=False)
# # end

# def do_something(sender, **kwargs):
    # the object which is saved can be accessed via kwargs 'instance' key.
    # obj = kwargs['instance']
    # print 'the object is now saved.'
    # ...do something else...

# here we connect a post_save signal for MyModel
# in other terms whenever an instance of MyModel is saved
# the 'do_something' function will be called.
# post_save.connect(do_something, sender=Tramite)
<|endoftext|>


<commit_before>from django.test import TestCase
from django.test.client import Client
from django.core.handlers.wsgi import WSGIRequest
from django.core.urlresolvers import reverse

from adzone.models import *
from django.contrib.auth.models import User

def datenow():
    from datetime import datetime
    return datetime.now()

class RequestFactory(Client):
    """
    Class that lets you create mock Request objects for use in testing.

    Usage:

    rf = RequestFactory()
    get_request = rf.get('/hello/')
    post_request = rf.post('/submit/', {'foo': 'bar'})

    This class re-uses the django.test.client.Client interface, docs here:
    http://www.djangoproject.com/documentation/testing/#the-test-client

    Once you have a request object you can pass it to any view function,
    just as if that view had been hooked up using a URLconf.

    """
    def request(self, **request):
        """
        Similar to parent class, but returns the request object as soon as it
        has created it.
        """
        environ = {
            'HTTP_COOKIE': self.cookies,
            'PATH_INFO': '/',
            'QUERY_STRING': '',
            'REQUEST_METHOD': 'GET',
            'SCRIPT_NAME': '',
            'SERVER_NAME': 'testserver',
            'SERVER_PORT': 80,
            'SERVER_PROTOCOL': 'HTTP/1.1',
            'REMOTE_ADDR': '127.0.0.1',
        }
        environ.update(self.defaults)
        environ.update(request)
        return WSGIRequest(environ)

class AdvertisingTestCase(TestCase):
    def setUp(self):
        self.request = RequestFactory().request()

        self.user = User.objects.create_user('test', 'test@example.com', 'testpass')

        self.advertiser = Advertiser.objects.create(
                company_name = 'teh_node Web Development',
                website = 'http://andre.smoenux.webfactional.com/',
                user = self.user)

        # Categories setup
        self.category = AdCategory.objects.create(
                title = 'Internet Services',
                slug = 'internet-services',
                description = 'All internet based services')

        self.category2 = AdCategory.objects.create(
                title = 'Category Two',
                slug = 'categorytwo',
                description = 'A Second Category')

        # Zones setup
        self.adzone = AdZone.objects.create(
                title = 'Sidebar',
                slug = 'sidebar',
                description = 'Side Bar Ads')

        self.adzone2 = AdZone.objects.create(
                title = 'Content Banner',
                slug = 'contentbanner',
                description = 'Content Adverts')

        # Ad setup
        self.ad = TextAd.objects.create(
                title = 'First Ad',
                content = 'For all your web design and development needs, at competitive rates.',
                url = 'http://www.teh-node.co.za/',
                enabled = True,
                advertiser = self.advertiser,
                category = self.category,
                zone = self.adzone)

        self.ad2 = TextAd.objects.create(
                title = 'Second Ad',
                content = 'A second advert.',
                url = 'http://www.teh-node.co.za/',
                enabled = True,
                advertiser = self.advertiser,
                category = self.category2,
                zone = self.adzone2)

        self.ad3 = TextAd.objects.create(
                title = 'A Third Ad',
                content = 'A third advert.',
                url = 'http://www.teh-node.co.za/',
                enabled = True,
                advertiser = self.advertiser,
                category = self.category2,
                zone = self.adzone2)

        # AdImpression Setup
        self.impression1 = AdImpression.objects.create(
                impression_date=datenow(),
                source_ip='127.0.0.2',
                ad=self.ad)
        self.impression2 = AdImpression.objects.create(
                impression_date=datenow(),
                source_ip='127.0.0.3',
                ad=self.ad2)

        # Clicks Setup
        self.adclick1 = AdClick.objects.create(
                click_date=datenow(),
                source_ip='127.0.0.1',
                ad=self.ad)

class AdvertiserTestCase(AdvertisingTestCase):
    def testAdvertiser(self):
        self.assertEquals(self.advertiser.get_website_url(), 'http://andre.smoenux.webfactional.com/')

class CategoryTestCase(AdvertisingTestCase):
    def testAdCategory(self):
        self.assertEquals(self.category.__unicode__(), 'Internet Services')

class ZoneTestCase(AdvertisingTestCase):
    def testAdZone(self):
        self.assertEquals(self.adzone.__unicode__(), 'Sidebar')

    def testAdinZone(self):
        ads = TextAd.objects.filter(zone__slug='sidebar')
        self.assertEquals(len(ads), 1)

class AdvertTestCase(AdvertisingTestCase):
    def testAd(self):
        self.assertEquals(reverse(
            'adzone_ad_view',
            args=['1']
        )[-8:], '/view/1/')
        adimpressions = AdImpression.objects.filter(ad=self.ad)
        self.assertEquals(len(adimpressions), 1)
        self.assertEquals(adimpressions[0].source_ip, '127.0.0.2')

    def testAdAdvertiser(self):
        self.assertEquals(self.ad.advertiser.__unicode__(), 'teh_node Web Development')
        self.assertEquals(self.ad.advertiser.company_name, 'teh_node Web Development')

    def testAddsInCategory(self):
        ads = TextAd.objects.filter(category__slug='internet-services')
        self.assertEquals(len(ads), 1)
        self.assertEquals(ads[0].title, 'First Ad')

    def testRandomAd(self):
        ad = AdBase.objects.get_random_ad(
            ad_category='internet-services',
            ad_zone='sidebar'
        )
        self.assertEquals(ad.title, 'First Ad')

class ImpressionTestCase(AdvertisingTestCase):
    def testImpression(self):
        impressions = AdImpression.objects.all()
        self.assertEquals(len(impressions), 2)

class ClickTestCase(AdvertisingTestCase):
    def testClicks(self):
        clicks = AdClick.objects.all()
        self.assertEquals(len(clicks), 1)

    def testClickOnAds(self):
        c = Client(REMOTE_ADDR='127.0.0.1')
        response = c.get(reverse(
            'adzone_ad_view',
            args=['1']
        ))
        self.assertEquals(len(AdClick.objects.all()), 2)
        click = AdClick.objects.all()[1]
        self.assertEquals(click.source_ip, '127.0.0.1')

    def testInvalidAdURL(self):
        c = Client(REMOTE_ADDR='127.0.0.1')
        response = c.get('/te/10')
        self.assertEquals(len(AdClick.objects.all()), 1)
<commit_msg>Refactoring the tests<commit_after>from datetime import datetime

from django.test import TestCase
# from django.test.client import RequestFactory
# from django.core.urlresolvers import reverse
from django.contrib.auth.models import User

from adzone.models import Advertiser, AdCategory, AdZone, AdBase

user = lambda: User.objects.create_user('test', 'test@example.com', 'secret')

def datenow():
    return datetime.now()

def create_objects():
    advertiser = Advertiser.objects.create(
        company_name='Advertiser Name 1',
        website='http://example.com/', user=user())

    category = AdCategory.objects.create(
        title='Internet Services',
        slug='internet-services',
        description='All internet based services')

    adzone = AdZone.objects.create(
        title='Sidebar',
        slug='sidebar',
        description='Sidebar Zone Description')

    return advertiser, category, adzone

    # ad = TextAd.objects.create(
    #     title='First Ad',
    #     content='For all your web design and development needs, at competitive rates.',
    #     url='http://www.teh-node.co.za/',
    #     enabled=True,
    #     advertiser=self.advertiser,
    #     category=self.category,
    #     zone=self.adzone)

    # ad2 = TextAd.objects.create(
    #     title='Second Ad',
    #     content='A second advert.',
    #     url='http://www.teh-node.co.za/',
    #     enabled=True,
    #     advertiser=self.advertiser,
    #     category=self.category2,
    #     zone=self.adzone2)

    # ad3 = TextAd.objects.create(
    #     title='A Third Ad',
    #     content='A third advert.',
    #     url='http://www.teh-node.co.za/',
    #     enabled=True,
    #     advertiser=self.advertiser,
    #     category=self.category2,
    #     zone=self.adzone2)

    # # AdImpression Setup
    # impression1 = AdImpression.objects.create(
    #     impression_date=datenow(),
    #     source_ip='127.0.0.2',
    #     ad=self.ad)

    # impression2 = AdImpression.objects.create(
    #     impression_date=datenow(),
    #     source_ip='127.0.0.3',
    #     ad=self.ad2)

    # # Clicks Setup
    # adclick1 = AdClick.objects.create(
    #     click_date=datenow(),
    #     source_ip='127.0.0.1',
    #     ad=self.ad)

class AdvertiserTestCase(TestCase):

    def test_model(self):
        Advertiser(
            company_name='Advertiser Name 1',
            website='http://example.com/',
            user=user())

    def test_get_website_url(self):
        advertiser = Advertiser(
            company_name='Advertiser Name 1',
            website='http://example.com/',
            user=user())

        self.assertEqual(
            'http://example.com/',
            advertiser.get_website_url())

class AdCategoryTestCase(TestCase):

    def test_model(self):
        AdCategory(
            title='Internet Services',
            slug='internet-services',
            description='All internet based services')

class AdZoneTestCase(TestCase):

    def test_model(self):
        AdZone(
            title='Ad Zone Title',
            slug='adzone',
            description='Ad Zone Description')

class AdBaseTestCase(TestCase):

    def test_model(self):
        advertiser, category, zone = create_objects()
        AdBase(
            title='Ad Title',
            url='www.example.com',
            advertiser=advertiser,
            category=category,
            zone=zone
        )

# class AdvertTestCase(AdvertisingTestCase):
#     def testAd(self):
#         self.assertEquals(reverse(
#             'adzone_ad_view',
#             args=['1']
#         )[-8:], '/view/1/')
#         adimpressions = AdImpression.objects.filter(ad=self.ad)
#         self.assertEquals(len(adimpressions), 1)
#         self.assertEquals(adimpressions[0].source_ip, '127.0.0.2')
#
#     def testAdAdvertiser(self):
#         self.assertEquals(self.ad.advertiser.__unicode__(), 'teh_node Web Development')
#         self.assertEquals(self.ad.advertiser.company_name, 'teh_node Web Development')
#
#     def testAddsInCategory(self):
#         ads = TextAd.objects.filter(category__slug='internet-services')
#         self.assertEquals(len(ads), 1)
#         self.assertEquals(ads[0].title, 'First Ad')
#
#     def testRandomAd(self):
#         ad = AdBase.objects.get_random_ad(
#             ad_category='internet-services',
#             ad_zone='sidebar'
#         )
#         self.assertEquals(ad.title, 'First Ad')
#
#
# class ImpressionTestCase(AdvertisingTestCase):
#     def testImpression(self):
#         impressions = AdImpression.objects.all()
#         self.assertEquals(len(impressions), 2)
#
#
# class ClickTestCase(AdvertisingTestCase):
#
#     def testClicks(self):
#         clicks = AdClick.objects.all()
#         self.assertEquals(len(clicks), 1)
#
#     def testClickOnAds(self):
#         c = Client(REMOTE_ADDR='127.0.0.1')
#         response = c.get(reverse(
#             'adzone_ad_view',
#             args=['1']
#         ))
#         self.assertEquals(len(AdClick.objects.all()), 2)
#         click = AdClick.objects.all()[1]
#         self.assertEquals(click.source_ip, '127.0.0.1')
#
#     def testInvalidAdURL(self):
#         c = Client(REMOTE_ADDR='127.0.0.1')
#         response = c.get('/te/10')
#         self.assertEquals(len(AdClick.objects.all()), 1)
<|endoftext|>


<commit_before>from ..utils import cleanup_name, asciireplace, limit_text
from ..exceptions.messages import RowObjectMsg as msg
from .formulas import Formula

from tabulate import tabulate

class BaseRow(dict):
    """This Base Class represents a row in a spreadsheet"""

    __slots__ = ["__output__"]

    def __init__(self, data, *args, **kwargs):
        super(BaseRow, self).__init__(data)
        self.construct(*args, **kwargs)
        self.__output__ = True

    def construct(self, *args, **kwargs):
        pass

    def formula(self, columnname, func, rowref=None, **kwargs):
        if rowref is None:
            rowref = self
        self.setcolumn(columnname, Formula(func, rowref, **kwargs))

    def getformula(self, columnname):
        """Get Formula"""
        try:
            return self.getcolumn(columnname).call()
        except ValueError:
            raise ValueError(
                msg.getformulamsg.format(columnname))

    @property
    def outputrow(self):
        return self.__output__

    @outputrow.setter
    def outputrow(self, v):
        if not isinstance(v, bool):
            raise TypeError(msg.outputrowmsg.format(bool, type(v)))
        self.__output__ = v

    def getcolumn(self, v):
        "Get column by the orginal column name"
        s = cleanup_name(v)
        if s in self.keys():
            return getattr(self, s)
        else:
            raise KeyError("{}".format(v))

    def setcolumn(self, a, v):
        s = cleanup_name(a)
        if s in self.keys():
            self.__setattr__(s, v)
        else:
            raise Exception("{}".format(a))

    def delcolumn(self, v):
        s = cleanup_name(a)
        if s in self.keys():
            self.__delattr__(s, v)
        else:
            raise KeyError("{}".format(a))

    def __repr__(self):
        return "<'{rowname}':{columnamount}>".format(
            rowname=self.__class__.__name__,
            columnamount=len(self.keys())
            )

    def __str__(self):
        data = self.longcolumn()
        sortedcolumns = sorted(data)

        return tabulate(
            [sortedcolumns] + [[data[c] for c in sortedcolumns]],
            headers="firstrow",
            tablefmt="grid")

    def __pos__(self):
        self.outputrow = True
        return self

    def __neg__(self):
        self.outputrow = False
        return self

    def __invert__(self):
        self.outputrow = not (self.outputrow)
        return self

    def __getattribute__(self, attr):
        if not (attr in super(BaseRow, self).keys()):
            return super(dict, self).__getattribute__(attr)
        else:
            return (self[attr]["value"])

    def __getattr__(self, attr):
        s = cleanup_name(attr)
        if s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                '\'{}\' has no attribute \'{}\''.format(
                    type(self), attr),
                ". However, '{s}' is an existing condensed ".format(s=s) +
                "column name. Only the condensed version is supported."
                .format(s=s)
                )))
        else:
            raise AttributeError('\'{}\' has no attribute \'{}\''.format(
        type(self), attr))

    def __setattr__(self, attr, v):
        """Allows setting new data to row
        is called if row does not have specially
        defined set behavior"""
        s = cleanup_name(attr)
        if attr in super(BaseRow, self).keys():
            self[attr]["value"] = v
        elif s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                '\'{}\' has no attribute \'{}\''.format(
                    type(self), attr),
                ". However, '{s}' is an existing condensed ".format(s=s) +
                "column name. Only the condensed version is supported."
                .format(s=s)
                )))
        else:
            super(BaseRow, self).__setattr__(attr, v)

    def __delattr__(self, attr):
        """Allows deletion of rows and attributes (Makes a row empty) by using
        del statement"""
        s = cleanup_name(attr)
        if attr in super(BaseRow, self).keys():
            self[attr]["value"] = ""
        elif s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                '\'{}\' has no attribute \'{}\''.format(
                    type(self), attr),
                ". However, '{s}' is an existing condensed ".format(s=s) +
                "column name. Only the condensed version is supported."
                .format(s=s)
                )))
        else:
            super(BaseRow, self).__delattr__(attr)

    def addcolumn(self, columnname, columndata=""):
        """Adds a column"""
        short_cn = cleanup_name(columnname)
        if not self.get(short_cn):
            self[short_cn] = {"org_name":columnname, "value":columndata}
        else:
            raise Exception("Column already exists.")

    def longcolumn(self, columns=None):
        """Generates a Dict that uses orginal names of
        the column, used for output"""
        newdict = {}
        if columns:
            shortcolumns_check = [cleanup_name(x) for x in columns]
        for k in self.keys():
            if columns:
                if not (k in shortcolumns_check):
                    continue
            newdict.update({self[k]["org_name"]: self[k]["value"]})

        return newdict

    def tabulate(self, format="grid", only_ascii=True, columns=None, text_limit=None):
        data = self.longcolumn()
        sortedcolumns = sorted(data) if not columns else columns

        result = tabulate(
            [sortedcolumns] + [[limit_text(data[c], text_limit) for c in sortedcolumns]],
            headers="firstrow",
            tablefmt=format)
        if only_ascii:
            return asciireplace(result)
        else:
            return result<commit_msg>Removed Default Tabulation for row objects.<commit_after>from ..utils import cleanup_name, asciireplace, limit_text
from ..exceptions.messages import RowObjectMsg as msg
from .formulas import Formula

from tabulate import tabulate

class BaseRow(dict):
    """This Base Class represents a row in a spreadsheet"""

    __slots__ = ["__output__"]

    def __init__(self, data, *args, **kwargs):
        super(BaseRow, self).__init__(data)
        self.construct(*args, **kwargs)
        self.__output__ = True

    def construct(self, *args, **kwargs):
        pass

    def formula(self, columnname, func, rowref=None, **kwargs):
        if rowref is None:
            rowref = self
        self.setcolumn(columnname, Formula(func, rowref, **kwargs))

    def getformula(self, columnname):
        """Get Formula"""
        try:
            return self.getcolumn(columnname).call()
        except ValueError:
            raise ValueError(
                msg.getformulamsg.format(columnname))

    @property
    def outputrow(self):
        return self.__output__

    @outputrow.setter
    def outputrow(self, v):
        if not isinstance(v, bool):
            raise TypeError(msg.outputrowmsg.format(bool, type(v)))
        self.__output__ = v

    def getcolumn(self, v):
        "Get column by the orginal column name"
        s = cleanup_name(v)
        if s in self.keys():
            return getattr(self, s)
        else:
            raise KeyError("{}".format(v))

    def setcolumn(self, a, v):
        s = cleanup_name(a)
        if s in self.keys():
            self.__setattr__(s, v)
        else:
            raise Exception("{}".format(a))

    def delcolumn(self, v):
        s = cleanup_name(a)
        if s in self.keys():
            self.__delattr__(s, v)
        else:
            raise KeyError("{}".format(a))

    def __repr__(self):
        return "<'{rowname}':{columnamount}>".format(
            rowname=self.__class__.__name__,
            columnamount=len(self.keys())
            )

    def __str__(self):
        return "<'{rowname}':{columnamount}>".format(
            rowname=self.__class__.__name__,
            columnamount=len(self.keys())
            )

    def __pos__(self):
        self.outputrow = True
        return self

    def __neg__(self):
        self.outputrow = False
        return self

    def __invert__(self):
        self.outputrow = not (self.outputrow)
        return self

    def __getattribute__(self, attr):
        if not (attr in super(BaseRow, self).keys()):
            return super(dict, self).__getattribute__(attr)
        else:
            return (self[attr]["value"])

    def __getattr__(self, attr):
        s = cleanup_name(attr)
        if s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                '\'{}\' has no attribute \'{}\''.format(
                    type(self), attr),
                ". However, '{s}' is an existing condensed ".format(s=s) +
                "column name. Only the condensed version is supported."
                .format(s=s)
                )))
        else:
            raise AttributeError('\'{}\' has no attribute \'{}\''.format(
        type(self), attr))

    def __setattr__(self, attr, v):
        """Allows setting new data to row
        is called if row does not have specially
        defined set behavior"""
        s = cleanup_name(attr)
        if attr in super(BaseRow, self).keys():
            self[attr]["value"] = v
        elif s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                '\'{}\' has no attribute \'{}\''.format(
                    type(self), attr),
                ". However, '{s}' is an existing condensed ".format(s=s) +
                "column name. Only the condensed version is supported."
                .format(s=s)
                )))
        else:
            super(BaseRow, self).__setattr__(attr, v)

    def __delattr__(self, attr):
        """Allows deletion of rows and attributes (Makes a row empty) by using
        del statement"""
        s = cleanup_name(attr)
        if attr in super(BaseRow, self).keys():
            self[attr]["value"] = ""
        elif s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                '\'{}\' has no attribute \'{}\''.format(
                    type(self), attr),
                ". However, '{s}' is an existing condensed ".format(s=s) +
                "column name. Only the condensed version is supported."
                .format(s=s)
                )))
        else:
            super(BaseRow, self).__delattr__(attr)

    def addcolumn(self, columnname, columndata=""):
        """Adds a column"""
        short_cn = cleanup_name(columnname)
        if not self.get(short_cn):
            self[short_cn] = {"org_name":columnname, "value":columndata}
        else:
            raise Exception("Column already exists.")

    def longcolumn(self, columns=None):
        """Generates a Dict that uses orginal names of
        the column, used for output"""
        newdict = {}
        if columns:
            shortcolumns_check = [cleanup_name(x) for x in columns]
        for k in self.keys():
            if columns:
                if not (k in shortcolumns_check):
                    continue
            newdict.update({self[k]["org_name"]: self[k]["value"]})

        return newdict

    def tabulate(self, format="grid", only_ascii=True, columns=None, text_limit=None):
        data = self.longcolumn()
        sortedcolumns = sorted(data) if not columns else columns

        result = tabulate(
            [sortedcolumns] + [[limit_text(data[c], text_limit) for c in sortedcolumns]],
            headers="firstrow",
            tablefmt=format)
        if only_ascii:
            return asciireplace(result)
        else:
            return result<|endoftext|>


<commit_before><commit_msg>Nest with-statements<commit_after><|endoftext|>


<commit_before>"""
strategy.py

Defines a class that contains a strategy definition and the functions to trade on that strategy.
"""
class Strategy(object):
    """
    Defines a particular trade strategy
    """

    def __init__(self, portfolio, commission_min=1.00, commission=0.0075, buy_percent=1.0, sell_percent=1.0):
        """create base Strategy class"""

        self.commission_min = commission_min
        self.commission = commission
        self.portfolio = portfolio
        self.buy_percent = buy_percent
        self.sell_percent = sell_percent
        self.long_open = {symbol:False for symbol in portfolio.assets.keys()}
        self.short_open = {symbol:False for symbol in portfolio.assets.keys()}
        self.display_data = []

        recordings = [
            'buy price', 'buy shares', 'buy fees', 'buy date',
            'sell price', 'sell shares', 'sell fees', 'sell date',
            'gain', 'profit', 'loss', 'return', 'win/loose',
            'min', 'min date', 'max', 'max date',
            'drawdown', 'drawdown days'
        ]
        self.record = {symbol:pd.DataFrame(columns=recordings) for symbol in portfolio.assets.keys()}
        self.max = {symbol:[0, None] for symbol in portfolio.assets.keys()}
        self.min = {symbol:[999999999999999, None] for symbol in portfolio.assets.keys()}
        self.drawdown = {symbol:[999999999999999, None] for symbol in portfolio.assets.keys()}

    def on_date(self, date, i):
        """called for each date of portfolio data, implements trading logic"""
        pass

    def enter_long(self, symbol, date, shares):

        self.long_open[symbol] = True
        self.portfolio.trade(symbol, date, shares, self.commission_min, self.commission)

        # set min, max & drawdown
        self.max[symbol] = [self.portfolio.assets[symbol].c[date], date]
        self.min[symbol] = [self.portfolio.assets[symbol].c[date], date]
        self.drawdown[symbol] = [self.portfolio.assets[symbol].c[date], date]

        record = {
            'buy price' : [self.portfolio.assets[symbol].c[date]],
            'buy shares' : [shares],
            'buy fees' : [max(self.commission_min, abs(self.commission * shares))],
            'buy date' : [date],
            'sell price' : [None],
            'sell shares' : [None],
            'sell fees' : [None],
            'sell date' : [None],
            'gain' : [None],
            'profit' : [None],
            'loss' : [None],
            'return' : [None],
            'win/loose' : [None],
            'min':[None],
            'max':[None],
            'min date':[None],
            'max date':[None],
            'drawdown':[None],
            'drawdown days':[None]
        }

        self.record[symbol] = self.record[symbol].append(pd.DataFrame(record), ignore_index=True)

    def exit_long(self, symbol, date, shares):

        self.long_open[symbol] = False
        self.portfolio.trade(symbol, date, -1.0 * shares, self.commission_min, self.commission)

        i = self.record[symbol].index[-1]

        self.record[symbol].loc[i, 'sell price'] = self.portfolio.assets[symbol].c[date]
        self.record[symbol].loc[i, 'sell shares'] = shares
        self.record[symbol].loc[i, 'sell fees'] =  max(self.commission_min, abs(self.commission * shares))
        self.record[symbol].loc[i, 'sell date'] = date

        initial = self.record[symbol].loc[i, 'buy price'] * self.record[symbol].loc[i, 'buy shares']
        final = self.record[symbol].loc[i, 'sell price'] * self.record[symbol].loc[i, 'sell shares']
        fees = self.record[symbol].loc[i, 'sell fees'] + self.record[symbol].loc[i, 'sell fees']
        gain = (final - initial)

        self.record[symbol].loc[i, 'gain'] = gain
        self.record[symbol].loc[i, 'return'] = gain / initial

        self.record[symbol].loc[i, 'profit'] = gain if gain > 0 else 0
        self.record[symbol].loc[i, 'loss'] = (abs(gain) + fees) if gain < 0 else fees

        if gain > 0:
            self.record[symbol].loc[i, 'win/loose'] = 'w'
        elif gain < 0:
            self.record[symbol].loc[i, 'win/loose'] = 'l'
        else:
            self.record[symbol].loc[i, 'win/loose'] = '-'

        self.record[symbol].loc[i, 'min'] = self.min[symbol][0] * self.record[symbol].loc[i, 'buy shares']
        self.record[symbol].loc[i, 'max'] = self.max[symbol][0] * self.record[symbol].loc[i, 'buy shares']

        self.record[symbol].loc[i, 'min date'] = self.min[symbol][1]
        self.record[symbol].loc[i, 'max date'] = self.max[symbol][1]

        self.record[symbol].loc[i, 'drawdown'] = (self.max[symbol][0] - self.drawdown[symbol][0]) * self.record[symbol].loc[i, 'buy shares']
        self.record[symbol].loc[i, 'drawdown days'] = (self.drawdown[symbol][1] - self.max[symbol][1]).days

    def enter_short(self):
        self.short_open[symbol] = True
        pass

    def exit_short(self):
        self.short_open[symbol] = False
        pass

    def run(self):
        """iterates over data"""

        self.before_run()

        for date in self.portfolio.cash.index:
            self.on_date(date)

        self.after_run()

    def before_run(self):
        """called at the start of run"""
        self.display_data = []

    def after_run(self):
        """called at the end of run"""

        print tabulate.tabulate(self.display_data, headers=['trade', 'symbol', 'date', 'price', 'shares', 'value', 'cash'])
        print
        performance = self.calc_performance()

        assets = [' '] + performance.keys()
        perf = [performance.values()[0].keys()] + [[p for p in asset_performance.values()] for asset_performance in performance.values()]
        perf = [list(x) for x in zip(*perf)]

        print tabulate.tabulate(perf, headers=assets, floatfmt=".2f")

    def calc_performance(self):
        """calculate performance"""

        performance = {}
        for symbol in self.portfolio.assets.keys():

            trades = len(self.record[symbol])
            profit = self.record[symbol]['profit'].sum()
            loss = self.record[symbol]['loss'].sum()
            try:
                wins = len(self.record[symbol].groupby('win/loose').groups['w'])
            except:
                wins = 0
            try:
                losses = len(self.record[symbol].groupby('win/loose').groups['l'])
            except:
                losses = 0
            try:
                washes = len(self.record[symbol].groupby('win/loose').groups['-'])
            except:
                washes = 0
            max_drawdown = self.record[symbol]['drawdown'].max()
            average_drawdown = self.record[symbol]['drawdown'].mean()
            max_drawdown_time = self.record[symbol]['drawdown days'].max()
            average_drawdown_time = self.record[symbol]['drawdown days'].mean()

            performance[symbol] = {
                'trades': trades,
                'wins': wins,
                'losses': losses,
                'washes': washes,
                'profit': profit,
                'loss': loss,
                'net_profit': profit - loss,
                'profit_factor': profit / loss,
                'percent_profitable': wins / trades,
                'average_trade_net_profit' : (profit - loss) / trades,
                'max_drawdown' : max_drawdown,
                'average_drawdown' : average_drawdown,
                'max_drawdown_days' : max_drawdown_time,
                'average_drawdown_days' : average_drawdown_time
            }

        return performance

################################################################################################################################

class Sma_Xover(Strategy):
    """
    buy when sma(50) crosses above sma(200), sell when sma(50) crosses below sma(200)
    """

    def on_date(self, date):
        """iterates over dates"""

        for symbol in self.portfolio.assets.keys():

            if len(self.portfolio.assets[symbol].c[:date]) > 200:

                price = self.portfolio.assets[symbol].c[date]

                # record min, max & drawdown
                if price > self.max[symbol][0]:
                    self.max[symbol] = [price, date]
                    # reset drawdown
                    self.drawdown[symbol] = [price, date]

                if price < self.min[symbol][0]:
                    self.min[symbol] = [price, date]

                if price < self.drawdown[symbol][0]:
                    self.drawdown[symbol] = [price, date]

                # calculate moving averages
                sma50 = sma(self.portfolio.assets[symbol].c[:date], 50).iloc[-2:]
                sma200 = sma(self.portfolio.assets[symbol].c[:date], 200).iloc[-2:]

                # determine crossover and direction
                xover = np.sign(np.sign(sma50 - sma200).diff()).iloc[-1]

                if xover > 0 and not self.long_open[symbol]:

                    shares = (self.buy_percent * self.portfolio.cash[date]) / price
                    self.enter_long(symbol, date, shares)

                    self.display_data.append(['buy', symbol, str(date.date()), price, shares, price * shares, self.portfolio.cash[date]])

                elif (xover < 0 and self.long_open[symbol]) or (date == self.portfolio.cash.index[-1] and self.long_open[symbol]):
                    shares = self.sell_percent * self.portfolio.positions[symbol][date]
                    self.exit_long(symbol, date, shares)

                    self.display_data.append(['sell', symbol, str(date.date()), price, shares, price * shares, self.portfolio.cash[date]])

class BuyAndHold(Strategy):
    """
    buy on the first day, never sell
    """

    def on_date(self, date):
        """iterates over dates"""

        for symbol in self.portfolio.assets.keys():

            if date == self.portfolio.assets[symbol].c.index[0]:

                price = self.portfolio.assets[symbol].c[date]
                shares = (self.buy_percent * self.portfolio.cash[date]) / price
                self.enter_long(symbol, date, shares)

                self.display_data.append(['buy', symbol, str(date.date()), price, shares, price * shares, self.portfolio.cash[date]])
<commit_msg>add comments<commit_after>"""
strategy.py

Defines a class that contains a strategy definition and the functions to trade on that strategy.
"""
class Strategy(object):
    """
    Defines a particular trade strategy
    """

    def __init__(self, portfolio, commission_min=1.00, commission=0.0075, buy_percent=1.0, sell_percent=1.0):
        """create base Strategy class"""

        self.commission_min = commission_min
        self.commission = commission
        self.portfolio = portfolio
        self.buy_percent = buy_percent
        self.sell_percent = sell_percent
        self.long_open = {symbol:False for symbol in portfolio.assets.keys()}
        self.short_open = {symbol:False for symbol in portfolio.assets.keys()}
        self.display_data = []

        recordings = [
            'buy price', 'buy shares', 'buy fees', 'buy date',
            'sell price', 'sell shares', 'sell fees', 'sell date',
            'gain', 'profit', 'loss', 'return', 'win/loose',
            'min', 'min date', 'max', 'max date',
            'drawdown', 'drawdown days'
        ]
        self.record = {symbol:pd.DataFrame(columns=recordings) for symbol in portfolio.assets.keys()}
        self.max = {symbol:[0, None] for symbol in portfolio.assets.keys()}
        self.min = {symbol:[999999999999999, None] for symbol in portfolio.assets.keys()}
        self.drawdown = {symbol:[999999999999999, None] for symbol in portfolio.assets.keys()}

    def on_date(self, date, i):
        """called for each date of portfolio data, implements trading logic"""
        pass

    def enter_long(self, symbol, date, shares):

        self.long_open[symbol] = True
        self.portfolio.trade(symbol, date, shares, self.commission_min, self.commission)

        # set min, max & drawdown
        self.max[symbol] = [self.portfolio.assets[symbol].c[date], date]
        self.min[symbol] = [self.portfolio.assets[symbol].c[date], date]
        self.drawdown[symbol] = [self.portfolio.assets[symbol].c[date], date]

        record = {
            'buy price' : [self.portfolio.assets[symbol].c[date]],
            'buy shares' : [shares],
            'buy fees' : [max(self.commission_min, abs(self.commission * shares))],
            'buy date' : [date],
            'sell price' : [None],
            'sell shares' : [None],
            'sell fees' : [None],
            'sell date' : [None],
            'gain' : [None],
            'profit' : [None],
            'loss' : [None],
            'return' : [None],
            'win/loose' : [None],
            'min':[None],
            'max':[None],
            'min date':[None],
            'max date':[None],
            'drawdown':[None],
            'drawdown days':[None]
        }

        self.record[symbol] = self.record[symbol].append(pd.DataFrame(record), ignore_index=True)

    def exit_long(self, symbol, date, shares):

        self.long_open[symbol] = False
        self.portfolio.trade(symbol, date, -1.0 * shares, self.commission_min, self.commission)

        i = self.record[symbol].index[-1]

        self.record[symbol].loc[i, 'sell price'] = self.portfolio.assets[symbol].c[date]
        self.record[symbol].loc[i, 'sell shares'] = shares
        self.record[symbol].loc[i, 'sell fees'] =  max(self.commission_min, abs(self.commission * shares))
        self.record[symbol].loc[i, 'sell date'] = date

        initial = self.record[symbol].loc[i, 'buy price'] * self.record[symbol].loc[i, 'buy shares']
        final = self.record[symbol].loc[i, 'sell price'] * self.record[symbol].loc[i, 'sell shares']
        fees = self.record[symbol].loc[i, 'sell fees'] + self.record[symbol].loc[i, 'sell fees']
        gain = (final - initial)

        self.record[symbol].loc[i, 'gain'] = gain
        self.record[symbol].loc[i, 'return'] = gain / initial

        self.record[symbol].loc[i, 'profit'] = gain if gain > 0 else 0
        self.record[symbol].loc[i, 'loss'] = (abs(gain) + fees) if gain < 0 else fees

        if gain > 0:
            self.record[symbol].loc[i, 'win/loose'] = 'w'
        elif gain < 0:
            self.record[symbol].loc[i, 'win/loose'] = 'l'
        else:
            self.record[symbol].loc[i, 'win/loose'] = '-'

        self.record[symbol].loc[i, 'min'] = self.min[symbol][0] * self.record[symbol].loc[i, 'buy shares']
        self.record[symbol].loc[i, 'max'] = self.max[symbol][0] * self.record[symbol].loc[i, 'buy shares']

        self.record[symbol].loc[i, 'min date'] = self.min[symbol][1]
        self.record[symbol].loc[i, 'max date'] = self.max[symbol][1]

        self.record[symbol].loc[i, 'drawdown'] = (self.max[symbol][0] - self.drawdown[symbol][0]) * self.record[symbol].loc[i, 'buy shares']
        self.record[symbol].loc[i, 'drawdown days'] = (self.drawdown[symbol][1] - self.max[symbol][1]).days

    def enter_short(self):
        self.short_open[symbol] = True
        pass

    def exit_short(self):
        self.short_open[symbol] = False
        pass

    def run(self):
        """iterates over data"""

        self.before_run()

        for date in self.portfolio.cash.index:
            self.on_date(date)

        self.after_run()

    def before_run(self):
        """called at the start of run"""
        self.display_data = []

    def after_run(self):
        """called at the end of run"""

        # display trading data
        print tabulate.tabulate(self.display_data, headers=['trade', 'symbol', 'date', 'price', 'shares', 'value', 'cash'])
        print

        # calculate performance
        performance = self.calc_performance()
        assets = [' '] + performance.keys()
        perf = [performance.values()[0].keys()] + [[p for p in asset_performance.values()] for asset_performance in performance.values()]
        perf = [list(x) for x in zip(*perf)]

        # display performance
        print tabulate.tabulate(perf, headers=assets, floatfmt=".2f")

    def calc_performance(self):
        """calculate performance"""

        performance = {}
        for symbol in self.portfolio.assets.keys():

            trades = len(self.record[symbol])
            profit = self.record[symbol]['profit'].sum()
            loss = self.record[symbol]['loss'].sum()
            try:
                wins = len(self.record[symbol].groupby('win/loose').groups['w'])
            except:
                wins = 0
            try:
                losses = len(self.record[symbol].groupby('win/loose').groups['l'])
            except:
                losses = 0
            try:
                washes = len(self.record[symbol].groupby('win/loose').groups['-'])
            except:
                washes = 0
            max_drawdown = self.record[symbol]['drawdown'].max()
            average_drawdown = self.record[symbol]['drawdown'].mean()
            max_drawdown_time = self.record[symbol]['drawdown days'].max()
            average_drawdown_time = self.record[symbol]['drawdown days'].mean()

            performance[symbol] = {
                'trades': trades,
                'wins': wins,
                'losses': losses,
                'washes': washes,
                'profit': profit,
                'loss': loss,
                'net_profit': profit - loss,
                'profit_factor': profit / loss,
                'percent_profitable': wins / trades,
                'average_trade_net_profit' : (profit - loss) / trades,
                'max_drawdown' : max_drawdown,
                'average_drawdown' : average_drawdown,
                'max_drawdown_days' : max_drawdown_time,
                'average_drawdown_days' : average_drawdown_time
            }

        return performance

################################################################################################################################

class Sma_Xover(Strategy):
    """
    buy when sma(50) crosses above sma(200), sell when sma(50) crosses below sma(200)
    """

    def on_date(self, date):
        """iterates over dates"""

        for symbol in self.portfolio.assets.keys():

            if len(self.portfolio.assets[symbol].c[:date]) > 200:

                price = self.portfolio.assets[symbol].c[date]

                # record min, max & drawdown
                if price > self.max[symbol][0]:
                    self.max[symbol] = [price, date]
                    # reset drawdown
                    self.drawdown[symbol] = [price, date]

                if price < self.min[symbol][0]:
                    self.min[symbol] = [price, date]

                if price < self.drawdown[symbol][0]:
                    self.drawdown[symbol] = [price, date]

                # calculate moving averages
                sma50 = sma(self.portfolio.assets[symbol].c[:date], 50).iloc[-2:]
                sma200 = sma(self.portfolio.assets[symbol].c[:date], 200).iloc[-2:]

                # determine crossover and direction
                xover = np.sign(np.sign(sma50 - sma200).diff()).iloc[-1]

                if xover > 0 and not self.long_open[symbol]:

                    shares = (self.buy_percent * self.portfolio.cash[date]) / price
                    self.enter_long(symbol, date, shares)

                    self.display_data.append(['buy', symbol, str(date.date()), price, shares, price * shares, self.portfolio.cash[date]])

                elif (xover < 0 and self.long_open[symbol]) or (date == self.portfolio.cash.index[-1] and self.long_open[symbol]):
                    shares = self.sell_percent * self.portfolio.positions[symbol][date]
                    self.exit_long(symbol, date, shares)

                    self.display_data.append(['sell', symbol, str(date.date()), price, shares, price * shares, self.portfolio.cash[date]])

class BuyAndHold(Strategy):
    """
    buy on the first day, never sell
    """

    def on_date(self, date):
        """iterates over dates"""

        for symbol in self.portfolio.assets.keys():

            if date == self.portfolio.assets[symbol].c.index[0]:

                price = self.portfolio.assets[symbol].c[date]
                shares = (self.buy_percent * self.portfolio.cash[date]) / price
                self.enter_long(symbol, date, shares)

                self.display_data.append(['buy', symbol, str(date.date()), price, shares, price * shares, self.portfolio.cash[date]])
<|endoftext|>


<commit_before>from __future__ import absolute_import, division, print_function

import subprocess
from os.path import islink, isfile
import sys

NO_EXT = (
    '.py', '.pyc', '.pyo', '.h', '.a', '.c', '.txt', '.html',
    '.xml', '.png', '.jpg', '.gif', '.class',
)

MAGIC = {
    b'\xca\xfe\xba\xbe': 'MachO-universal',
    b'\xce\xfa\xed\xfe': 'MachO-i386',
    b'\xcf\xfa\xed\xfe': 'MachO-x86_64',
    b'\xfe\xed\xfa\xce': 'MachO-ppc',
    b'\xfe\xed\xfa\xcf': 'MachO-ppc64',
}

FILETYPE = {
    b'\x01\x00\x00\x00': 'MH_OBJECT',
    b'\x02\x00\x00\x00': 'MH_EXECUTE',
    b'\x03\x00\x00\x00': 'MH_FVMLIB',
    b'\x04\x00\x00\x00': 'MH_CORE',
    b'\x05\x00\x00\x00': 'MH_PRELOAD',
    b'\x06\x00\x00\x00': 'MH_DYLIB',
    b'\x07\x00\x00\x00': 'MH_DYLINKER',
    b'\x08\x00\x00\x00': 'MH_BUNDLE',
    b'\x09\x00\x00\x00': 'MH_DYLIB_STUB',
    b'\x0a\x00\x00\x00': 'MH_DSYM',
    b'\x0b\x00\x00\x00': 'MH_KEXT_BUNDLE',
}

def is_macho(path):
    if path.endswith(NO_EXT) or islink(path) or not isfile(path):
        return False
    with open(path, 'rb') as fi:
        head = fi.read(4)
    return bool(head in MAGIC)

def is_dylib(path):
    with open(path, 'rb') as fi:
        # file type indicated by fourth 32-bit constant in the mach header
        header = fi.read(16)[-4:]
        return header in FILETYPE and FILETYPE[header] == 'MH_DYLIB'

def otool(path):
    "thin wrapper around otool -L"
    lines = subprocess.check_output(['otool', '-L', path]).decode('utf-8').splitlines()
    assert lines[0].startswith(path), path
    res = []
    for line in lines[1:]:
        assert line[0] == '\t'
        res.append(line.split()[0])
    return res

def install_name_change(path, cb_func):
    """
    change dynamic shared library install names of Mach-O binary `path`.

    `cb_func` is a callback function which called for each shared library name.
    It is called with `path` and the current shared library install name,
    and return the new name (or None if the name should be unchanged).
    """
    changes = []
    for link in otool(path):
        new_link = cb_func(path, link)
        if new_link:
            changes.append((link, new_link))

    ret = True
    for old, new in changes:
        args = ['install_name_tool', '-change', old, new, path]
        print(' '.join(args))
        p = subprocess.Popen(args, stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        stderr = stderr.decode('utf-8')
        if "Mach-O dynamic shared library stub file" in stderr:
            print("Skipping Mach-O dynamic shared library stub file %s" % path)
            ret = False
            continue
        else:
            print(stderr, file=sys.stderr)
        if p.returncode:
            raise RuntimeError("install_name_tool failed with exit status %d"
                % p.returncode)
    return ret

if __name__ == '__main__':
    import sys
    if sys.platform == 'darwin':
        for path in '/bin/ls', '/etc/locate.rc':
            print(path, is_macho(path))
<commit_msg>Add debug output to assert<commit_after>from __future__ import absolute_import, division, print_function

import subprocess
from os.path import islink, isfile
import sys

NO_EXT = (
    '.py', '.pyc', '.pyo', '.h', '.a', '.c', '.txt', '.html',
    '.xml', '.png', '.jpg', '.gif', '.class',
)

MAGIC = {
    b'\xca\xfe\xba\xbe': 'MachO-universal',
    b'\xce\xfa\xed\xfe': 'MachO-i386',
    b'\xcf\xfa\xed\xfe': 'MachO-x86_64',
    b'\xfe\xed\xfa\xce': 'MachO-ppc',
    b'\xfe\xed\xfa\xcf': 'MachO-ppc64',
}

FILETYPE = {
    b'\x01\x00\x00\x00': 'MH_OBJECT',
    b'\x02\x00\x00\x00': 'MH_EXECUTE',
    b'\x03\x00\x00\x00': 'MH_FVMLIB',
    b'\x04\x00\x00\x00': 'MH_CORE',
    b'\x05\x00\x00\x00': 'MH_PRELOAD',
    b'\x06\x00\x00\x00': 'MH_DYLIB',
    b'\x07\x00\x00\x00': 'MH_DYLINKER',
    b'\x08\x00\x00\x00': 'MH_BUNDLE',
    b'\x09\x00\x00\x00': 'MH_DYLIB_STUB',
    b'\x0a\x00\x00\x00': 'MH_DSYM',
    b'\x0b\x00\x00\x00': 'MH_KEXT_BUNDLE',
}

def is_macho(path):
    if path.endswith(NO_EXT) or islink(path) or not isfile(path):
        return False
    with open(path, 'rb') as fi:
        head = fi.read(4)
    return bool(head in MAGIC)

def is_dylib(path):
    with open(path, 'rb') as fi:
        # file type indicated by fourth 32-bit constant in the mach header
        header = fi.read(16)[-4:]
        return header in FILETYPE and FILETYPE[header] == 'MH_DYLIB'

def otool(path):
    "thin wrapper around otool -L"
    lines = subprocess.check_output(['otool', '-L', path]).decode('utf-8').splitlines()
    assert lines[0].startswith(path), path
    res = []
    for line in lines[1:]:
        assert line[0] == '\t', path
        res.append(line.split()[0])
    return res

def install_name_change(path, cb_func):
    """
    change dynamic shared library install names of Mach-O binary `path`.

    `cb_func` is a callback function which called for each shared library name.
    It is called with `path` and the current shared library install name,
    and return the new name (or None if the name should be unchanged).
    """
    changes = []
    for link in otool(path):
        new_link = cb_func(path, link)
        if new_link:
            changes.append((link, new_link))

    ret = True
    for old, new in changes:
        args = ['install_name_tool', '-change', old, new, path]
        print(' '.join(args))
        p = subprocess.Popen(args, stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        stderr = stderr.decode('utf-8')
        if "Mach-O dynamic shared library stub file" in stderr:
            print("Skipping Mach-O dynamic shared library stub file %s" % path)
            ret = False
            continue
        else:
            print(stderr, file=sys.stderr)
        if p.returncode:
            raise RuntimeError("install_name_tool failed with exit status %d"
                % p.returncode)
    return ret

if __name__ == '__main__':
    import sys
    if sys.platform == 'darwin':
        for path in '/bin/ls', '/etc/locate.rc':
            print(path, is_macho(path))
<|endoftext|>


<commit_before>import datetime
import re
import sys
from urllib.error import URLError
from urllib.parse import urlparse
from urllib.request import urlopen

from Bio import Entrez
import xmltodict

class Publication(object):
    """
    Use a PubMedLookup record to make a Publication object with info about
    a scientific publication.
    """
    def __init__(self, pubmed_record):
        """
        Upon init: set Publication attributes (record, pmid, pubmed_url,
        title, authors, first_author, last_author, journal, volume, issue,
        pages, url, abstract, year, month, and day)
        """
        self.record = pubmed_record.record
        self.pmid = self.record.get('Id')
        self.pubmed_url = 'http://www.ncbi.nlm.nih.gov/pubmed/{}' \
                          .format(self.pmid)
        self.title = self.record.get('Title')
        self.authors = ", ".join(self.record.get('AuthorList'))
        self.first_author = self.record.get('AuthorList')[0]
        self.last_author = self.record.get('AuthorList')[-1]
        self.journal = self.record.get('Source')
        self.volume = self.record.get('Volume')
        self.issue = self.record.get('Issue')
        self.pages = self.record.get('Pages')
        self.set_article_url()

        xml_dict = self.get_pubmed_xml()
        self.set_abstract(xml_dict)
        self.set_pub_year_month_day(xml_dict)

    def authors_et_al(self, max_authors=5):
        """
        Return string with a truncated author list followed by 'et al.'
        """
        author_list = self.record.get('AuthorList')
        if len(author_list) <= max_authors:
            authors_et_al = self.authors
        else:
            authors_et_al = ", ".join(
                self.record.get('AuthorList')[:max_authors]) + ", et al."
        return authors_et_al

    def cite(self, max_authors=5):
        """
        Return string with a citation for the record, formatted as:
        '{authors} ({year}). {title} {journal} {volume}({issue}): {pages}.'
        """
        citation_data = {
            'title': self.title,
            'authors': self.authors_et_al(max_authors),
            'year': self.year,
            'journal': self.journal,
            'volume': self.volume,
            'issue': self.issue,
            'pages': self.pages,
        }
        citation = "{authors} ({year}). {title} {journal}".format(**citation_data)
        if self.volume and self.issue and self.pages:
            citation += " {volume}({issue}): {pages}.".format(**citation_data)
        elif self.volume and self.issue:
            citation += " {volume}({issue}).".format(**citation_data)
        elif self.volume and self.pages:
            citation += " {volume}: {pages}.".format(**citation_data)
        elif self.volume:
            citation += " {volume}.".format(**citation_data)
        elif self.pages:
            citation += " {pages}.".format(**citation_data)
        else:
            citation += "."

        return citation

    def cite_mini(self):
        """
        Return string with a citation for the record, formatted as:
        '{first_author} - {year} - {journal}'
        """
        citation_data = [self.first_author]

        if len(self.record.get('AuthorList')) > 1:
            citation_data.append(self.last_author)

        citation_data.extend([self.year, self.journal])

        return " - ".join(citation_data)

    @staticmethod
    def parse_abstract(xml_dict):
        """
        Parse PubMed XML dictionary to retrieve abstract.
        """
        abstract_paragraphs = []
        abstract_xml = xml_dict['PubmedArticleSet']['PubmedArticle'] \
            ['MedlineCitation']['Article']['Abstract']['AbstractText']

        if isinstance(abstract_xml, str):
            abstract_paragraphs.append(abstract_xml)

        elif isinstance(abstract_xml, dict):
            abstract_text = abstract_xml.get('#text')
            try:
                abstract_label = abstract_xml['@Label']
            except KeyError:
                abstract_paragraphs.append(abstract_text)
            else:
                abstract_paragraphs.append("{}: {}"
                    .format(abstract_label, abstract_text))

        elif isinstance(abstract_xml, list):
            for abstract_section in abstract_xml:
                try:
                    abstract_text = abstract_section['#text']
                except KeyError:
                    abstract_text = abstract_section

                try:
                    abstract_label = abstract_section['@Label']
                except KeyError:
                    abstract_paragraphs.append(abstract_text)
                else:
                    abstract_paragraphs.append("{}: {}"
                        .format(abstract_label, abstract_text))

        else:
            raise RuntimeError("Error parsing abstract.")

        return "\n\n".join(abstract_paragraphs)

    def get_pubmed_xml(self):
        """
        Use a PubMed ID to retrieve PubMed metadata in XML form.
        """
        url = 'http://eutils.ncbi.nlm.nih.gov/entrez/eutils/' \
              'efetch.fcgi?db=pubmed&rettype=abstract&id={}' \
              .format(self.pmid)

        try:
            response = urlopen(url)
        except URLError:
            xml_dict = ''
        else:
            xml = response.read().decode()
            xml_dict = xmltodict.parse(xml)

        return xml_dict

    def set_abstract(self, xml_dict):
        """
        If record has an abstract, get it extract it from PubMed's XML data
        """
        if self.record.get('HasAbstract') == 1 and xml_dict:
                self.abstract = self.parse_abstract(xml_dict)
        else:
            self.abstract = ''

    def set_article_url(self):
        """
        If record has a DOI, set article URL based on where the DOI points.
        """
        if 'DOI' in self.record:
            doi_url = "/".join(['http://dx.doi.org', self.record['DOI']])

            try:
                response = urlopen(doi_url)
            except URLError:
                self.url = ''
            else:
                self.url = response.geturl()
        else:
            self.url = ''

    def set_pub_year_month_day(self, xml_dict):
        """
        Set publication year, month, day from PubMed's XML data
        """
        try:
            pubdate_xml = xml_dict['PubmedArticleSet']['PubmedArticle'] \
                ['MedlineCitation']['Article']['Journal']['JournalIssue'] \
                ['PubDate']
        except KeyError:
            year = ''
            month = ''
            day = ''
        else:
            year = pubdate_xml.get('Year')
            month_short = pubdate_xml.get('Month')
            day = pubdate_xml.get('Day')

            try:
                month = datetime.datetime.strptime(month_short, "%b").month
            except ValueError:
                month = ''

        self.year, self.month, self.day = year, month, day

class PubMedLookup(object):
    """
    Retrieve a PubMed record using its PubMed ID or PubMed URL.
    (e.g., '22331878' or 'http://www.ncbi.nlm.nih.gov/pubmed/22331878')
    """

    def __init__(self, query, user_email):
        """
        Upon init: set email as required by API, determine whether query
        is PubMed ID or PubMed URL and retrieve PubMed record accordingly.
        """
        Entrez.email = user_email

        pmid_pattern = r'^\d+$'
        pmurl_pattern = r'^https?://www\.ncbi\.nlm\.nih\.gov/pubmed/\d+$'
        if re.match(pmid_pattern, str(query)):
            pmid = query
        elif re.match(pmurl_pattern, query):
            pmid = self.parse_pubmed_url(query)
        else:
            raise RuntimeError(
                "Query ({}) doesn't appear to be a PubMed ID or PubMed URL"
                .format(query))

        self.record = self.get_pubmed_record(pmid)[0]

    @staticmethod
    def parse_pubmed_url(pubmed_url):
        """Get PubMed ID (pmid) from PubMed URL."""
        parse_result = urlparse(pubmed_url)
        pattern = re.compile('^/pubmed/(\d+)$')
        pmid = pattern.match(parse_result.path).group(1)
        return pmid

    @staticmethod
    def get_pubmed_record(pmid):
        """Get PubMed record from PubMed ID."""
        handle = Entrez.esummary(db="pubmed", id=pmid)
        record = Entrez.read(handle)
        return record
<commit_msg>Only get author list once<commit_after>import datetime
import re
import sys
from urllib.error import URLError
from urllib.parse import urlparse
from urllib.request import urlopen

from Bio import Entrez
import xmltodict

class Publication(object):
    """
    Use a PubMedLookup record to make a Publication object with info about
    a scientific publication.
    """
    def __init__(self, pubmed_record):
        """
        Upon init: set Publication attributes (record, pmid, pubmed_url,
        title, authors, first_author, last_author, journal, volume, issue,
        pages, url, abstract, year, month, and day)
        """
        self.record = pubmed_record.record
        self.pmid = self.record.get('Id')
        self.pubmed_url = 'http://www.ncbi.nlm.nih.gov/pubmed/{}' \
                          .format(self.pmid)
        self.title = self.record.get('Title')
        self._author_list = self.record.get('AuthorList')
        self.authors = ", ".join(self._author_list)
        self.first_author = self._author_list[0]
        self.last_author = self._author_list[-1]
        self.journal = self.record.get('Source')
        self.volume = self.record.get('Volume')
        self.issue = self.record.get('Issue')
        self.pages = self.record.get('Pages')
        self.set_article_url()

        xml_dict = self.get_pubmed_xml()
        self.set_abstract(xml_dict)
        self.set_pub_year_month_day(xml_dict)

    def authors_et_al(self, max_authors=5):
        """
        Return string with a truncated author list followed by 'et al.'
        """
        author_list = self._author_list
        if len(author_list) <= max_authors:
            authors_et_al = self.authors
        else:
            authors_et_al = ", ".join(
                self._author_list[:max_authors]) + ", et al."
        return authors_et_al

    def cite(self, max_authors=5):
        """
        Return string with a citation for the record, formatted as:
        '{authors} ({year}). {title} {journal} {volume}({issue}): {pages}.'
        """
        citation_data = {
            'title': self.title,
            'authors': self.authors_et_al(max_authors),
            'year': self.year,
            'journal': self.journal,
            'volume': self.volume,
            'issue': self.issue,
            'pages': self.pages,
        }
        citation = "{authors} ({year}). {title} {journal}".format(**citation_data)
        if self.volume and self.issue and self.pages:
            citation += " {volume}({issue}): {pages}.".format(**citation_data)
        elif self.volume and self.issue:
            citation += " {volume}({issue}).".format(**citation_data)
        elif self.volume and self.pages:
            citation += " {volume}: {pages}.".format(**citation_data)
        elif self.volume:
            citation += " {volume}.".format(**citation_data)
        elif self.pages:
            citation += " {pages}.".format(**citation_data)
        else:
            citation += "."

        return citation

    def cite_mini(self):
        """
        Return string with a citation for the record, formatted as:
        '{first_author} - {year} - {journal}'
        """
        citation_data = [self.first_author]

        if len(self._author_list) > 1:
            citation_data.append(self.last_author)

        citation_data.extend([self.year, self.journal])

        return " - ".join(citation_data)

    @staticmethod
    def parse_abstract(xml_dict):
        """
        Parse PubMed XML dictionary to retrieve abstract.
        """
        abstract_paragraphs = []
        abstract_xml = xml_dict['PubmedArticleSet']['PubmedArticle'] \
            ['MedlineCitation']['Article']['Abstract']['AbstractText']

        if isinstance(abstract_xml, str):
            abstract_paragraphs.append(abstract_xml)

        elif isinstance(abstract_xml, dict):
            abstract_text = abstract_xml.get('#text')
            try:
                abstract_label = abstract_xml['@Label']
            except KeyError:
                abstract_paragraphs.append(abstract_text)
            else:
                abstract_paragraphs.append("{}: {}"
                    .format(abstract_label, abstract_text))

        elif isinstance(abstract_xml, list):
            for abstract_section in abstract_xml:
                try:
                    abstract_text = abstract_section['#text']
                except KeyError:
                    abstract_text = abstract_section

                try:
                    abstract_label = abstract_section['@Label']
                except KeyError:
                    abstract_paragraphs.append(abstract_text)
                else:
                    abstract_paragraphs.append("{}: {}"
                        .format(abstract_label, abstract_text))

        else:
            raise RuntimeError("Error parsing abstract.")

        return "\n\n".join(abstract_paragraphs)

    def get_pubmed_xml(self):
        """
        Use a PubMed ID to retrieve PubMed metadata in XML form.
        """
        url = 'http://eutils.ncbi.nlm.nih.gov/entrez/eutils/' \
              'efetch.fcgi?db=pubmed&rettype=abstract&id={}' \
              .format(self.pmid)

        try:
            response = urlopen(url)
        except URLError:
            xml_dict = ''
        else:
            xml = response.read().decode()
            xml_dict = xmltodict.parse(xml)

        return xml_dict

    def set_abstract(self, xml_dict):
        """
        If record has an abstract, get it extract it from PubMed's XML data
        """
        if self.record.get('HasAbstract') == 1 and xml_dict:
                self.abstract = self.parse_abstract(xml_dict)
        else:
            self.abstract = ''

    def set_article_url(self):
        """
        If record has a DOI, set article URL based on where the DOI points.
        """
        if 'DOI' in self.record:
            doi_url = "/".join(['http://dx.doi.org', self.record['DOI']])

            try:
                response = urlopen(doi_url)
            except URLError:
                self.url = ''
            else:
                self.url = response.geturl()
        else:
            self.url = ''

    def set_pub_year_month_day(self, xml_dict):
        """
        Set publication year, month, day from PubMed's XML data
        """
        try:
            pubdate_xml = xml_dict['PubmedArticleSet']['PubmedArticle'] \
                ['MedlineCitation']['Article']['Journal']['JournalIssue'] \
                ['PubDate']
        except KeyError:
            year = ''
            month = ''
            day = ''
        else:
            year = pubdate_xml.get('Year')
            month_short = pubdate_xml.get('Month')
            day = pubdate_xml.get('Day')

            try:
                month = datetime.datetime.strptime(month_short, "%b").month
            except ValueError:
                month = ''

        self.year, self.month, self.day = year, month, day

class PubMedLookup(object):
    """
    Retrieve a PubMed record using its PubMed ID or PubMed URL.
    (e.g., '22331878' or 'http://www.ncbi.nlm.nih.gov/pubmed/22331878')
    """

    def __init__(self, query, user_email):
        """
        Upon init: set email as required by API, determine whether query
        is PubMed ID or PubMed URL and retrieve PubMed record accordingly.
        """
        Entrez.email = user_email

        pmid_pattern = r'^\d+$'
        pmurl_pattern = r'^https?://www\.ncbi\.nlm\.nih\.gov/pubmed/\d+$'
        if re.match(pmid_pattern, str(query)):
            pmid = query
        elif re.match(pmurl_pattern, query):
            pmid = self.parse_pubmed_url(query)
        else:
            raise RuntimeError(
                "Query ({}) doesn't appear to be a PubMed ID or PubMed URL"
                .format(query))

        self.record = self.get_pubmed_record(pmid)[0]

    @staticmethod
    def parse_pubmed_url(pubmed_url):
        """Get PubMed ID (pmid) from PubMed URL."""
        parse_result = urlparse(pubmed_url)
        pattern = re.compile('^/pubmed/(\d+)$')
        pmid = pattern.match(parse_result.path).group(1)
        return pmid

    @staticmethod
    def get_pubmed_record(pmid):
        """Get PubMed record from PubMed ID."""
        handle = Entrez.esummary(db="pubmed", id=pmid)
        record = Entrez.read(handle)
        return record
<|endoftext|>