<commit_before>"""
Provides functions for syncing entities and their relationships to the
Entity and EntityRelationship tables.
"""
import logging
import sys

import wrapt
from collections import defaultdict
from itertools import chain

from django.contrib.contenttypes.models import ContentType
import manager_utils
from django.db import transaction, connection

from entity.config import entity_registry
from entity.models import Entity, EntityRelationship, EntityKind

LOG = logging.getLogger(__name__)

@wrapt.decorator
def defer_entity_syncing(wrapped, instance, args, kwargs):
    """
    A decorator that can be used to defer the syncing of entities until after the method has been run
    This is being introduced to help avoid deadlocks in the meantime as we attempt to better understand
    why they are happening
    """

    # Defer entity syncing while we run our method
    sync_entities.defer = True

    # Run the method
    try:
        return wrapped(*args, **kwargs)

    # After we run the method disable the deferred syncing
    # and sync all the entities that have been buffered to be synced
    finally:
        # Enable entity syncing again
        sync_entities.defer = False

        # Get the models that need to be synced
        model_objs = list(sync_entities.buffer.values())

        # If none is in the model objects we need to sync all
        if None in model_objs:
            model_objs = list()

        # Clear the buffer
        sync_entities.buffer = {}

        # Sync the entities that were deferred
        sync_entities(*model_objs)

class SyncLock(object):
    """
    This will create a postgres advisory lock to be used during the syncing process
    This is being introduced to help avoid deadlocks in the meantime as we attempt to better understand
    why they are happening
    """
    def __init__(self, model, name):
        # Save the model
        self._model = model

        # Save the name
        self._name = name

        # Create an empty transaction
        self._transaction = None

        # Call the parent
        super(SyncLock, self).__init__()

    def __enter__(self):
        # Create the transaction
        self._transaction = transaction.atomic()

        lock_sql = (
            'SELECT '
            'pg_advisory_xact_lock(\'{table_name}\'::regclass::integer, hashtext(\'{lock_name}\'));'
        ).format(
            table_name=self._model._meta.db_table,
            lock_name=self._name
        )

        # Start the transaction
        self._transaction.__enter__()

        # Keep a reference to the exception
        exception = None

        # Create the connection
        try:
            with connection.cursor() as cursor:
                cursor.execute(lock_sql)
        except Exception as e:
            exception = e

        # If we have an exception, raise it
        if exception is not None:
            self.__exit__(*sys.exc_info())
            raise exception

        # Return the transaction
        return transaction

    def __exit__(self, *args, **kwargs):
        # Exit the transaction
        try:
            if self._transaction:
                self._transaction.__exit__(*args, **kwargs)
        except:
            pass

def _get_super_entities_by_ctype(model_objs_by_ctype, model_ids_to_sync, sync_all):
    """
    Given model objects organized by content type and a dictionary of all model IDs that need
    to be synced, organize all super entity relationships that need to be synced.

    Ensure that the model_ids_to_sync dict is updated with any new super entities
    that need to be part of the overall entity sync
    """
    super_entities_by_ctype = defaultdict(lambda: defaultdict(list))  # pragma: no cover
    for ctype, model_objs_for_ctype in model_objs_by_ctype.items():
        entity_config = entity_registry.entity_registry.get(ctype.model_class())
        super_entities = entity_config.get_super_entities(model_objs_for_ctype, sync_all)
        super_entities_by_ctype[ctype] = {
            ContentType.objects.get_for_model(model_class, for_concrete_model=False): relationships
            for model_class, relationships in super_entities.items()
        }

        # Continue adding to the set of entities that need to be synced
        for super_entity_ctype, relationships in super_entities_by_ctype[ctype].items():
            for sub_entity_id, super_entity_id in relationships:
                model_ids_to_sync[ctype].add(sub_entity_id)
                model_ids_to_sync[super_entity_ctype].add(super_entity_id)

    return super_entities_by_ctype

def _get_model_objs_to_sync(model_ids_to_sync, model_objs_map, sync_all):
    """
    Given the model IDs to sync, fetch all model objects to sync
    """
    model_objs_to_sync = {}
    for ctype, model_ids_to_sync_for_ctype in model_ids_to_sync.items():
        model_qset = entity_registry.entity_registry.get(ctype.model_class()).queryset

        if not sync_all:
            model_objs_to_sync[ctype] = model_qset.filter(id__in=model_ids_to_sync_for_ctype)
        else:
            model_objs_to_sync[ctype] = [
                model_objs_map[ctype, model_id] for model_id in model_ids_to_sync_for_ctype
            ]

    return model_objs_to_sync

def sync_entities(*model_objs):
    """
    Syncs entities

    Args:
        model_objs (List[Model]): The model objects to sync. If empty, all entities will be synced
    """

    # Check if we are deferring processing
    if sync_entities.defer:
        # If we dont have any model objects passed add a none to let us know that we need to sync all
        if not model_objs:
            sync_entities.buffer[None] = None
        else:
            # Add each model obj to the buffer
            for model_obj in model_objs:
                sync_entities.buffer[(model_obj.__class__, model_obj.pk)] = model_obj

        # Return false that we did not do anything
        return False

    # Log what we are syncing
    LOG.debug('sync_entities', model_objs)

    # Determine if we are syncing all
    sync_all = not model_objs
    model_objs_map = {
        (ContentType.objects.get_for_model(model_obj, for_concrete_model=False), model_obj.id): model_obj
        for model_obj in model_objs
    }
    if not model_objs_map:
        # Sync everything
        for model_class, entity_config in entity_registry.entity_registry.items():
            model_qset = entity_config.queryset
            model_objs_map.update({
                (ContentType.objects.get_for_model(model_class, for_concrete_model=False), model_obj.id): model_obj
                for model_obj in model_qset.all()
            })

    # Organize by content type
    model_objs_by_ctype = defaultdict(list)
    for (ctype, model_id), model_obj in model_objs_map.items():
        model_objs_by_ctype[ctype].append(model_obj)

    # Build a dict of all entities that need to be synced. These include the original models
    # and any super entities from super_entities_by_ctype. This dict is keyed on ctype with
    # a list of IDs of each model
    model_ids_to_sync = defaultdict(set)
    for (ctype, model_id), model_obj in model_objs_map.items():
        model_ids_to_sync[ctype].add(model_obj.id)

    # For each ctype, obtain super entities. This is a dict keyed on ctype. Each value
    # is a dict keyed on the ctype of the super entity with a list of tuples for
    # IDs of sub/super entity relationships
    super_entities_by_ctype = _get_super_entities_by_ctype(model_objs_by_ctype, model_ids_to_sync, sync_all)

    # Now that we have all models we need to sync, fetch them so that we can extract
    # metadata and entity kinds. If we are syncing all entities, we've already fetched
    # everything and can fill in this data struct without doing another DB hit
    model_objs_to_sync = _get_model_objs_to_sync(model_ids_to_sync, model_objs_map, sync_all)

    # Obtain all entity kind tuples associated with the models
    entity_kind_tuples_to_sync = set()
    for ctype, model_objs_to_sync_for_ctype in model_objs_to_sync.items():
        entity_config = entity_registry.entity_registry.get(ctype.model_class())
        for model_obj in model_objs_to_sync_for_ctype:
            entity_kind_tuples_to_sync.add(entity_config.get_entity_kind(model_obj))

    # Upsert all entity kinds and obtain the map of them
    entity_kinds_to_upsert = [
        EntityKind(name=name, display_name=display_name)
        for name, display_name in entity_kind_tuples_to_sync
    ]
    with SyncLock(EntityKind, 'sync_entities__entity_kind'):
        upserted_entity_kinds = manager_utils.bulk_upsert2(
            EntityKind.all_objects.all(),
            entity_kinds_to_upsert,
            ['name'],
            ['display_name'],
            returning=['id', 'name'],
            ignore_duplicate_updates=True,
            return_untouched=True
        )
    entity_kinds_map = {
        entity_kind.name: entity_kind
        for entity_kind in upserted_entity_kinds
    }

    # Now that we have all entity kinds, build all entities that need to be synced
    entities_to_upsert = []
    for ctype, model_objs_to_sync_for_ctype in model_objs_to_sync.items():
        entity_config = entity_registry.entity_registry.get(ctype.model_class())
        entities_to_upsert.extend([
            Entity(
                entity_id=model_obj.id,
                entity_type_id=ctype.id,
                entity_kind_id=entity_kinds_map[entity_config.get_entity_kind(model_obj)[0]].id,
                entity_meta=entity_config.get_entity_meta(model_obj),
                display_name=entity_config.get_display_name(model_obj),
                is_active=entity_config.get_is_active(model_obj)
            )
            for model_obj in model_objs_to_sync_for_ctype
        ])

    # Upsert entities (or do a sync if we are updating all entities)
    with SyncLock(Entity, 'sync_entities__entity'):
        if model_objs:
            upserted_entities = manager_utils.bulk_upsert2(
                Entity.all_objects.all(),
                entities_to_upsert,
                ['entity_type_id', 'entity_id'],
                ['entity_kind_id', 'entity_meta', 'display_name', 'is_active'],
                returning=['id', 'entity_type_id', 'entity_id'],
                ignore_duplicate_updates=True,
                return_untouched=True
            )
        else:
            upserted_entities = manager_utils.sync2(
                Entity.all_objects.all(),
                entities_to_upsert,
                ['entity_type_id', 'entity_id'],
                ['entity_kind_id', 'entity_meta', 'display_name', 'is_active'],
                returning=['id', 'entity_type_id', 'entity_id'],
                ignore_duplicate_updates=True
            )

    entities_map = {
        (entity.entity_type_id, entity.entity_id): entity
        for entity in chain(
            upserted_entities.updated, upserted_entities.created, upserted_entities.untouched
        )
    }

    # Now that all entities are upserted, sync entity relationships
    entity_relationships_to_sync = [
        EntityRelationship(
            sub_entity_id=entities_map[sub_ctype.id, sub_entity_id].id,
            super_entity_id=entities_map[super_ctype.id, super_entity_id].id,
        )
        for sub_ctype, super_entities_by_sub_ctype in super_entities_by_ctype.items()
        for super_ctype, relationships in super_entities_by_sub_ctype.items()
        for sub_entity_id, super_entity_id in relationships
    ]
    # Find the entities of the original model objects we were syncing. These
    # are needed to properly sync entity relationships
    original_entity_ids = [
        entities_map[ctype.id, model_obj.id].id
        for ctype, model_objs_for_ctype in model_objs_by_ctype.items()
        for model_obj in model_objs_for_ctype
    ]

    if not model_objs:
        # If we're syncing everything, just sync against the entire entity relationship
        # table instead of doing a complex __in query
        sync_against = EntityRelationship.objects.all()
    else:
        sync_against = EntityRelationship.objects.filter(sub_entity_id__in=original_entity_ids)

    # Sync the relations
    with SyncLock(EntityRelationship, 'sync_entities__entity_relationship'):
        manager_utils.sync2(
            sync_against,
            entity_relationships_to_sync,
            ['sub_entity_id', 'super_entity_id'],
            [],
            ignore_duplicate_updates=True
        )

# Add a defer and buffer method to the sync entities method
# This is used by the defer_entity_syncing decorator
sync_entities.defer = False
sync_entities.buffer = {}

def sync_entities_watching(instance):
    """
    Syncs entities watching changes of a model instance.
    """
    for entity_model, entity_model_getter in entity_registry.entity_watching[instance.__class__]:
        model_objs = list(entity_model_getter(instance))
        if model_objs:
            sync_entities(*model_objs)
<commit_msg>Fix logging<commit_after>"""
Provides functions for syncing entities and their relationships to the
Entity and EntityRelationship tables.
"""
import logging
import sys

import wrapt
from collections import defaultdict
from itertools import chain

from django.contrib.contenttypes.models import ContentType
import manager_utils
from django.db import transaction, connection

from entity.config import entity_registry
from entity.models import Entity, EntityRelationship, EntityKind

LOG = logging.getLogger(__name__)

@wrapt.decorator
def defer_entity_syncing(wrapped, instance, args, kwargs):
    """
    A decorator that can be used to defer the syncing of entities until after the method has been run
    This is being introduced to help avoid deadlocks in the meantime as we attempt to better understand
    why they are happening
    """

    # Defer entity syncing while we run our method
    sync_entities.defer = True

    # Run the method
    try:
        return wrapped(*args, **kwargs)

    # After we run the method disable the deferred syncing
    # and sync all the entities that have been buffered to be synced
    finally:
        # Enable entity syncing again
        sync_entities.defer = False

        # Get the models that need to be synced
        model_objs = list(sync_entities.buffer.values())

        # If none is in the model objects we need to sync all
        if None in model_objs:
            model_objs = list()

        # Clear the buffer
        sync_entities.buffer = {}

        # Sync the entities that were deferred
        sync_entities(*model_objs)

class SyncLock(object):
    """
    This will create a postgres advisory lock to be used during the syncing process
    This is being introduced to help avoid deadlocks in the meantime as we attempt to better understand
    why they are happening
    """
    def __init__(self, model, name):
        # Save the model
        self._model = model

        # Save the name
        self._name = name

        # Create an empty transaction
        self._transaction = None

        # Call the parent
        super(SyncLock, self).__init__()

    def __enter__(self):
        # Create the transaction
        self._transaction = transaction.atomic()

        lock_sql = (
            'SELECT '
            'pg_advisory_xact_lock(\'{table_name}\'::regclass::integer, hashtext(\'{lock_name}\'));'
        ).format(
            table_name=self._model._meta.db_table,
            lock_name=self._name
        )

        # Start the transaction
        self._transaction.__enter__()

        # Keep a reference to the exception
        exception = None

        # Create the connection
        try:
            with connection.cursor() as cursor:
                cursor.execute(lock_sql)
        except Exception as e:
            exception = e

        # If we have an exception, raise it
        if exception is not None:
            self.__exit__(*sys.exc_info())
            raise exception

        # Return the transaction
        return transaction

    def __exit__(self, *args, **kwargs):
        # Exit the transaction
        try:
            if self._transaction:
                self._transaction.__exit__(*args, **kwargs)
        except:
            pass

def _get_super_entities_by_ctype(model_objs_by_ctype, model_ids_to_sync, sync_all):
    """
    Given model objects organized by content type and a dictionary of all model IDs that need
    to be synced, organize all super entity relationships that need to be synced.

    Ensure that the model_ids_to_sync dict is updated with any new super entities
    that need to be part of the overall entity sync
    """
    super_entities_by_ctype = defaultdict(lambda: defaultdict(list))  # pragma: no cover
    for ctype, model_objs_for_ctype in model_objs_by_ctype.items():
        entity_config = entity_registry.entity_registry.get(ctype.model_class())
        super_entities = entity_config.get_super_entities(model_objs_for_ctype, sync_all)
        super_entities_by_ctype[ctype] = {
            ContentType.objects.get_for_model(model_class, for_concrete_model=False): relationships
            for model_class, relationships in super_entities.items()
        }

        # Continue adding to the set of entities that need to be synced
        for super_entity_ctype, relationships in super_entities_by_ctype[ctype].items():
            for sub_entity_id, super_entity_id in relationships:
                model_ids_to_sync[ctype].add(sub_entity_id)
                model_ids_to_sync[super_entity_ctype].add(super_entity_id)

    return super_entities_by_ctype

def _get_model_objs_to_sync(model_ids_to_sync, model_objs_map, sync_all):
    """
    Given the model IDs to sync, fetch all model objects to sync
    """
    model_objs_to_sync = {}
    for ctype, model_ids_to_sync_for_ctype in model_ids_to_sync.items():
        model_qset = entity_registry.entity_registry.get(ctype.model_class()).queryset

        if not sync_all:
            model_objs_to_sync[ctype] = model_qset.filter(id__in=model_ids_to_sync_for_ctype)
        else:
            model_objs_to_sync[ctype] = [
                model_objs_map[ctype, model_id] for model_id in model_ids_to_sync_for_ctype
            ]

    return model_objs_to_sync

def sync_entities(*model_objs):
    """
    Syncs entities

    Args:
        model_objs (List[Model]): The model objects to sync. If empty, all entities will be synced
    """

    # Check if we are deferring processing
    if sync_entities.defer:
        # If we dont have any model objects passed add a none to let us know that we need to sync all
        if not model_objs:
            sync_entities.buffer[None] = None
        else:
            # Add each model obj to the buffer
            for model_obj in model_objs:
                sync_entities.buffer[(model_obj.__class__, model_obj.pk)] = model_obj

        # Return false that we did not do anything
        return False

    # Log what we are syncing
    LOG.info('sync_entities')
    LOG.info(model_objs)

    # Determine if we are syncing all
    sync_all = not model_objs
    model_objs_map = {
        (ContentType.objects.get_for_model(model_obj, for_concrete_model=False), model_obj.id): model_obj
        for model_obj in model_objs
    }
    if not model_objs_map:
        # Sync everything
        for model_class, entity_config in entity_registry.entity_registry.items():
            model_qset = entity_config.queryset
            model_objs_map.update({
                (ContentType.objects.get_for_model(model_class, for_concrete_model=False), model_obj.id): model_obj
                for model_obj in model_qset.all()
            })

    # Organize by content type
    model_objs_by_ctype = defaultdict(list)
    for (ctype, model_id), model_obj in model_objs_map.items():
        model_objs_by_ctype[ctype].append(model_obj)

    # Build a dict of all entities that need to be synced. These include the original models
    # and any super entities from super_entities_by_ctype. This dict is keyed on ctype with
    # a list of IDs of each model
    model_ids_to_sync = defaultdict(set)
    for (ctype, model_id), model_obj in model_objs_map.items():
        model_ids_to_sync[ctype].add(model_obj.id)

    # For each ctype, obtain super entities. This is a dict keyed on ctype. Each value
    # is a dict keyed on the ctype of the super entity with a list of tuples for
    # IDs of sub/super entity relationships
    super_entities_by_ctype = _get_super_entities_by_ctype(model_objs_by_ctype, model_ids_to_sync, sync_all)

    # Now that we have all models we need to sync, fetch them so that we can extract
    # metadata and entity kinds. If we are syncing all entities, we've already fetched
    # everything and can fill in this data struct without doing another DB hit
    model_objs_to_sync = _get_model_objs_to_sync(model_ids_to_sync, model_objs_map, sync_all)

    # Obtain all entity kind tuples associated with the models
    entity_kind_tuples_to_sync = set()
    for ctype, model_objs_to_sync_for_ctype in model_objs_to_sync.items():
        entity_config = entity_registry.entity_registry.get(ctype.model_class())
        for model_obj in model_objs_to_sync_for_ctype:
            entity_kind_tuples_to_sync.add(entity_config.get_entity_kind(model_obj))

    # Upsert all entity kinds and obtain the map of them
    entity_kinds_to_upsert = [
        EntityKind(name=name, display_name=display_name)
        for name, display_name in entity_kind_tuples_to_sync
    ]
    with SyncLock(EntityKind, 'sync_entities__entity_kind'):
        upserted_entity_kinds = manager_utils.bulk_upsert2(
            EntityKind.all_objects.all(),
            entity_kinds_to_upsert,
            ['name'],
            ['display_name'],
            returning=['id', 'name'],
            ignore_duplicate_updates=True,
            return_untouched=True
        )
    entity_kinds_map = {
        entity_kind.name: entity_kind
        for entity_kind in upserted_entity_kinds
    }

    # Now that we have all entity kinds, build all entities that need to be synced
    entities_to_upsert = []
    for ctype, model_objs_to_sync_for_ctype in model_objs_to_sync.items():
        entity_config = entity_registry.entity_registry.get(ctype.model_class())
        entities_to_upsert.extend([
            Entity(
                entity_id=model_obj.id,
                entity_type_id=ctype.id,
                entity_kind_id=entity_kinds_map[entity_config.get_entity_kind(model_obj)[0]].id,
                entity_meta=entity_config.get_entity_meta(model_obj),
                display_name=entity_config.get_display_name(model_obj),
                is_active=entity_config.get_is_active(model_obj)
            )
            for model_obj in model_objs_to_sync_for_ctype
        ])

    # Upsert entities (or do a sync if we are updating all entities)
    with SyncLock(Entity, 'sync_entities__entity'):
        if model_objs:
            upserted_entities = manager_utils.bulk_upsert2(
                Entity.all_objects.all(),
                entities_to_upsert,
                ['entity_type_id', 'entity_id'],
                ['entity_kind_id', 'entity_meta', 'display_name', 'is_active'],
                returning=['id', 'entity_type_id', 'entity_id'],
                ignore_duplicate_updates=True,
                return_untouched=True
            )
        else:
            upserted_entities = manager_utils.sync2(
                Entity.all_objects.all(),
                entities_to_upsert,
                ['entity_type_id', 'entity_id'],
                ['entity_kind_id', 'entity_meta', 'display_name', 'is_active'],
                returning=['id', 'entity_type_id', 'entity_id'],
                ignore_duplicate_updates=True
            )

    entities_map = {
        (entity.entity_type_id, entity.entity_id): entity
        for entity in chain(
            upserted_entities.updated, upserted_entities.created, upserted_entities.untouched
        )
    }

    # Now that all entities are upserted, sync entity relationships
    entity_relationships_to_sync = [
        EntityRelationship(
            sub_entity_id=entities_map[sub_ctype.id, sub_entity_id].id,
            super_entity_id=entities_map[super_ctype.id, super_entity_id].id,
        )
        for sub_ctype, super_entities_by_sub_ctype in super_entities_by_ctype.items()
        for super_ctype, relationships in super_entities_by_sub_ctype.items()
        for sub_entity_id, super_entity_id in relationships
    ]
    # Find the entities of the original model objects we were syncing. These
    # are needed to properly sync entity relationships
    original_entity_ids = [
        entities_map[ctype.id, model_obj.id].id
        for ctype, model_objs_for_ctype in model_objs_by_ctype.items()
        for model_obj in model_objs_for_ctype
    ]

    if not model_objs:
        # If we're syncing everything, just sync against the entire entity relationship
        # table instead of doing a complex __in query
        sync_against = EntityRelationship.objects.all()
    else:
        sync_against = EntityRelationship.objects.filter(sub_entity_id__in=original_entity_ids)

    # Sync the relations
    with SyncLock(EntityRelationship, 'sync_entities__entity_relationship'):
        manager_utils.sync2(
            sync_against,
            entity_relationships_to_sync,
            ['sub_entity_id', 'super_entity_id'],
            [],
            ignore_duplicate_updates=True
        )

# Add a defer and buffer method to the sync entities method
# This is used by the defer_entity_syncing decorator
sync_entities.defer = False
sync_entities.buffer = {}

def sync_entities_watching(instance):
    """
    Syncs entities watching changes of a model instance.
    """
    for entity_model, entity_model_getter in entity_registry.entity_watching[instance.__class__]:
        model_objs = list(entity_model_getter(instance))
        if model_objs:
            sync_entities(*model_objs)
<|endoftext|>


<commit_before>from Configuration import Configuration
# import pUtil
import sys
import os

def set_environment():

    env = Configuration()

    # env['version'] = pUtil.getPilotVersion()
    env['pilot_version_tag'] = 'RC' #'PR' #'RC'
    env['force_devpilot'] = False

    # get the python executable path, to be used by child job
    env['pyexe'] = sys.executable

    # default job dispatcher web server
    env['pshttpurl'] = "pandaserver.cern.ch"   # Should be set with pilot option -w <url>

    # cache URL used by some experiment classes
    env['cache'] = ""

    # URL for the schedconfig
    env['schedconfigURL'] = "http://" + env['pshttpurl'] # Could be set with pilot option -I <url>

    # global variables with some default (test) values
    env['psport'] = 25443                      # PanDA server listening port
    env['uflag'] = None                        # Flag if this pilot is to only get a user analysis job or ordinary job
    env['abortDQSpace'] = 10                   # Unit is in GB
    env['warnDQSpace'] = 100                   # Unit is in GB
    env['localsizelimit_stdout'] = 2*1024**2   # Size limit of payload stdout size during running. unit is in kB
    env['localspacelimit'] = 2*1024**2         # Space limit of remaining local disk size during running. unit is in kB
    env['localspacelimit_user'] = 7*1024**2    # Maximum size of user work area. unit is in kB
    env['localspacelimit0'] = 5*1024**2        # Initial space limit before asking for a job. unit is in kB
    env['outputlimit'] = 500*1024**3           # Maximum allowed size of an output file. unit is B
    env['pilotId'] = 'xtestP001'               # Pilot id
    env['jobSchedulerId'] = 'xtestJS001'       # Scheduler id
    env['heartbeatPeriod'] = 30*60             # 30 minutes limit for checking job state file
    env['maxjobrecDefault'] = 20               # Default maximum number of job recoveries
    env['maxjobrec'] = env['maxjobrecDefault'] # Maximum number of job recoveries
    env['zombieJobList'] = []                  # Zombie list
    env['prodJobDone'] = False                 # Will be set to true when production job has finished
    env['rmwkdir'] = None                      # If set to true, workdir will always be deleted
    env['jobrec'] = True                       # Perform job recovery per default
    env['jobRequestFlag'] = True               # Ask server for initial job (read from file otherwise)
    env['debugLevel'] = 0                      # 0: debug info off, 1: display function name when called, 2: full debug info
    env['maxNumberOfRecoveryAttempts'] = 15    # As attempted by the job recovery
    env['stagein'] = False                     # Set to True during stagein phase
    env['stageout'] = False                    # Set to True during stageout phase
    env['queuename'] = ""                      # Name of queue used to download config info
    env['loopingLimitDefaultProd'] = 12*3600   # If job does not write anything in N hours, it is considered a looping job (production jobs)
    env['loopingLimitDefaultUser'] = 3*3600    # If job does not write anything in N hours, it is considered a looping job (user analysis jobs)
    env['loopingLimitMinDefault'] = 2*3600     # Minimum allow looping limit
    env['stageinretry'] = 2                    # Number of stage-in tries
    env['stageoutretry'] = 2                   # Number of stage-out tries
    env['logFileDir'] = ""                     # Log file directory
    env['testLevel'] = "0"                     # Test suite control variable (0: no test, 1: put error, 2: ...)
    env['loggingMode'] = None                  # True puts the pilot in logging mode (will not run a job), False is a normal pilot. None means run as before
    env['updateServerFlag'] = True             # Switch on/off Panda server updates (for testing only), default is True
    env['globalSite'] = None                   # Site object
    env['globalWorkNode'] = None               # Work node object
    env['globalJob'] = None                    # Job object
    env['memory'] = None                       # For memory restrictions when asking the dispatcher for a new job (MB)
    env['proxycheckFlag'] = True               # True (default): perform proxy validity checks, False: no check
    env['wrapperFlag'] = False                 # True for wrappers that expect an exit code via return, False (default) when exit() can be used
    env['cleanupLimit'] = 2                    # Cleanup time limit in hours, see Cleaner.py
    env['logTransferred'] = False              # Boolean to keep track of whether the log has been transferred or not
    env['errorLabel'] = "WARNING"              # Set to FAILED when job recovery is not used
    env['nSent'] = 0                           # Throttle variable
    env['proxyguard'] = None                   # Global proxyguard object needed for the signal handler to be able to restore the proxy if necessary
    env['jobRecoveryMode'] = None              # When true, the pilot only runs job recovery, no payload download
    env['pilotToken'] = None                   # Pilot authentication token
    env['countryGroup'] = ""                   # Country group selector for getJob request
    env['workingGroup'] = ""                   # Working group selector for getJob request
    env['allowOtherCountry'] = False           #
    env['allowSameUser'] = True                # Boolean, multi-jobs will only come from same taskID (and thus same user)
    env['taskID'] = ""                         # TaskID of the current job or the first downloaded job in case allowSameUser is set to True
    env['inputDir'] = ""                       # Location of input files (source for mv site mover)
    env['outputDir'] = ""                      # Location of output files (destination for mv site mover)
    env['jobIds'] = []                         # Global job id list
    env['stageoutStartTime'] = None            # Set when pilot receives "stageout" state message, used by looping job killer
    env['timefloor_default'] = None            # Time limit for multi-jobs in minutes (turned off by default)
    env['useCoPilot'] = False                  # CERNVM Co-Pilot framework (on: let Co-Pilot finish job, off: let pilot finish job (default))
    env['update_freq_server'] = 30*60          # Server update frequency, 30 minutes
    env['experiment'] = "ATLAS"                # Current experiment (can be set with pilot option -F <experiment>)
    env['getjobmaxtime'] = 3*60                # Maximum time the pilot will attempt to download a single job (seconds)
    env['pandaJobDataFileName'] = "pandaJobData.out" # Job definition file name
    env['verifySoftwareDirectory'] = True      # Normally the software directory should be verified, but potentially not on an HPC system
    env['workerNode'] = None
    env['hasQueueData'] = None
    env['stdout_tail'] = ""
    env['stdout_path'] = ""
    env['maxtime'] = 999999                    # The maximum time a job is allowed to run (always set in pilot using schedconfig.maxtime)

    # panda proxy variables
    env['panda_proxy_url'] = ""
    env['panda_proxy_port'] = ""
    env['resourceType'] = ""

    # to test site mover
    env['copytool'] = 'gfal-copy'
    env['copytoolin'] = 'gfal-copy'

    # Harvester
    env['harvester'] = False                    # Harvester mode, True if pilot is launched by Harvester

    # In case the PILOT_INITDIR file has not been created yet, which means that this module is being
    # used by pilot.py, it will be created here using the current directory as init dir
    print "Current directory:", os.getcwd()
    path = os.path.join(os.getcwd(), "PILOT_INITDIR")
    if not os.path.exists(path):
        writeToFile(path, os.getcwd())
    else:
        print "File exists:", path
    env['pilot_initdir'] = readStringFromFile("PILOT_INITDIR")

    # data structure for job information in main loop
    # jobDic = {'prod/analy':[pid,job,os.getpgrp]}
    # (analysis label currently not used)
    env['jobDic'] = {}

    # some default values
    env['sitename'] = "testsite"
    env['workdir'] = os.getcwd() # "/tmp"
    env['appdir'] = ""
    env['return_code'] = None
    env['return'] = 0

    return env

def readStringFromFile(filename):
    """ read exit code from file <workdir>/EXITCODE """

    s = ""
    if os.path.exists(filename):
        try:
            f = open(filename, "r")
        except Exception, e:
            print "Failed to open %s: %s" % (filename, str(e))
        else:
            s = f.read()
            print "Found string %s in file %s" % (s, filename)
            f.close()
    else:
        print "No string to report (file %s does not exist)" % (filename)

    return s

def storePilotInitdir(targetdir, pilot_initdir):
    """ Store the pilot launch directory in a file used by environment.py """

    # This function is used to store the location of the init directory in the init directory itself as well as in the
    # site work directory. The location file is used by environment.py to set the global env['pilot_initdir'] used
    # by the pilot and the Monitor

    # This function must be called before the global env variable is instantiated in the pilot

    print "Storing PILOT_INITDIR2"
    path = os.path.join(targetdir, "PILOT_INITDIR2")
    writeToFile(path, pilot_initdir)

def writeToFile(filename, s):
    """ write string s to file """

    try:
        f = open(filename, "w")
    except Exception, e:
        print "!!WARNING!!2990!! Could not open: %s, %s" % (filename, str(e))
    else:
        f.write("%s" % (s))
        f.close()
        print 'Wrote string "%s" to file: %s' % (s, filename)
<commit_msg>Updated version tag<commit_after>from Configuration import Configuration
# import pUtil
import sys
import os

def set_environment():

    env = Configuration()

    # env['version'] = pUtil.getPilotVersion()
    env['pilot_version_tag'] = 'PR' #'RC'
    env['force_devpilot'] = False

    # get the python executable path, to be used by child job
    env['pyexe'] = sys.executable

    # default job dispatcher web server
    env['pshttpurl'] = "pandaserver.cern.ch"   # Should be set with pilot option -w <url>

    # cache URL used by some experiment classes
    env['cache'] = ""

    # URL for the schedconfig
    env['schedconfigURL'] = "http://" + env['pshttpurl'] # Could be set with pilot option -I <url>

    # global variables with some default (test) values
    env['psport'] = 25443                      # PanDA server listening port
    env['uflag'] = None                        # Flag if this pilot is to only get a user analysis job or ordinary job
    env['abortDQSpace'] = 10                   # Unit is in GB
    env['warnDQSpace'] = 100                   # Unit is in GB
    env['localsizelimit_stdout'] = 2*1024**2   # Size limit of payload stdout size during running. unit is in kB
    env['localspacelimit'] = 2*1024**2         # Space limit of remaining local disk size during running. unit is in kB
    env['localspacelimit_user'] = 7*1024**2    # Maximum size of user work area. unit is in kB
    env['localspacelimit0'] = 5*1024**2        # Initial space limit before asking for a job. unit is in kB
    env['outputlimit'] = 500*1024**3           # Maximum allowed size of an output file. unit is B
    env['pilotId'] = 'xtestP001'               # Pilot id
    env['jobSchedulerId'] = 'xtestJS001'       # Scheduler id
    env['heartbeatPeriod'] = 30*60             # 30 minutes limit for checking job state file
    env['maxjobrecDefault'] = 20               # Default maximum number of job recoveries
    env['maxjobrec'] = env['maxjobrecDefault'] # Maximum number of job recoveries
    env['zombieJobList'] = []                  # Zombie list
    env['prodJobDone'] = False                 # Will be set to true when production job has finished
    env['rmwkdir'] = None                      # If set to true, workdir will always be deleted
    env['jobrec'] = True                       # Perform job recovery per default
    env['jobRequestFlag'] = True               # Ask server for initial job (read from file otherwise)
    env['debugLevel'] = 0                      # 0: debug info off, 1: display function name when called, 2: full debug info
    env['maxNumberOfRecoveryAttempts'] = 15    # As attempted by the job recovery
    env['stagein'] = False                     # Set to True during stagein phase
    env['stageout'] = False                    # Set to True during stageout phase
    env['queuename'] = ""                      # Name of queue used to download config info
    env['loopingLimitDefaultProd'] = 12*3600   # If job does not write anything in N hours, it is considered a looping job (production jobs)
    env['loopingLimitDefaultUser'] = 3*3600    # If job does not write anything in N hours, it is considered a looping job (user analysis jobs)
    env['loopingLimitMinDefault'] = 2*3600     # Minimum allow looping limit
    env['stageinretry'] = 2                    # Number of stage-in tries
    env['stageoutretry'] = 2                   # Number of stage-out tries
    env['logFileDir'] = ""                     # Log file directory
    env['testLevel'] = "0"                     # Test suite control variable (0: no test, 1: put error, 2: ...)
    env['loggingMode'] = None                  # True puts the pilot in logging mode (will not run a job), False is a normal pilot. None means run as before
    env['updateServerFlag'] = True             # Switch on/off Panda server updates (for testing only), default is True
    env['globalSite'] = None                   # Site object
    env['globalWorkNode'] = None               # Work node object
    env['globalJob'] = None                    # Job object
    env['memory'] = None                       # For memory restrictions when asking the dispatcher for a new job (MB)
    env['proxycheckFlag'] = True               # True (default): perform proxy validity checks, False: no check
    env['wrapperFlag'] = False                 # True for wrappers that expect an exit code via return, False (default) when exit() can be used
    env['cleanupLimit'] = 2                    # Cleanup time limit in hours, see Cleaner.py
    env['logTransferred'] = False              # Boolean to keep track of whether the log has been transferred or not
    env['errorLabel'] = "WARNING"              # Set to FAILED when job recovery is not used
    env['nSent'] = 0                           # Throttle variable
    env['proxyguard'] = None                   # Global proxyguard object needed for the signal handler to be able to restore the proxy if necessary
    env['jobRecoveryMode'] = None              # When true, the pilot only runs job recovery, no payload download
    env['pilotToken'] = None                   # Pilot authentication token
    env['countryGroup'] = ""                   # Country group selector for getJob request
    env['workingGroup'] = ""                   # Working group selector for getJob request
    env['allowOtherCountry'] = False           #
    env['allowSameUser'] = True                # Boolean, multi-jobs will only come from same taskID (and thus same user)
    env['taskID'] = ""                         # TaskID of the current job or the first downloaded job in case allowSameUser is set to True
    env['inputDir'] = ""                       # Location of input files (source for mv site mover)
    env['outputDir'] = ""                      # Location of output files (destination for mv site mover)
    env['jobIds'] = []                         # Global job id list
    env['stageoutStartTime'] = None            # Set when pilot receives "stageout" state message, used by looping job killer
    env['timefloor_default'] = None            # Time limit for multi-jobs in minutes (turned off by default)
    env['useCoPilot'] = False                  # CERNVM Co-Pilot framework (on: let Co-Pilot finish job, off: let pilot finish job (default))
    env['update_freq_server'] = 30*60          # Server update frequency, 30 minutes
    env['experiment'] = "ATLAS"                # Current experiment (can be set with pilot option -F <experiment>)
    env['getjobmaxtime'] = 3*60                # Maximum time the pilot will attempt to download a single job (seconds)
    env['pandaJobDataFileName'] = "pandaJobData.out" # Job definition file name
    env['verifySoftwareDirectory'] = True      # Normally the software directory should be verified, but potentially not on an HPC system
    env['workerNode'] = None
    env['hasQueueData'] = None
    env['stdout_tail'] = ""
    env['stdout_path'] = ""
    env['maxtime'] = 999999                    # The maximum time a job is allowed to run (always set in pilot using schedconfig.maxtime)

    # panda proxy variables
    env['panda_proxy_url'] = ""
    env['panda_proxy_port'] = ""
    env['resourceType'] = ""

    # to test site mover
    env['copytool'] = 'gfal-copy'
    env['copytoolin'] = 'gfal-copy'

    # Harvester
    env['harvester'] = False                    # Harvester mode, True if pilot is launched by Harvester

    # In case the PILOT_INITDIR file has not been created yet, which means that this module is being
    # used by pilot.py, it will be created here using the current directory as init dir
    print "Current directory:", os.getcwd()
    path = os.path.join(os.getcwd(), "PILOT_INITDIR")
    if not os.path.exists(path):
        writeToFile(path, os.getcwd())
    else:
        print "File exists:", path
    env['pilot_initdir'] = readStringFromFile("PILOT_INITDIR")

    # data structure for job information in main loop
    # jobDic = {'prod/analy':[pid,job,os.getpgrp]}
    # (analysis label currently not used)
    env['jobDic'] = {}

    # some default values
    env['sitename'] = "testsite"
    env['workdir'] = os.getcwd() # "/tmp"
    env['appdir'] = ""
    env['return_code'] = None
    env['return'] = 0

    return env

def readStringFromFile(filename):
    """ read exit code from file <workdir>/EXITCODE """

    s = ""
    if os.path.exists(filename):
        try:
            f = open(filename, "r")
        except Exception, e:
            print "Failed to open %s: %s" % (filename, str(e))
        else:
            s = f.read()
            print "Found string %s in file %s" % (s, filename)
            f.close()
    else:
        print "No string to report (file %s does not exist)" % (filename)

    return s

def storePilotInitdir(targetdir, pilot_initdir):
    """ Store the pilot launch directory in a file used by environment.py """

    # This function is used to store the location of the init directory in the init directory itself as well as in the
    # site work directory. The location file is used by environment.py to set the global env['pilot_initdir'] used
    # by the pilot and the Monitor

    # This function must be called before the global env variable is instantiated in the pilot

    print "Storing PILOT_INITDIR2"
    path = os.path.join(targetdir, "PILOT_INITDIR2")
    writeToFile(path, pilot_initdir)

def writeToFile(filename, s):
    """ write string s to file """

    try:
        f = open(filename, "w")
    except Exception, e:
        print "!!WARNING!!2990!! Could not open: %s, %s" % (filename, str(e))
    else:
        f.write("%s" % (s))
        f.close()
        print 'Wrote string "%s" to file: %s' % (s, filename)
<|endoftext|>


<commit_before>'''
Analysis class:
Performs analysis on historical btc prices to create
a buy/sell sentiment. Might need to do linear regression
and hypothesis testing for sentiment
'''
import numpy
from account import Account
from data import Data

class Analysis:
    def __init__(self):
        # define mem var

<commit_msg>created linear regression models for 3 yr and 1 yr prices<commit_after>'''
Analysis class:
Performs analysis on historical btc prices to create
a buy/sell sentiment. Uses simple linear regression to
determine trend lines for extrapolated prediction
'''
from sklearn import linear_model
import numpy as np
import datetime as dt
import matplotlib.pyplot as plt
#from account import Account
#from data import Data

class Analysis:
    def __init__(self):
        self.total_days = [x for x in range(1094)]
        self.prices = []

        with open('prices.txt', 'r') as infile:
            infile.readline()
            data = infile.readlines()
            for x in data:
                line = x.split('\t')
                self.prices.append(float(line[1].strip('\n')))

        self.prices.reverse()

    def regr_three_yr(self):
        days = self.total_days
        prices = self.prices

        days = np.reshape(days, (len(days), 1))
        prices = np.reshape(prices, (len(prices), 1))

        regr = linear_model.LinearRegression()

        regr.fit(days, prices)

        data_points = []
        data_points.append(regr.predict(0)[0][0])
        data_points.append(regr.predict(1095)[0][0])

        return self.total_days, self.prices, [0, self.total_days[1093]], data_points

    def regr_one_yr(self):
        days = self.total_days[:365]
        prices = self.prices[729:]
        print(len(prices))

        days = np.reshape(days, (len(days), 1))
        prices = np.reshape(prices, (len(prices), 1))

        regr = linear_model.LinearRegression()

        regr.fit(days, prices)

        data_points = []
        data_points.append(regr.predict(0)[0][0])
        data_points.append(regr.predict(365)[0][0])

        return self.total_days[:364], self.prices[730:], [0, 364], data_points

if __name__ == '__main__':
    analysis_1 = Analysis()
    scatter_days, scatter_prices, regr_days, regr_prices = analysis_1.regr_three_yr()
    #scatter_days, scatter_prices, regr_days, regr_prices = analysis_1.regr_one_yr()

    plt.scatter(scatter_days, scatter_prices, s=1)
    plt.plot(regr_days, regr_prices, ls='--', color='#9a0000')
    plt.show()

'''
def main():
    infile = open('prices.txt', 'r')
    last_update = infile.readline()   # skip over the first line
    data = infile.readlines()   # read the data

    # create a list of training data
    days = []
    prices = []
    day = 1093
    for x in data:
        line = x.split('\t')
        days.append(day)
        prices.append(float(line[1].strip('\n')))
        day -= 1

    days.reverse()
    prices.reverse()

    plt.scatter(days, prices, s=1)

    regr = linear_model.LinearRegression()
    days = np.reshape(days, (len(days), 1))
    prices = np.reshape(prices, (len(prices), 1))

    regr.fit(days, prices)
    predicted_price = regr.predict(1458)
    print(regr.coef_)
    print(regr.intercept_)
    print(predicted_price)

    data_points = []
    data_points.append(regr.intercept_)
    data_points.append((regr.intercept_ + regr.coef_*len(days)))

    days_plot = [0, len(days)]

    plt.plot(days_plot, data_points, ls='--', color='#0de521')
    plt.show()

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


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

from __future__ import print_function

import curses
import argparse
import logging
import sys
import string

from data import *
from config import *
from view import *

def get_view(window, datum, config, file_num, total_files, position, show_help):
    cursor = position[:]
    link = position[:] if config.annotation_type == AnnType.link else [-1, -1]
    return View(window, cursor, link, datum, config, file_num, total_files, show_help)

def annotate(window, config, filenames):
    # Set color combinations
    for num, fore, back in COLORS:
        curses.init_pair(num, fore, back)

    # No blinking cursor
    curses.curs_set(0)

    cfilename = 0
    filename, start_pos, output_file, annotation_files = filenames[cfilename]
    datum = Datum(filename, config, output_file, annotation_files)
    view = get_view(window, datum, config, cfilename, len(filenames),
            start_pos, True)

    at_end = None
    search_term = ''
    typing_command = False
    nsteps = 0
    while True:
        if at_end is not None:
            # Draw screen
            view.render_edgecase(at_end)
            # Get input
            user_input = window.getch()
            nsteps += 1
            if at_end == 'start' and user_input in [ord('.'), ord("/")]:
                at_end = None
            elif at_end == 'end' and user_input in [ord(','), ord("\\")]:
                at_end = None
            elif user_input == ord("q"):
                break
        elif typing_command:
            # Draw screen
            view.render("/"+ search_term)
            view.must_show_linking_pos = False
            # Get input
            user_input = window.getch()
            nsteps += 1
            if nsteps % 100 == 0 and config.mode == Mode.annotate:
                datum.write_out()

###            logging.debug("Got:" + str(user_input))

            # TODO: Hacky, these numbers were worked out by hand.
            if user_input in [ord('?'), 10]:
                typing_command = False
            elif user_input in [ord('!'), 263, 127]:
                search_term = search_term[:-1]
            elif user_input in [ord(v) for v in string.printable]:
                search_term += chr(user_input)
        else:
            # Draw screen
            view.render(search_term)
            view.must_show_linking_pos = False
            # Get input
            user_input = window.getch()
            nsteps += 1
            if nsteps % 100 == 0 and config.mode == Mode.annotate:
                datum.write_out()

            # Note - First two are SHIFT + DOWN and SHIFT + UP, determined by
            # hand on two laptops.
            if user_input == ord('/'):
                typing_command = True
                search_term = ''
            elif user_input in [ord('c'), 337]:
                if config.annotation_type == AnnType.link:
                    view.move_up(True)
                    view.must_show_linking_pos = True
                else:
                    view.move_to_top()
            elif user_input in [ord('v'), 336]:
                if config.annotation_type == AnnType.link:
                    view.move_down(True)
                    view.must_show_linking_pos = True
                else:
                    view.move_to_bottom()
            elif user_input == curses.KEY_SLEFT:
                if config.annotation_type == AnnType.link:
                    view.move_left(True)
                else:
                    view.move_to_start()
            elif user_input == curses.KEY_SRIGHT:
                if config.annotation_type == AnnType.link:
                    view.move_right(True)
                else:
                    view.move_to_end()
            elif user_input == curses.KEY_UP:
                view.move_up()
            elif user_input == curses.KEY_DOWN:
                view.move_down()
            elif user_input == curses.KEY_LEFT:
                view.move_left()
            elif user_input == curses.KEY_RIGHT:
                view.move_right()
            elif user_input == ord("h"):
                view.toggle_help()
            elif user_input == ord("n"):
                if len(search_term) > 0:
                    view.next_match(search_term)
                else:
                    view.next_disagreement()
            elif user_input == ord("p"):
                if len(search_term) > 0:
                    view.previous_match(search_term)
                else:
                    view.previous_disagreement()
            elif user_input == ord("d") and config.mode == Mode.annotate:
                datum.modify_annotation(view.cursor, view.linking_pos)
                if config.annotation_type == AnnType.link:
                    if config.annotation == AnnScope.line:
                        view.move_down(True)
                        view.cursor[0] = view.linking_pos[0]
                        view.cursor[1] = view.linking_pos[1]
                        view.move_up()
                    else:
                        view.move_right(True)
                        view.cursor[0] = view.linking_pos[0]
                        view.cursor[1] = view.linking_pos[1]
                        view.move_left()
                    view.must_show_linking_pos = True
            elif user_input == ord("D") and config.mode == Mode.annotate:
                datum.modify_annotation(view.cursor, view.linking_pos)
            elif user_input == ord("u") and config.mode == Mode.annotate:
                datum.remove_annotation(view.cursor, view.linking_pos)
            elif user_input in [ord('s'), ord('b'), ord('r')]:
                if config.mode == Mode.annotate:
                    if config.annotation_type != AnnType.link:
                        datum.modify_annotation(view.cursor, view.linking_pos,
                                chr(user_input))
            elif user_input in [ord(c) for c in ",.q"]:
                # If we can get another file, do
                if config.mode == Mode.annotate:
                    datum.write_out()
                # TODO: Set to linking line rather than cursor where
                # appropriate
                filenames[cfilename][1][0] = view.cursor[0]
                filenames[cfilename][1][1] = view.cursor[1]
                if user_input == ord('q'):
                    break

                direction = 1 if user_input in [ord('.'), ord("/")] else -1
                if 0 <= cfilename + direction < len(filenames):
                    cfilename += direction
                    filename, start_pos, output_file, annotation_files = \
                            filenames[cfilename]
                    datum = Datum(filename, config, output_file,
                            annotation_files)
                    view = get_view(window, datum, config, cfilename,
                            len(filenames), start_pos, view.show_help)
                elif direction > 0:
                    at_end = 'end'
                else:
                    at_end = 'start'

        window.clear()

    out_filename = args.log_prefix + '.todo'
    out = open(out_filename, "w")
    for filename, start_pos, output_file, annotation_files in filenames:
        parts = [
            filename,
            output_file,
            str(start_pos[0]), str(start_pos[1]),
            ' '.join(annotation_files)
        ]
        print(" ".join(parts), file=out)
    out.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='A tool for annotating text data with extra information.')
    parser.add_argument('data', help='File containing a list of files to annotate')
    parser.add_argument('--log_prefix', help='Prefix for logging files (otherwise none)', default="annotation_log")
    parser.add_argument('--readonly', help='Do not allow changes or save annotations.', default=False)
    args = parser.parse_args()

    logging.basicConfig(filename=args.log_prefix + '.log', level=logging.DEBUG)

    ### Process configuration
    mode = Mode.read if args.readonly else Mode.annotate
    config = get_default_config(args, mode)
    filenames = read_filenames(args.data, config)
    if len(filenames) == 0:
        print("File '{}' contained no filenames".format(args.data))
        sys.exit(0)

    ### Start interface
    curses.wrapper(annotate, config, filenames)
<commit_msg>Debug -- Avoid edge case issue in saving<commit_after>#!/usr/bin/env python3

from __future__ import print_function

import curses
import argparse
import logging
import sys
import string

from data import *
from config import *
from view import *

def get_view(window, datum, config, file_num, total_files, position, show_help):
    cursor = position[:]
    link = position[:] if config.annotation_type == AnnType.link else [-1, -1]
    return View(window, cursor, link, datum, config, file_num, total_files, show_help)

def annotate(window, config, filenames):
    # Set color combinations
    for num, fore, back in COLORS:
        curses.init_pair(num, fore, back)

    # No blinking cursor
    curses.curs_set(0)

    cfilename = 0
    filename, start_pos, output_file, annotation_files = filenames[cfilename]
    datum = Datum(filename, config, output_file, annotation_files)
    view = get_view(window, datum, config, cfilename, len(filenames),
            start_pos, True)

    at_end = None
    search_term = ''
    typing_command = False
    nsteps = 0
    while True:
        if at_end is not None:
            # Draw screen
            view.render_edgecase(at_end)
            # Get input
            user_input = window.getch()
            if at_end == 'start' and user_input in [ord('.'), ord("/")]:
                at_end = None
            elif at_end == 'end' and user_input in [ord(','), ord("\\")]:
                at_end = None
            elif user_input == ord("q"):
                break
        elif typing_command:
            # Draw screen
            view.render("/"+ search_term)
            view.must_show_linking_pos = False
            # Get input
            user_input = window.getch()
            nsteps += 1
            if nsteps % 100 == 0 and config.mode == Mode.annotate:
                datum.write_out()

###            logging.debug("Got:" + str(user_input))

            # TODO: Hacky, these numbers were worked out by hand.
            if user_input in [ord('?'), 10]:
                typing_command = False
            elif user_input in [ord('!'), 263, 127]:
                search_term = search_term[:-1]
            elif user_input in [ord(v) for v in string.printable]:
                search_term += chr(user_input)
        else:
            # Draw screen
            view.render(search_term)
            view.must_show_linking_pos = False
            # Get input
            user_input = window.getch()
            nsteps += 1
            if nsteps % 100 == 0 and config.mode == Mode.annotate:
                datum.write_out()

            # Note - First two are SHIFT + DOWN and SHIFT + UP, determined by
            # hand on two laptops.
            if user_input == ord('/'):
                typing_command = True
                search_term = ''
            elif user_input in [ord('c'), 337]:
                if config.annotation_type == AnnType.link:
                    view.move_up(True)
                    view.must_show_linking_pos = True
                else:
                    view.move_to_top()
            elif user_input in [ord('v'), 336]:
                if config.annotation_type == AnnType.link:
                    view.move_down(True)
                    view.must_show_linking_pos = True
                else:
                    view.move_to_bottom()
            elif user_input == curses.KEY_SLEFT:
                if config.annotation_type == AnnType.link:
                    view.move_left(True)
                else:
                    view.move_to_start()
            elif user_input == curses.KEY_SRIGHT:
                if config.annotation_type == AnnType.link:
                    view.move_right(True)
                else:
                    view.move_to_end()
            elif user_input == curses.KEY_UP:
                view.move_up()
            elif user_input == curses.KEY_DOWN:
                view.move_down()
            elif user_input == curses.KEY_LEFT:
                view.move_left()
            elif user_input == curses.KEY_RIGHT:
                view.move_right()
            elif user_input == ord("h"):
                view.toggle_help()
            elif user_input == ord("n"):
                if len(search_term) > 0:
                    view.next_match(search_term)
                else:
                    view.next_disagreement()
            elif user_input == ord("p"):
                if len(search_term) > 0:
                    view.previous_match(search_term)
                else:
                    view.previous_disagreement()
            elif user_input == ord("d") and config.mode == Mode.annotate:
                datum.modify_annotation(view.cursor, view.linking_pos)
                if config.annotation_type == AnnType.link:
                    if config.annotation == AnnScope.line:
                        view.move_down(True)
                        view.cursor[0] = view.linking_pos[0]
                        view.cursor[1] = view.linking_pos[1]
                        view.move_up()
                    else:
                        view.move_right(True)
                        view.cursor[0] = view.linking_pos[0]
                        view.cursor[1] = view.linking_pos[1]
                        view.move_left()
                    view.must_show_linking_pos = True
            elif user_input == ord("D") and config.mode == Mode.annotate:
                datum.modify_annotation(view.cursor, view.linking_pos)
            elif user_input == ord("u") and config.mode == Mode.annotate:
                datum.remove_annotation(view.cursor, view.linking_pos)
            elif user_input in [ord('s'), ord('b'), ord('r')]:
                if config.mode == Mode.annotate:
                    if config.annotation_type != AnnType.link:
                        datum.modify_annotation(view.cursor, view.linking_pos,
                                chr(user_input))
            elif user_input in [ord(c) for c in ",.q"]:
                # If we can get another file, do
                if config.mode == Mode.annotate:
                    datum.write_out()
                # TODO: Set to linking line rather than cursor where
                # appropriate
                filenames[cfilename][1][0] = view.cursor[0]
                filenames[cfilename][1][1] = view.cursor[1]
                if user_input == ord('q'):
                    break

                direction = 1 if user_input in [ord('.'), ord("/")] else -1
                if 0 <= cfilename + direction < len(filenames):
                    cfilename += direction
                    filename, start_pos, output_file, annotation_files = \
                            filenames[cfilename]
                    datum = Datum(filename, config, output_file,
                            annotation_files)
                    view = get_view(window, datum, config, cfilename,
                            len(filenames), start_pos, view.show_help)
                elif direction > 0:
                    at_end = 'end'
                else:
                    at_end = 'start'

        window.clear()

    out_filename = args.log_prefix + '.todo'
    out = open(out_filename, "w")
    for filename, start_pos, output_file, annotation_files in filenames:
        parts = [
            filename,
            output_file,
            str(start_pos[0]), str(start_pos[1]),
            ' '.join(annotation_files)
        ]
        print(" ".join(parts), file=out)
    out.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='A tool for annotating text data with extra information.')
    parser.add_argument('data', help='File containing a list of files to annotate')
    parser.add_argument('--log_prefix', help='Prefix for logging files (otherwise none)', default="annotation_log")
    parser.add_argument('--readonly', help='Do not allow changes or save annotations.', default=False)
    args = parser.parse_args()

    logging.basicConfig(filename=args.log_prefix + '.log', level=logging.DEBUG)

    ### Process configuration
    mode = Mode.read if args.readonly else Mode.annotate
    config = get_default_config(args, mode)
    filenames = read_filenames(args.data, config)
    if len(filenames) == 0:
        print("File '{}' contained no filenames".format(args.data))
        sys.exit(0)

    ### Start interface
    curses.wrapper(annotate, config, filenames)
<|endoftext|>


<commit_before>#
# Author    : Manuel Bernal Llinares
# Project   : trackhub-creator
# Timestamp : 28-09-2017 14:20
# ---
# © 2017 Manuel Bernal Llinares <mbdebian@gmail.com>
# All rights reserved.
#

"""
This module offers converters between different formats
"""

import abc
# App imports
import config_manager
from parallel.models import ParallelRunner, CommandLineRunnerFactory

# Factories
class DataFormatConverterFactory:
    pass

# Possible base class for data format converters
class DataFormatConverter(ParallelRunner):
    def __init__(self):
        super().__init__()

class FileDataFormatConverter(DataFormatConverter):
    def __init__(self):
        super().__init__()
        self.file_path_source = ''

class BedToBigBedConverter(FileDataFormatConverter):
    def __init__(self):
        super().__init__()

    @abc.abstractmethod
    def _get_command_line_runner(self):
        ...

    def _run(self):
        pass
        # TODO - Conversion algorithm goes here

class BedToBigBedMultithreadedConverter(BedToBigBedConverter):
    def __init__(self):
        self._logger = config_manager.get_app_config_manager()\
            .get_logger_for("{}.{}".format(__name__, type(self).__name__))

    def _get_command_line_runner(self):
        return CommandLineRunnerFactory.get_multithread_command_line_runner()
<commit_msg>prevent the script from running separatedly from the application<commit_after>#
# Author    : Manuel Bernal Llinares
# Project   : trackhub-creator
# Timestamp : 28-09-2017 14:20
# ---
# © 2017 Manuel Bernal Llinares <mbdebian@gmail.com>
# All rights reserved.
#

"""
This module offers converters between different formats
"""

import abc
# App imports
import config_manager
from parallel.models import ParallelRunner, CommandLineRunnerFactory

# Factories
class DataFormatConverterFactory:
    pass

# Possible base class for data format converters
class DataFormatConverter(ParallelRunner):
    def __init__(self):
        super().__init__()

class FileDataFormatConverter(DataFormatConverter):
    def __init__(self):
        super().__init__()
        self.file_path_source = ''

class BedToBigBedConverter(FileDataFormatConverter):
    def __init__(self):
        super().__init__()

    @abc.abstractmethod
    def _get_command_line_runner(self):
        ...

    def _run(self):
        pass
        # TODO - Conversion algorithm goes here

class BedToBigBedMultithreadedConverter(BedToBigBedConverter):
    def __init__(self):
        self._logger = config_manager.get_app_config_manager() \
            .get_logger_for("{}.{}".format(__name__, type(self).__name__))

    def _get_command_line_runner(self):
        return CommandLineRunnerFactory.get_multithread_command_line_runner()

if __name__ == '__main__':
    print("ERROR: This script is part of a pipeline collection and it is not meant to be run in stand alone mode")
<|endoftext|>


<commit_before>from collections import OrderedDict
from sqlalchemy import or_, and_, func
from sqlalchemy.orm import outerjoin
from sqlalchemy.sql.expression import case

from dataactcore.models.domainModels import SF133, TASLookup, CGAC, FREC

gtas_model = SF133
tas_model = TASLookup

mapping = OrderedDict([
    ('allocation_transfer_agency', ['AllocationTransferAgencyIdentifier']),
    ('agency_identifier', ['AgencyIdentifier']),
    ('beginning_period_of_availa', ['BeginningPeriodOfAvailability']),
    ('ending_period_of_availabil', ['EndingPeriodOfAvailability']),
    ('availability_type_code', ['AvailabilityTypeCode']),
    ('main_account_code', ['MainAccountCode']),
    ('sub_account_code', ['SubAccountCode']),
    ('total_budgetary_resources_cpe', ['TotalBudgetaryResources_CPE']),
    ('budget_authority_appropria_cpe', ['BudgetAuthorityAppropriatedAmount_CPE']),
    ('budget_authority_unobligat_fyb', ['BudgetAuthorityUnobligatedBalanceBroughtForward_FYB']),
    ('adjustments_to_unobligated_cpe', ['AdjustmentsToUnobligatedBalanceBroughtForward_CPE']),
    ('other_budgetary_resources_cpe', ['OtherBudgetaryResourcesAmount_CPE']),
    ('contract_authority_amount_cpe', ['ContractAuthorityAmountTotal_CPE']),
    ('borrowing_authority_amount_cpe', ['BorrowingAuthorityAmountTotal_CPE']),
    ('spending_authority_from_of_cpe', ['SpendingAuthorityfromOffsettingCollectionsAmountTotal_CPE']),
    ('status_of_budgetary_resour_cpe', ['StatusOfBudgetaryResourcesTotal_CPE']),
    ('obligations_incurred_total_cpe', ['ObligationsIncurredTotalByTAS_CPE']),
    ('gross_outlay_amount_by_tas_cpe', ['GrossOutlayAmountByTAS_CPE']),
    ('unobligated_balance_cpe', ['UnobligatedBalance_CPE']),
    ('deobligations_recoveries_r_cpe', ['DeobligationsRecoveriesRefundsByTAS_CPE'])
])
db_columns = [key for key in mapping]

def query_data(session, agency_code, period, year):
    """ Request A file data

        Args:
            session: DB session
            agency_code: FREC or CGAC code for generation
            period: The period for which to get GTAS data
            year: The year for which to get GTAS data

        Returns:
            The rows using the provided dates for the given agency.
    """
    # set a boolean to determine if the original agency code is frec or cgac
    frec_provided = len(agency_code) == 4
    tas_gtas = tas_gtas_combo(session, period, year)
    # Make a list of FRECs to compare to for 011 AID entries
    frec_list = []
    if not frec_provided:
        frec_list = session.query(FREC.frec_code).select_from(outerjoin(CGAC, FREC, CGAC.cgac_id == FREC.cgac_id)).\
            filter(CGAC.cgac_code == agency_code).all()
        # Put the frec list in a format that can be read by a filter
        frec_list = [frec.frec_code for frec in frec_list]
    # Group agencies together that need to be grouped
    # NOTE: If these change, update A33.1 to match
    agency_array = []
    if agency_code == '097':
        agency_array = ['017', '021', '057', '097']
    elif agency_code == '020':
        agency_array = ['020', '580', '373']
    elif agency_code == '077':
        agency_array = ['077', '071']
    elif agency_code == '089':
        agency_array = ['089', '486']
    elif agency_code == '1601':
        agency_array = ['1601', '016']
    elif agency_code == '1125':
        agency_array = ['1125', '011']
    elif agency_code == '1100':
        agency_array = ['1100', '256']

    # Save the ATA filter
    agency_filters = []
    if not agency_array:
        agency_filters.append(tas_gtas.c.allocation_transfer_agency == agency_code)
    else:
        agency_filters.append(tas_gtas.c.allocation_transfer_agency.in_(agency_array))

    # Save the AID filter
    if agency_code in ['097', '020', '077', '089']:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier.in_(agency_array)))
    elif agency_code == '1100':
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   or_(tas_gtas.c.agency_identifier == '256',
                                       tas_gtas.c.fr_entity_type == '1100')))
    elif not frec_provided:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier == agency_code))
    else:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.fr_entity_type == agency_code))

    # If we're checking a CGAC, we want to filter on all of the related FRECs for AID 011
    if frec_list:
        agency_filters.append(and_(tas_gtas.c.allocation_transfer_agency.is_(None),
                                   tas_gtas.c.agency_identifier == '011',
                                   tas_gtas.c.fr_entity_type.in_(frec_list)))

    rows = initial_query(session, tas_gtas.c, year).\
        filter(func.coalesce(tas_gtas.c.financial_indicator2, '') != 'F').\
        filter(or_(*agency_filters)).\
        group_by(tas_gtas.c.allocation_transfer_agency,
                 tas_gtas.c.agency_identifier,
                 tas_gtas.c.beginning_period_of_availa,
                 tas_gtas.c.ending_period_of_availabil,
                 tas_gtas.c.availability_type_code,
                 tas_gtas.c.main_account_code,
                 tas_gtas.c.sub_account_code)

    return rows

def tas_gtas_combo(session, period, year):
    """ Creates a combined list of TAS and GTAS data filtered by the given period/year

        Args:
            session: DB session
            period: The period for which to get GTAS data
            year: The year for which to get GTAS data

        Returns:
            A WITH clause to use with other queries
    """
    query = session.query(
        gtas_model.allocation_transfer_agency.label('allocation_transfer_agency'),
        gtas_model.agency_identifier.label('agency_identifier'),
        gtas_model.beginning_period_of_availa.label('beginning_period_of_availa'),
        gtas_model.ending_period_of_availabil.label('ending_period_of_availabil'),
        gtas_model.availability_type_code.label('availability_type_code'),
        gtas_model.main_account_code.label('main_account_code'),
        gtas_model.sub_account_code.label('sub_account_code'),
        gtas_model.amount.label('amount'),
        gtas_model.line.label('line'),
        tas_model.financial_indicator2.label('financial_indicator2'),
        tas_model.fr_entity_type.label('fr_entity_type')).\
        join(tas_model, gtas_model.tas == func.concat(func.coalesce(tas_model.allocation_transfer_agency, '000'),
                                                      func.coalesce(tas_model.agency_identifier, '000'),
                                                      func.coalesce(tas_model.beginning_period_of_availa, '0000'),
                                                      func.coalesce(tas_model.ending_period_of_availabil, '0000'),
                                                      func.coalesce(tas_model.availability_type_code, ' '),
                                                      func.coalesce(tas_model.main_account_code, '0000'),
                                                      func.coalesce(tas_model.sub_account_code, '000'))).\
        filter(gtas_model.period == period).\
        filter(gtas_model.fiscal_year == year)
    return query.cte('tas_gtas')

def initial_query(session, model, year):
    """ Creates the initial query for D2 files.

        Args:
            session: The current DB session
            model: subquery model to get data from
            year: the year for which to get data

        Returns:
            The base query (a select from the tas/gtas tables with the specified columns).
    """
    budget_authority_line_max = 1042 if year <= 2020 else 1066
    return session.query(
        model.allocation_transfer_agency,
        model.agency_identifier,
        model.beginning_period_of_availa,
        model.ending_period_of_availabil,
        model.availability_type_code,
        model.main_account_code,
        model.sub_account_code,
        func.sum(case([(model.line == 1910, model.amount)], else_=0)).label('total_budgetary_resources_cpe'),
        func.sum(case([(model.line.in_([1160, 1180, 1260, 1280]), model.amount)],
                      else_=0)).label('budget_authority_appropria_cpe'),
        func.sum(case([(model.line == 1000, model.amount)], else_=0)).label('budget_authority_unobligat_fyb'),
        func.sum(case([(and_(model.line >= 1010, model.line <= budget_authority_line_max), model.amount)],
                      else_=0)).label('adjustments_to_unobligated_cpe'),
        func.sum(case([(model.line.in_([1540, 1640, 1340, 1440, 1750, 1850]), model.amount)],
                      else_=0)).label('other_budgetary_resources_cpe'),
        func.sum(case([(model.line.in_([1540, 1640]), model.amount)], else_=0)).label('contract_authority_amount_cpe'),
        func.sum(case([(model.line.in_([1340, 1440]), model.amount)], else_=0)).label('borrowing_authority_amount_cpe'),
        func.sum(case([(model.line.in_([1750, 1850]), model.amount)], else_=0)).label('spending_authority_from_of_cpe'),
        func.sum(case([(model.line == 2500, model.amount)], else_=0)).label('status_of_budgetary_resour_cpe'),
        func.sum(case([(model.line == 2190, model.amount)], else_=0)).label('obligations_incurred_total_cpe'),
        func.sum(case([(model.line == 3020, model.amount)], else_=0)).label('gross_outlay_amount_by_tas_cpe'),
        func.sum(case([(model.line == 2490, model.amount)], else_=0)).label('unobligated_balance_cpe'),
        func.sum(case([(model.line.in_([1021, 1033]), model.amount)], else_=0)).label('deobligations_recoveries_r_cpe'))
<commit_msg>Adding gtas_status column to file A generation<commit_after>import datetime

from collections import OrderedDict
from sqlalchemy import or_, and_, func, null
from sqlalchemy.orm import outerjoin
from sqlalchemy.sql.expression import case, literal_column
from dataactbroker.helpers.generic_helper import generate_raw_quoted_query

from dataactcore.models.domainModels import SF133, TASLookup, CGAC, FREC, TASFailedEdits
from dataactcore.models.jobModels import SubmissionWindowSchedule

gtas_model = SF133
tas_model = TASLookup
fail_model = TASFailedEdits

mapping = OrderedDict([
    ('allocation_transfer_agency', ['AllocationTransferAgencyIdentifier']),
    ('agency_identifier', ['AgencyIdentifier']),
    ('beginning_period_of_availa', ['BeginningPeriodOfAvailability']),
    ('ending_period_of_availabil', ['EndingPeriodOfAvailability']),
    ('availability_type_code', ['AvailabilityTypeCode']),
    ('main_account_code', ['MainAccountCode']),
    ('sub_account_code', ['SubAccountCode']),
    ('total_budgetary_resources_cpe', ['TotalBudgetaryResources_CPE']),
    ('budget_authority_appropria_cpe', ['BudgetAuthorityAppropriatedAmount_CPE']),
    ('budget_authority_unobligat_fyb', ['BudgetAuthorityUnobligatedBalanceBroughtForward_FYB']),
    ('adjustments_to_unobligated_cpe', ['AdjustmentsToUnobligatedBalanceBroughtForward_CPE']),
    ('other_budgetary_resources_cpe', ['OtherBudgetaryResourcesAmount_CPE']),
    ('contract_authority_amount_cpe', ['ContractAuthorityAmountTotal_CPE']),
    ('borrowing_authority_amount_cpe', ['BorrowingAuthorityAmountTotal_CPE']),
    ('spending_authority_from_of_cpe', ['SpendingAuthorityfromOffsettingCollectionsAmountTotal_CPE']),
    ('status_of_budgetary_resour_cpe', ['StatusOfBudgetaryResourcesTotal_CPE']),
    ('obligations_incurred_total_cpe', ['ObligationsIncurredTotalByTAS_CPE']),
    ('gross_outlay_amount_by_tas_cpe', ['GrossOutlayAmountByTAS_CPE']),
    ('unobligated_balance_cpe', ['UnobligatedBalance_CPE']),
    ('deobligations_recoveries_r_cpe', ['DeobligationsRecoveriesRefundsByTAS_CPE']),
    ('gtas_status', ['GTASStatus'])
])
db_columns = [key for key in mapping]

def query_data(session, agency_code, period, year):
    """ Request A file data

        Args:
            session: DB session
            agency_code: FREC or CGAC code for generation
            period: The period for which to get GTAS data
            year: The year for which to get GTAS data

        Returns:
            The rows using the provided dates for the given agency.
    """
    # set a boolean to determine if the original agency code is frec or cgac
    frec_provided = len(agency_code) == 4
    tas_gtas = tas_gtas_combo(session, period, year)
    tas_gtas_fail = failed_edits_details(session, tas_gtas, period, year)
    # Make a list of FRECs to compare to for 011 AID entries
    frec_list = []
    if not frec_provided:
        frec_list = session.query(FREC.frec_code).select_from(outerjoin(CGAC, FREC, CGAC.cgac_id == FREC.cgac_id)).\
            filter(CGAC.cgac_code == agency_code).all()
        # Put the frec list in a format that can be read by a filter
        frec_list = [frec.frec_code for frec in frec_list]
    # Group agencies together that need to be grouped
    # NOTE: If these change, update A33.1 to match
    agency_array = []
    if agency_code == '097':
        agency_array = ['017', '021', '057', '097']
    elif agency_code == '020':
        agency_array = ['020', '580', '373']
    elif agency_code == '077':
        agency_array = ['077', '071']
    elif agency_code == '089':
        agency_array = ['089', '486']
    elif agency_code == '1601':
        agency_array = ['1601', '016']
    elif agency_code == '1125':
        agency_array = ['1125', '011']
    elif agency_code == '1100':
        agency_array = ['1100', '256']

    # Save the ATA filter
    agency_filters = []
    if not agency_array:
        agency_filters.append(tas_gtas_fail.c.allocation_transfer_agency == agency_code)
    else:
        agency_filters.append(tas_gtas_fail.c.allocation_transfer_agency.in_(agency_array))

    # Save the AID filter
    if agency_code in ['097', '020', '077', '089']:
        agency_filters.append(and_(tas_gtas_fail.c.allocation_transfer_agency.is_(None),
                                   tas_gtas_fail.c.agency_identifier.in_(agency_array)))
    elif agency_code == '1100':
        agency_filters.append(and_(tas_gtas_fail.c.allocation_transfer_agency.is_(None),
                                   or_(tas_gtas_fail.c.agency_identifier == '256',
                                       tas_gtas_fail.c.fr_entity_type == '1100')))
    elif not frec_provided:
        agency_filters.append(and_(tas_gtas_fail.c.allocation_transfer_agency.is_(None),
                                   tas_gtas_fail.c.agency_identifier == agency_code))
    else:
        agency_filters.append(and_(tas_gtas_fail.c.allocation_transfer_agency.is_(None),
                                   tas_gtas_fail.c.fr_entity_type == agency_code))

    # If we're checking a CGAC, we want to filter on all of the related FRECs for AID 011
    if frec_list:
        agency_filters.append(and_(tas_gtas_fail.c.allocation_transfer_agency.is_(None),
                                   tas_gtas_fail.c.agency_identifier == '011',
                                   tas_gtas_fail.c.fr_entity_type.in_(frec_list)))

    rows = initial_query(session, tas_gtas_fail.c, year).\
        filter(func.coalesce(tas_gtas_fail.c.financial_indicator2, '') != 'F').\
        filter(or_(*agency_filters)).\
        group_by(tas_gtas_fail.c.allocation_transfer_agency,
                 tas_gtas_fail.c.agency_identifier,
                 tas_gtas_fail.c.beginning_period_of_availa,
                 tas_gtas_fail.c.ending_period_of_availabil,
                 tas_gtas_fail.c.availability_type_code,
                 tas_gtas_fail.c.main_account_code,
                 tas_gtas_fail.c.sub_account_code,
                 tas_gtas_fail.c.gtas_status)

    return rows

def tas_gtas_combo(session, period, year):
    """ Creates a combined list of TAS and GTAS data filtered by the given period/year

        Args:
            session: DB session
            period: The period for which to get GTAS data
            year: The year for which to get GTAS data

        Returns:
            A WITH clause to use with other queries
    """
    query = session.query(
        gtas_model.allocation_transfer_agency.label('allocation_transfer_agency'),
        gtas_model.agency_identifier.label('agency_identifier'),
        gtas_model.beginning_period_of_availa.label('beginning_period_of_availa'),
        gtas_model.ending_period_of_availabil.label('ending_period_of_availabil'),
        gtas_model.availability_type_code.label('availability_type_code'),
        gtas_model.main_account_code.label('main_account_code'),
        gtas_model.sub_account_code.label('sub_account_code'),
        gtas_model.amount.label('amount'),
        gtas_model.line.label('line'),
        tas_model.financial_indicator2.label('financial_indicator2'),
        tas_model.fr_entity_type.label('fr_entity_type'),
        gtas_model.tas.label('tas')).\
        join(tas_model, gtas_model.tas == tas_model.tas).\
        filter(gtas_model.period == period).\
        filter(gtas_model.fiscal_year == year)
    return query.cte('tas_gtas')

def failed_edits_details(session, tas_gtas, period, year):
    """ Creates a combined list of tas_failed_edits details

        Args:
            session: DB session
            tas_gtas: A CTE containing the filtered tas_gtas models
            period: The period for which to get GTAS data
            year: The year for which to get GTAS data

        Returns:
            A WITH clause to use with other queries
    """
    submission_period = session.query(SubmissionWindowSchedule).filter_by(period=period, year=year).one_or_none()
    # If submission period doesn't exist, the gtas_status will always be blank
    if not submission_period:
        query = session.query(tas_gtas, null().label('gtas_status'))
    elif submission_period.period_start > datetime.datetime.utcnow():
        query = session.query(tas_gtas, literal_column("'GTAS window open'").label('gtas_status'))
    else:
        has_period = session.query(fail_model).filter_by(period=period, fiscal_year=year).first() is not None
        # if the period doesn't exist in the TASFailingEdits model but does exist in our list, leave it blank
        if not has_period:
            query = session.query(tas_gtas, null().label('gtas_status'))
        else:
            query = session.query(
                    tas_gtas,
                    case([
                        (and_(fail_model.severity == 'fatal', fail_model.approved_override_exists.is_(False)),
                            literal_column("'failed fatal edit - no override'")),
                        (and_(fail_model.atb_submission_status == 'F',
                              fail_model.severity == 'fatal',
                              fail_model.approved_override_exists.is_(True)),
                            literal_column("'failed fatal edit - override'")),
                        (and_(fail_model.atb_submission_status == 'E',
                              fail_model.severity == 'fatal',
                              fail_model.approved_override_exists.is_(True)),
                            literal_column("'passed required edits - override'")),
                        (and_(fail_model.atb_submission_status == 'P',
                              fail_model.severity == 'fatal',
                              fail_model.approved_override_exists.is_(True)),
                            literal_column("'pending certification - override'")),
                        (and_(fail_model.atb_submission_status == 'C',
                              fail_model.severity == 'fatal',
                              fail_model.approved_override_exists.is_(True)),
                            literal_column("'certified - override'")),
                        (fail_model.edit_id.isnot(None), literal_column("'passed required edits - override'"))
                        ],
                        else_=literal_column("'passed required edits'")).label('gtas_status')).\
                join(fail_model, and_(tas_gtas.c.tas == fail_model.tas,
                                      fail_model.fiscal_year == year,
                                      fail_model.period == period), isouter=True)
    return query.cte('tas_gtas_fail')

def initial_query(session, model, year):
    """ Creates the initial query for D2 files.

        Args:
            session: The current DB session
            model: subquery model to get data from
            year: the year for which to get data

        Returns:
            The base query (a select from the tas/gtas tables with the specified columns).
    """
    budget_authority_line_max = 1042 if year <= 2020 else 1066
    return session.query(
        model.allocation_transfer_agency,
        model.agency_identifier,
        model.beginning_period_of_availa,
        model.ending_period_of_availabil,
        model.availability_type_code,
        model.main_account_code,
        model.sub_account_code,
        func.sum(case([(model.line == 1910, model.amount)], else_=0)).label('total_budgetary_resources_cpe'),
        func.sum(case([(model.line.in_([1160, 1180, 1260, 1280]), model.amount)],
                      else_=0)).label('budget_authority_appropria_cpe'),
        func.sum(case([(model.line == 1000, model.amount)], else_=0)).label('budget_authority_unobligat_fyb'),
        func.sum(case([(and_(model.line >= 1010, model.line <= budget_authority_line_max), model.amount)],
                      else_=0)).label('adjustments_to_unobligated_cpe'),
        func.sum(case([(model.line.in_([1540, 1640, 1340, 1440, 1750, 1850]), model.amount)],
                      else_=0)).label('other_budgetary_resources_cpe'),
        func.sum(case([(model.line.in_([1540, 1640]), model.amount)], else_=0)).label('contract_authority_amount_cpe'),
        func.sum(case([(model.line.in_([1340, 1440]), model.amount)], else_=0)).label('borrowing_authority_amount_cpe'),
        func.sum(case([(model.line.in_([1750, 1850]), model.amount)], else_=0)).label('spending_authority_from_of_cpe'),
        func.sum(case([(model.line == 2500, model.amount)], else_=0)).label('status_of_budgetary_resour_cpe'),
        func.sum(case([(model.line == 2190, model.amount)], else_=0)).label('obligations_incurred_total_cpe'),
        func.sum(case([(model.line == 3020, model.amount)], else_=0)).label('gross_outlay_amount_by_tas_cpe'),
        func.sum(case([(model.line == 2490, model.amount)], else_=0)).label('unobligated_balance_cpe'),
        func.sum(case([(model.line.in_([1021, 1033]), model.amount)], else_=0)).label('deobligations_recoveries_r_cpe'),
        model.gtas_status.label('gtas_status'))
<|endoftext|>


<commit_before>#!/usr/bin/env python 3
from accessoryFunctions.accessoryFunctions import make_path
from databasesetup import rest_auth_class
from Bio import SeqIO
from argparse import ArgumentParser
from glob import glob
import os

__author__ = 'adamkoziol'

class Get(object):

    def getrmlsthelper(self):
        """
        Makes a system call to rest_auth.py, a Python script modified from
        https://github.com/kjolley/BIGSdb/tree/develop/scripts/test
        And downloads the most up-to-date rMLST profile and alleles
        """
        # Extract the path of the current script from the full path + file name
        homepath = os.path.split(os.path.abspath(__file__))[0]
        # Set the path/name of the folder to contain the new alleles and profile
        newfolder = os.path.join(self.path, self.analysistype)
        # Create the path
        make_path(newfolder)
        # Create arguments to feed into the rest_auth_class script
        args = ArgumentParser
        args.secret_file = os.path.join(homepath, 'secret.txt')
        args.file_path = homepath
        args.output_path = newfolder
        args.logging = self.logging
        rmlst = rest_auth_class.REST(args)
        # Download the profile and alleles
        rmlst.main()
        # Get the new alleles into a list, and create the combinedAlleles file
        alleles = glob(os.path.join(newfolder, '*.tfa'))
        self.combinealleles(newfolder, alleles)

    def combinealleles(self, allelepath, alleles):
        self.logging.info('Creating combined rMLST allele file')
        with open(os.path.join(allelepath, 'rMLST_combined.fasta'), 'w') as combinedfile:
            # Open each allele file
            for allele in sorted(alleles):
                # with open(allele, 'rU') as fasta:
                for record in SeqIO.parse(open(allele, "rU"), "fasta"):
                    # Extract the sequence record from each entry in the multifasta
                    # Replace and dashes in the record.id with underscores
                    record.id = record.id.replace('-', '_')
                    # Remove and dashes or 'N's from the sequence data - makeblastdb can't handle sequences
                    # with gaps
                    # noinspection PyProtectedMember
                    record.seq._data = record.seq._data.replace('-', '').replace('N', '')
                    # Clear the name and description attributes of the record
                    record.name = ''
                    record.description = ''
                    # Write each record to the combined file
                    SeqIO.write(record, combinedfile, 'fasta')

    def __init__(self, args):
        self.path = os.path.join(args.path)
        self.logging = args.logging
        self.analysistype = 'rMLST'
        self.getrmlsthelper()

if __name__ == '__main__':
    # Argument parser for user-inputted values, and a nifty help menu
    # Parser for arguments
    parser = ArgumentParser(description='')
    parser.add_argument('path',
                        help='Specify input directory')

    # Get the arguments into an object
    arguments = parser.parse_args()
    # Run the script
    Get(arguments)
<commit_msg>Modified get_rmlst to not assume that the credentials will be in the same path as the script itself. They are now an argument to be provided<commit_after>#!/usr/bin/env python 3
from accessoryFunctions.accessoryFunctions import make_path
from databasesetup import rest_auth_class
from Bio import SeqIO
from argparse import ArgumentParser
from glob import glob
import os

__author__ = 'adamkoziol'

class Get(object):

    def getrmlsthelper(self):
        """
        Makes a system call to rest_auth.py, a Python script modified from
        https://github.com/kjolley/BIGSdb/tree/develop/scripts/test
        And downloads the most up-to-date rMLST profile and alleles
        """
        # Set the path/name of the folder to contain the new alleles and profile
        newfolder = os.path.join(self.path, self.analysistype)
        # Create the path
        make_path(newfolder)
        # Create arguments to feed into the rest_auth_class script
        args = ArgumentParser
        args.secret_file = os.path.join(self.credentials, 'secret.txt')
        args.file_path = os.path.join(self.credentials)
        args.output_path = newfolder
        args.logging = self.logging
        rmlst = rest_auth_class.REST(args)
        # Download the profile and alleles
        rmlst.main()
        # Get the new alleles into a list, and create the combinedAlleles file
        alleles = glob(os.path.join(newfolder, '*.tfa'))
        self.combinealleles(newfolder, alleles)

    def combinealleles(self, allelepath, alleles):
        self.logging.info('Creating combined rMLST allele file')
        with open(os.path.join(allelepath, 'rMLST_combined.fasta'), 'w') as combinedfile:
            # Open each allele file
            for allele in sorted(alleles):
                # with open(allele, 'rU') as fasta:
                for record in SeqIO.parse(open(allele, "rU"), "fasta"):
                    # Extract the sequence record from each entry in the multifasta
                    # Replace and dashes in the record.id with underscores
                    record.id = record.id.replace('-', '_')
                    # Remove and dashes or 'N's from the sequence data - makeblastdb can't handle sequences
                    # with gaps
                    # noinspection PyProtectedMember
                    record.seq._data = record.seq._data.replace('-', '').replace('N', '')
                    # Clear the name and description attributes of the record
                    record.name = ''
                    record.description = ''
                    # Write each record to the combined file
                    SeqIO.write(record, combinedfile, 'fasta')

    def __init__(self, args):
        self.path = os.path.join(args.path)
        self.logging = args.logging
        self.analysistype = 'rMLST'
        self.credentials = args.credentials
        self.getrmlsthelper()

if __name__ == '__main__':
    # Argument parser for user-inputted values, and a nifty help menu
    # Parser for arguments
    parser = ArgumentParser(description='')
    parser.add_argument('path',
                        help='Specify input directory')

    # Get the arguments into an object
    arguments = parser.parse_args()
    # Run the script
    Get(arguments)
<|endoftext|>