<commit_before>''' imports from installed packages '''
from django.core.management.base import BaseCommand, CommandError

from django_mongokit import get_database

try:
  from bson import ObjectId
except ImportError:  # old pymongo
  from pymongo.objectid import ObjectId

''' imports from application folders/files '''
from gnowsys_ndf.ndf.models import Node

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

class Command(BaseCommand):

  help = " This script will add the new field(s) into already existing documents (only if they doesn't exists) in your database."

  def handle(self, *args, **options):
    collection = get_database()[Node.collection_name]
    # Keep latest fields to be added at top

    # Updates the value of object_cardinality to 100. So that teaches will behave as 1:M (one-to-many) relation.
    teaches = collection.Node.one({'_type': "RelationType", 'name': "teaches"})
    res = collection.update({'_id': teaches._id, 'object_cardinality': {'$ne': 100}},
            {'$set': {'object_cardinality': 100}},
            upsert=False, multi=False
        )
    if res["updatedExisting"]:
        print "\n 'teaches' RelationType updated with object_cardinality: 100. Changed document: ", res['n']
    else:
        print "\n 'teaches' RelationType: no need to update."

    # Replacing object_type of "has_course" relationship from "NUSSD Course" to "Announced Course"
    ann_course = collection.Node.one({'_type': "GSystemType", 'name': "Announced Course"})
    res = collection.update({'_type': "RelationType", 'name': "has_course"},
            {'$set': {'object_type': [ann_course._id]}},
            upsert=False, multi=False
          )
    print "\n Replaced object_type of 'has_course' relationship from 'NUSSD Course' to 'Announced Course'."

    # Adds "relation_set" field (with default value as []) to all documents belonging to GSystems.
    res = collection.update({'_type': {'$nin': ["MetaType", "GSystemType", "RelationType", "AttributeType", "GRelation", "GAttribute", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'relation_set': {'$exists': False}},
                            {'$set': {'relation_set': []}},
                            upsert=False, multi=True
    )
    print "\n 'relation_set' field added to following no. of documents: ", res['n']

    # Adds "attribute_set" field (with default value as []) to all documents belonging to GSystems.
    res = collection.update({'_type': {'$nin': ["MetaType", "GSystemType", "RelationType", "AttributeType", "GRelation", "GAttribute", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'attribute_set': {'$exists': False}},
                            {'$set': {'attribute_set': []}},
                            upsert=False, multi=True
    )
    print "\n 'attribute_set' field added to following no. of documents: ", res['n']

    # Adds "license" field (with default value as "") to all documents belonging to GSystems (except Author).
    res = collection.update({'_type': {'$nin': ["MetaType", "Author", "GSystemType", "RelationType", "AttributeType", "GRelation", "GAttribute", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'license': {'$exists': False}},
                            {'$set': {'license': ""}},
                            upsert=False, multi=True
    )
    print "\n 'license' field added to following no. of documents: ", res['n']

    # Adding "Agency_type" field adding to group documents with default values
    res = collection.update({'_type': {'$in': ['Group']}, 'agency_type': {'$exists': False}},
                            {'$set': {'agency_type': "Project" }},
                            upsert=False, multi=True
    )
    if res['n']:
           print "\n 'agency_type' field added to 'Group' documents totalling to : ", res['n']

    # Adding "Agency_type" field adding to author documents with default values
    res = collection.update({'_type': {'$in': ['Author']}, 'agency_type': {'$exists': False}},
                            {'$set': {'agency_type': "Others" }},
                            upsert=False, multi=True
    )
    if res['n']:
           print "\n 'agency_type' field added to 'Author' documents totalling to : ", res['n']

    # Modify language field with unicode value if any document has language with dict datatype
    res = collection.update({'language': {}},
                            {'$set': {'language': u""}},
                            upsert=False, multi=True
    )

    # Removing existing "cr_or_xcr" field with no default value
    res = collection.update({'_type': {'$in': ['Group']}, 'cr_or_xcr': {'$exists': True}},
                            {'$unset': {'cr_or_xcr': False }},
                            upsert=False, multi=True
    )
    if res['n']:
           print "\n Already existing 'cr_or_xcr' field removed from documents totalling to : ", res['n']

    # Adding "curricular" field with no default value
    res = collection.update({'_type': {'$in': ['Group']}, 'curricular': {'$exists': False}},
                            {'$set': {'curricular': False }},
                            upsert=False, multi=True
    )
    print "\n 'curricular' field added to all Group documents totalling to : ", res['n']

    # Removing existing "partners" field with no default value
    res = collection.update({'_type': {'$in': ['Group']}, 'partners': {'$exists': True}},
                            {'$unset': {'partners': False }},
                            upsert=False, multi=True
    )
    if res['n']:
           print "\n Already existing 'partners' field removed from documents totalling to : ", res['n']

    # Adding "partner" field with no default value
    res = collection.update({'_type': {'$in': ['Group']}, 'partner': {'$exists': False}},
                            {'$set': {'partner': False }},
                            upsert=False, multi=True
    )
    print "\n 'partner' field added to all Group documents totalling to : ", res['n']

    # Adding "preferred_languages" field with no default value
    res = collection.update({'_type': {'$in': ['Author']}, 'preferred_languages': {'$exists': False}},
                            {'$set': {'preferred_languages': {}}},
                            upsert=False, multi=True
    )
    print "\n 'preferred_languages' field added to all author documents totalling to : ", res['n']

    # Adding "rating" field with no default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'rating': {'$exists': False}},
                            {'$set': {'rating': []}},
                            upsert=False, multi=True
    )
    print "\n 'rating' field added to following no. of documents: ", res['n']

    # Adds 'subject_scope', 'attribute_type_scope', 'object_value_scope' field (with default value as "") to all documents which belongs to GAttribute
    res = collection.update({'_type': {'$in': ["Group", "Author"]}, 'group_admin': {'$exists': False}},
                            {'$set': {'group_admin': []}},
                            upsert=False, multi=True
    )
    print "\n 'group_admin' field added to following no. of documents: ", res['n']

    # Adds 'subject_scope', 'attribute_type_scope', 'object_value_scope' field (with default value as "") to all documents which belongs to GAttribute
    res = collection.update({'_type': "GAttribute", 'subject_scope': {'$exists': False}, 'attribute_type_scope': {'$exists': False}, 'object_value_scope': {'$exists': False}},
                            {'$set': {'subject_scope':"", 'attribute_type_scope':"", 'object_value_scope': ""}},
                            upsert=False, multi=True
    )
    print "\n 'subject_scope', 'attribute_type_scope', 'object_value_scope' fields added to following no. of documents: ", res['n']

    # Adds 'subject_scope', 'relation_type_scope', 'right_subject_scope' field (with default value as "") to all documents which belongs to GRelation
    res = collection.update({'_type': "GRelation", 'subject_scope': {'$exists': False}, 'relation_type_scope': {'$exists': False}, 'right_subject_scope': {'$exists': False}},
                            {'$set': {'subject_scope':"", 'relation_type_scope':"", 'right_subject_scope': ""}},
                            upsert=False, multi=True
    )
    print "\n 'subject_scope', 'relation_type_scope', 'right_subject_scope' fields added to following no. of documents: ", res['n']

    # Adds "annotations" field (with default value as []) to all documents belonging to GSystems
    res = collection.update({'_type': {'$nin': ["MetaType", "GSystemType", "RelationType", "AttributeType", "GRelation", "GAttribute", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'annotations': {'$exists': False}},
                            {'$set': {'annotations': []}},
                            upsert=False, multi=True
    )
    print "\n annotations field added to following no. of documents: ", res['n']

    # Adds "group_set" field (with default value as []) to all documents except those which belongs to either GAttribute or GRelation
    res = collection.update({'_type': {'$nin': ["GAttribute", "GRelation", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'group_set': {'$exists': False}},
                            {'$set': {'group_set': []}},
                            upsert=False, multi=True
    )
    print "\n group_set field added to following no. of documents: ", res['n']

    # Adds "property_order" field (with default value as []) to all documents except those which belongs to either GAttribute or GRelation
    res = collection.update({'_type': {'$nin': ["GAttribute", "GRelation", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'property_order': {'$exists': False}},
                            {'$set': {'property_order': []}},
                            upsert=False, multi=True
    )
    print "\n property_order field added to following no. of documents: ", res['n']

    # Adding "modified_by" field with None as it's default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'modified_by': {'$exists': False}},
                            {'$set': {'modified_by': None}},
                            upsert=False, multi=True
    )
    print "\n modified_by field added to following no. of documents: ", res['n']

    # Adding "complex_data_type" field with empty list as it's default value
    res = collection.update({'_type': 'AttributeType', 'complex_data_type': {'$exists': False}},
                            {'$set': {'complex_data_type': []}},
                            upsert=False, multi=True
    )
    print "\n complex_data_type field added to following no. of documents: ", res['n']

    # Adding "post_node" field with empty list as it's default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'post_node': {'$exists': False}},
                            {'$set': {'post_node': []}},
                            upsert=False, multi=True
    )
    print "\n post_node field added to following no. of documents: ", res['n']

    # Adding "collection_set" field with empty list as it's default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'collection_set': {'$exists': False}},
                            {'$set': {'collection_set': []}},
                            upsert=False, multi=True
    )
    print "\n collection_set field added to following no. of documents: ", res['n']

    # Adding "location" field with no default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'location': {'$exists': False}},
                            {'$set': {'location': []}},
                            upsert=False, multi=True
    )
    print "\n location field added to following no. of documents: ", res['n'], "\n"

    # Adding "language" field with no default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'language': {'$exists': False}},
                            {'$set': {'language': unicode('')}},
                            upsert=False, multi=True
    )

    # Adding "access_policy" field
    # For Group documents, access_policy value is set depending upon their
    # group_type values, i.e. either PRIVATE/PUBLIC whichever is there
    collection.update({'_type': 'Group', 'group_type': 'PRIVATE'}, {'$set': {'access_policy': u"PRIVATE"}}, upsert=False, multi=True)
    collection.update({'_type': 'Group', 'group_type': 'PUBLIC'}, {'$set': {'access_policy': u"PUBLIC"}}, upsert=False, multi=True)

    # For Non-Group documents which doesn't consits of access_policy field, add it with PUBLIC as it's default value
    collection.update({'_type': {'$nin': ['Group', 'GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'access_policy': {'$exists': False}}, {'$set': {'access_policy': u"PUBLIC"}}, upsert=False, multi=True)

    collection.update({'_type': {'$nin': ['Group', 'GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'access_policy': {'$in': [None, "PUBLIC"]}}, {'$set': {'access_policy': u"PUBLIC"}}, upsert=False, multi=True)
    collection.update({'_type': {'$nin': ['Group', 'GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'access_policy': "PRIVATE"}, {'$set': {'access_policy': u"PRIVATE"}}, upsert=False, multi=True)

<commit_msg>bug resolved in sync_existing_documents<commit_after>''' imports from installed packages '''
from django.core.management.base import BaseCommand, CommandError

from django_mongokit import get_database

try:
  from bson import ObjectId
except ImportError:  # old pymongo
  from pymongo.objectid import ObjectId

''' imports from application folders/files '''
from gnowsys_ndf.ndf.models import Node

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

class Command(BaseCommand):

  help = " This script will add the new field(s) into already existing documents (only if they doesn't exists) in your database."

  def handle(self, *args, **options):
    collection = get_database()[Node.collection_name]
    # Keep latest fields to be added at top

    # Updates the value of object_cardinality to 100. So that teaches will behave as 1:M (one-to-many) relation.
    teaches = collection.Node.one({'_type': "RelationType", 'name': "teaches"})
    res = collection.update({'_id': teaches._id, 'object_cardinality': {'$ne': 100}},
            {'$set': {'object_cardinality': 100}},
            upsert=False, multi=False
        )
    if res["updatedExisting"]:
        print "\n 'teaches' RelationType updated with object_cardinality: 100. Changed document: ", res['n']
    else:
        print "\n 'teaches' RelationType: no need to update."

    # Replacing object_type of "has_course" relationship from "NUSSD Course" to "Announced Course"
    ann_course = collection.Node.one({'_type': "GSystemType", 'name': "Announced Course"})
    if ann_course:
        res = collection.update({'_type': "RelationType", 'name': "has_course"},
                {'$set': {'object_type': [ann_course._id]}},
                upsert=False, multi=False
              )
        print "\n Replaced object_type of 'has_course' relationship from 'NUSSD Course' to 'Announced Course'."

    # Adds "relation_set" field (with default value as []) to all documents belonging to GSystems.
    res = collection.update({'_type': {'$nin': ["MetaType", "GSystemType", "RelationType", "AttributeType", "GRelation", "GAttribute", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'relation_set': {'$exists': False}},
                            {'$set': {'relation_set': []}},
                            upsert=False, multi=True
    )
    print "\n 'relation_set' field added to following no. of documents: ", res['n']

    # Adds "attribute_set" field (with default value as []) to all documents belonging to GSystems.
    res = collection.update({'_type': {'$nin': ["MetaType", "GSystemType", "RelationType", "AttributeType", "GRelation", "GAttribute", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'attribute_set': {'$exists': False}},
                            {'$set': {'attribute_set': []}},
                            upsert=False, multi=True
    )
    print "\n 'attribute_set' field added to following no. of documents: ", res['n']

    # Adds "license" field (with default value as "") to all documents belonging to GSystems (except Author).
    res = collection.update({'_type': {'$nin': ["MetaType", "Author", "GSystemType", "RelationType", "AttributeType", "GRelation", "GAttribute", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'license': {'$exists': False}},
                            {'$set': {'license': ""}},
                            upsert=False, multi=True
    )
    print "\n 'license' field added to following no. of documents: ", res['n']

    # Adding "Agency_type" field adding to group documents with default values
    res = collection.update({'_type': {'$in': ['Group']}, 'agency_type': {'$exists': False}},
                            {'$set': {'agency_type': "Project" }},
                            upsert=False, multi=True
    )
    if res['n']:
           print "\n 'agency_type' field added to 'Group' documents totalling to : ", res['n']

    # Adding "Agency_type" field adding to author documents with default values
    res = collection.update({'_type': {'$in': ['Author']}, 'agency_type': {'$exists': False}},
                            {'$set': {'agency_type': "Others" }},
                            upsert=False, multi=True
    )
    if res['n']:
           print "\n 'agency_type' field added to 'Author' documents totalling to : ", res['n']

    # Modify language field with unicode value if any document has language with dict datatype
    res = collection.update({'language': {}},
                            {'$set': {'language': u""}},
                            upsert=False, multi=True
    )

    # Removing existing "cr_or_xcr" field with no default value
    res = collection.update({'_type': {'$in': ['Group']}, 'cr_or_xcr': {'$exists': True}},
                            {'$unset': {'cr_or_xcr': False }},
                            upsert=False, multi=True
    )
    if res['n']:
           print "\n Already existing 'cr_or_xcr' field removed from documents totalling to : ", res['n']

    # Adding "curricular" field with no default value
    res = collection.update({'_type': {'$in': ['Group']}, 'curricular': {'$exists': False}},
                            {'$set': {'curricular': False }},
                            upsert=False, multi=True
    )
    print "\n 'curricular' field added to all Group documents totalling to : ", res['n']

    # Removing existing "partners" field with no default value
    res = collection.update({'_type': {'$in': ['Group']}, 'partners': {'$exists': True}},
                            {'$unset': {'partners': False }},
                            upsert=False, multi=True
    )
    if res['n']:
           print "\n Already existing 'partners' field removed from documents totalling to : ", res['n']

    # Adding "partner" field with no default value
    res = collection.update({'_type': {'$in': ['Group']}, 'partner': {'$exists': False}},
                            {'$set': {'partner': False }},
                            upsert=False, multi=True
    )
    print "\n 'partner' field added to all Group documents totalling to : ", res['n']

    # Adding "preferred_languages" field with no default value
    res = collection.update({'_type': {'$in': ['Author']}, 'preferred_languages': {'$exists': False}},
                            {'$set': {'preferred_languages': {}}},
                            upsert=False, multi=True
    )
    print "\n 'preferred_languages' field added to all author documents totalling to : ", res['n']

    # Adding "rating" field with no default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'rating': {'$exists': False}},
                            {'$set': {'rating': []}},
                            upsert=False, multi=True
    )
    print "\n 'rating' field added to following no. of documents: ", res['n']

    # Adds 'subject_scope', 'attribute_type_scope', 'object_value_scope' field (with default value as "") to all documents which belongs to GAttribute
    res = collection.update({'_type': {'$in': ["Group", "Author"]}, 'group_admin': {'$exists': False}},
                            {'$set': {'group_admin': []}},
                            upsert=False, multi=True
    )
    print "\n 'group_admin' field added to following no. of documents: ", res['n']

    # Adds 'subject_scope', 'attribute_type_scope', 'object_value_scope' field (with default value as "") to all documents which belongs to GAttribute
    res = collection.update({'_type': "GAttribute", 'subject_scope': {'$exists': False}, 'attribute_type_scope': {'$exists': False}, 'object_value_scope': {'$exists': False}},
                            {'$set': {'subject_scope':"", 'attribute_type_scope':"", 'object_value_scope': ""}},
                            upsert=False, multi=True
    )
    print "\n 'subject_scope', 'attribute_type_scope', 'object_value_scope' fields added to following no. of documents: ", res['n']

    # Adds 'subject_scope', 'relation_type_scope', 'right_subject_scope' field (with default value as "") to all documents which belongs to GRelation
    res = collection.update({'_type': "GRelation", 'subject_scope': {'$exists': False}, 'relation_type_scope': {'$exists': False}, 'right_subject_scope': {'$exists': False}},
                            {'$set': {'subject_scope':"", 'relation_type_scope':"", 'right_subject_scope': ""}},
                            upsert=False, multi=True
    )
    print "\n 'subject_scope', 'relation_type_scope', 'right_subject_scope' fields added to following no. of documents: ", res['n']

    # Adds "annotations" field (with default value as []) to all documents belonging to GSystems
    res = collection.update({'_type': {'$nin': ["MetaType", "GSystemType", "RelationType", "AttributeType", "GRelation", "GAttribute", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'annotations': {'$exists': False}},
                            {'$set': {'annotations': []}},
                            upsert=False, multi=True
    )
    print "\n annotations field added to following no. of documents: ", res['n']

    # Adds "group_set" field (with default value as []) to all documents except those which belongs to either GAttribute or GRelation
    res = collection.update({'_type': {'$nin': ["GAttribute", "GRelation", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'group_set': {'$exists': False}},
                            {'$set': {'group_set': []}},
                            upsert=False, multi=True
    )
    print "\n group_set field added to following no. of documents: ", res['n']

    # Adds "property_order" field (with default value as []) to all documents except those which belongs to either GAttribute or GRelation
    res = collection.update({'_type': {'$nin': ["GAttribute", "GRelation", "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'property_order': {'$exists': False}},
                            {'$set': {'property_order': []}},
                            upsert=False, multi=True
    )
    print "\n property_order field added to following no. of documents: ", res['n']

    # Adding "modified_by" field with None as it's default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'modified_by': {'$exists': False}},
                            {'$set': {'modified_by': None}},
                            upsert=False, multi=True
    )
    print "\n modified_by field added to following no. of documents: ", res['n']

    # Adding "complex_data_type" field with empty list as it's default value
    res = collection.update({'_type': 'AttributeType', 'complex_data_type': {'$exists': False}},
                            {'$set': {'complex_data_type': []}},
                            upsert=False, multi=True
    )
    print "\n complex_data_type field added to following no. of documents: ", res['n']

    # Adding "post_node" field with empty list as it's default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'post_node': {'$exists': False}},
                            {'$set': {'post_node': []}},
                            upsert=False, multi=True
    )
    print "\n post_node field added to following no. of documents: ", res['n']

    # Adding "collection_set" field with empty list as it's default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'collection_set': {'$exists': False}},
                            {'$set': {'collection_set': []}},
                            upsert=False, multi=True
    )
    print "\n collection_set field added to following no. of documents: ", res['n']

    # Adding "location" field with no default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'location': {'$exists': False}},
                            {'$set': {'location': []}},
                            upsert=False, multi=True
    )
    print "\n location field added to following no. of documents: ", res['n'], "\n"

    # Adding "language" field with no default value
    res = collection.update({'_type': {'$nin': ['GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'language': {'$exists': False}},
                            {'$set': {'language': unicode('')}},
                            upsert=False, multi=True
    )

    # Adding "access_policy" field
    # For Group documents, access_policy value is set depending upon their
    # group_type values, i.e. either PRIVATE/PUBLIC whichever is there
    collection.update({'_type': 'Group', 'group_type': 'PRIVATE'}, {'$set': {'access_policy': u"PRIVATE"}}, upsert=False, multi=True)
    collection.update({'_type': 'Group', 'group_type': 'PUBLIC'}, {'$set': {'access_policy': u"PUBLIC"}}, upsert=False, multi=True)

    # For Non-Group documents which doesn't consits of access_policy field, add it with PUBLIC as it's default value
    collection.update({'_type': {'$nin': ['Group', 'GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'access_policy': {'$exists': False}}, {'$set': {'access_policy': u"PUBLIC"}}, upsert=False, multi=True)

    collection.update({'_type': {'$nin': ['Group', 'GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'access_policy': {'$in': [None, "PUBLIC"]}}, {'$set': {'access_policy': u"PUBLIC"}}, upsert=False, multi=True)
    collection.update({'_type': {'$nin': ['Group', 'GAttribute', 'GRelation', "ReducedDocs", "ToReduceDocs", "IndexedWordList", "node_holder"]}, 'access_policy': "PRIVATE"}, {'$set': {'access_policy': u"PRIVATE"}}, upsert=False, multi=True)

<|endoftext|>


<commit_before>from django.http import (HttpResponse,
                         HttpResponseNotFound,
                         HttpResponseServerError)
from whats_fresh_api.models import Vendor, Product, VendorProduct
from django.forms.models import model_to_dict
import json

def vendor_list(request):
    data = {}
    vendor_list = Vendor.objects.all()

    if len(vendor_list) == 0:
        data['error'] = {
            'debug': '',
            'status': True,
            'level': 'Error',
            'text': 'No Vendors found',
            'name': 'No Vendors'
        }
        return HttpResponseNotFound(
            json.dumps(data),
            content_type="application/json"
        )
    try:
        data['vendors'] = []
        for vendor in vendor_list:
            data['vendors'].append(model_to_dict(vendor, fields=[], exclude=[]))
            data['vendors'][-1]['phone'] = data['vendors'][-1]['phone'].national_number

            try:
                data['vendors'][-1]['phone'] = data['vendors'][-1]['phone'].national_number
            except AttributeError:
                data['vendors'][-1]['phone'] = None

            data['vendors'][-1]['created'] = str(vendor.created)
            data['vendors'][-1]['updated'] = str(vendor.modified)
            data['vendors'][-1]['ext'] = {}
            data['vendors'][-1]['story'] = data['vendors'][-1].pop('story_id')
            data['vendors'][-1]['id'] = vendor.id

            data['vendors'][-1]['lat'] = vendor.location.y
            data['vendors'][-1]['long'] = vendor.location.x

            del data['vendors'][-1]['products_preparations']

            vendor_products = vendor.vendorproduct_set.all()
            data['vendors'][-1]['products'] = []
            for vendor_product in vendor_products:
                product_data = {
                    'product_id': vendor_product.product_preparation.product.id,
                    'preparation_id': vendor_product.product_preparation.preparation.id,
                    'preparation': vendor_product.product_preparation.preparation.name,
                    'name': vendor_product.product_preparation.product.name
                }
                data['vendors'][-1]['products'].append(product_data)

        data['error'] = {
            'debug': None,
            'status': False,
            'level': None,
            'text': None,
            'name': None
        }
        return HttpResponse(json.dumps(data), content_type="application/json")
    except Exception as e:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Error',
            'text': e,
            'name': 'Unknown'
        }
        return HttpResponseServerError(
            json.dumps(data),
            content_type="application/json"
        )

def vendors_products(request, id=None):
    data = {}
    try:
        vendor_list = Vendor.objects.filter(
            vendorproduct__product_preparation__product__id__exact=id)
    except Exception as e:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Error',
            'text': 'Product id is invalid',
            'name': 'Invalid product'
        }
        return HttpResponseNotFound(
            json.dumps(data),
            content_type="application/json"
        )

    if len(vendor_list) == 0:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Important',
            'text': 'Could not find any vendors for product %s!' % id,
            'name': 'No Vendors for product %s' % id
        }
        return HttpResponse(
            json.dumps(data),
            content_type="application/json"
        )

    try:
        data['vendors'] = []
        for vendor in vendor_list:
            data['vendors'].append(model_to_dict(vendor, fields=[], exclude=[]))

            try:
                data['vendors'][-1]['phone'] = data['vendors'][-1]['phone'].national_number
            except AttributeError:
                data['vendors'][-1]['phone'] = None

            data['vendors'][-1]['created'] = str(vendor.created)
            data['vendors'][-1]['updated'] = str(vendor.modified)
            data['vendors'][-1]['ext'] = {}
            data['vendors'][-1]['id'] = vendor.id

            data['vendors'][-1]['story'] = data['vendors'][-1].pop('story_id')
            data['vendors'][-1]['lat'] = vendor.location.y
            data['vendors'][-1]['long'] = vendor.location.x

            del data['vendors'][-1]['products_preparations']

            vendor_products = vendor.vendorproduct_set.all()
            data['vendors'][-1]['products'] = []
            for vendor_product in vendor_products:
                product_data = {
                    'product_id': vendor_product.product_preparation.product.id,
                    'preparation_id': vendor_product.product_preparation.preparation.id,
                    'preparation': vendor_product.product_preparation.preparation.name,
                    'name': vendor_product.product_preparation.product.name
                }
                data['vendors'][-1]['products'].append(product_data)

        data['error'] = {
            'debug': None,
            'status': False,
            'level': None,
            'text': None,
            'name': None
        }
        return HttpResponse(json.dumps(data), content_type="application/json")

    except Exception as e:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Error',
            'text': 'Error {0} occurred processing the vendors for product {1}'.format(e, id),
            'name': 'Unknown'
        }
        return HttpResponseServerError(
            json.dumps(data),
            content_type="application/json"
        )

def vendor_details(request, id=None):
    data = {}

    try:
        vendor = Vendor.objects.get(id=id)
    except Exception as e:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Error',
            'text': 'Vendor id %s was not found.' % id,
            'name': 'Vendor Not Found'
        }
        return HttpResponseNotFound(
            json.dumps(data),
            content_type="application/json"
        )

    try:
        data = model_to_dict(vendor, fields=[], exclude=[])
<<<<<<< HEAD
=======

        data['story_id'] = vendor.story_id.id
        try:
            data['phone'] = vendor.phone.national_number
       except AttributeError:
            data['phone'] = None
>>>>>>> Write location-based tests

<<<<<<< HEAD
        data['story_id'] = vendor.story_id.id
        if data['phone']:
            data['phone'] = data['phone'].national_number
        else:
            data['phone'] = None
=======
        data['lat'] = vendor.location.y
        data['long'] = vendor.location.x
        del data['location']

>>>>>>> Update Vendor to use GeoDjango point rather than lat/long
        data['created'] = str(vendor.created)
        data['updated'] = str(vendor.modified)
        data['ext'] = {}
        data['story'] = data.pop('story_id')
        data['id'] = vendor.id

        del data['products_preparations']

        vendor_products = vendor.vendorproduct_set.all()
        data['products'] = []
        for vendor_product in vendor_products:
            product_data = {
                'product_id': vendor_product.product_preparation.product.id,
                'preparation_id': vendor_product.product_preparation.preparation.id,
                'preparation': vendor_product.product_preparation.preparation.name,
                'name': vendor_product.product_preparation.product.name
            }
            data['products'].append(product_data)

        data['error'] = {
            'debug': None,
            'status': False,
            'level': None,
            'text': None,
            'name': None
        }
        return HttpResponse(json.dumps(data), content_type="application/json")

    except Exception as e:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Error',
            'text': 'An unknown error occurred processing vendor %s'
            % id,
            'name': e
        }

        return HttpResponseServerError(
            json.dumps(data),
            content_type="application/json"
        )
<commit_msg>Rebase/merge off of develop<commit_after>from django.http import (HttpResponse,
                         HttpResponseNotFound,
                         HttpResponseServerError)
from whats_fresh_api.models import Vendor, Product, VendorProduct
from django.forms.models import model_to_dict
import json

def vendor_list(request):
    data = {}
    vendor_list = Vendor.objects.all()

    if len(vendor_list) == 0:
        data['error'] = {
            'debug': '',
            'status': True,
            'level': 'Error',
            'text': 'No Vendors found',
            'name': 'No Vendors'
        }
        return HttpResponseNotFound(
            json.dumps(data),
            content_type="application/json"
        )
    try:
        data['vendors'] = []
        for vendor in vendor_list:
            data['vendors'].append(
                model_to_dict(
                   vendor,
                   fields=[],
                   exclude=['location', 'phone', 'products_preparations']))

            try:
                data['vendors'][-1]['phone'] = vendor.phone.national_number
            except AttributeError:
                data['vendors'][-1]['phone'] = None

            data['vendors'][-1]['created'] = str(vendor.created)
            data['vendors'][-1]['updated'] = str(vendor.modified)
            data['vendors'][-1]['ext'] = {}
            data['vendors'][-1]['story'] = data['vendors'][-1].pop('story_id')
            data['vendors'][-1]['id'] = vendor.id

            data['vendors'][-1]['lat'] = vendor.location.y
            data['vendors'][-1]['long'] = vendor.location.x

            vendor_products = vendor.vendorproduct_set.all()
            data['vendors'][-1]['products'] = []
            for vendor_product in vendor_products:
                product_data = {
<<<<<<< HEAD
                    'product_id': vendor_product.product_preparation.product.id,
                    'preparation_id': vendor_product.product_preparation.preparation.id,
                    'preparation': vendor_product.product_preparation.preparation.name,
                    'name': vendor_product.product_preparation.product.name
=======
                    'id': vendor_product.product_preparation.product.id,
                    'name': vendor_product.product_preparation.product.name,
                    'preparation':
                        vendor_product.product_preparation.preparation.name
>>>>>>> Rebase/merge off of develop
                }
                data['vendors'][-1]['products'].append(product_data)

        data['error'] = {
            'debug': None,
            'status': False,
            'level': None,
            'text': None,
            'name': None
        }
        return HttpResponse(json.dumps(data), content_type="application/json")
    except Exception as e:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Error',
            'text': e,
            'name': 'Unknown'
        }
        return HttpResponseServerError(
            json.dumps(data),
            content_type="application/json"
        )

def vendors_products(request, id=None):
    data = {}
    try:
        vendor_list = Vendor.objects.filter(
            vendorproduct__product_preparation__product__id__exact=id)
    except Exception as e:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Error',
            'text': 'Product id is invalid',
            'name': 'Invalid product'
        }
        return HttpResponseNotFound(
            json.dumps(data),
            content_type="application/json"
        )

    if len(vendor_list) == 0:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Important',
            'text': 'Could not find any vendors for product %s!' % id,
            'name': 'No Vendors for product %s' % id
        }
        return HttpResponse(
            json.dumps(data),
            content_type="application/json"
        )

    try:
        data['vendors'] = []
        for vendor in vendor_list:
            data['vendors'].append(
                model_to_dict(
                    vendor,
                    fields=[],
                    exclude=['location', 'phone', 'products_preparations']))

            try:
                data['vendors'][-1]['phone'] = vendor.phone.national_number
            except AttributeError:
                data['vendors'][-1]['phone'] = None

            data['vendors'][-1]['created'] = str(vendor.created)
            data['vendors'][-1]['updated'] = str(vendor.modified)
            data['vendors'][-1]['ext'] = {}
            data['vendors'][-1]['id'] = vendor.id

            data['vendors'][-1]['story'] = data['vendors'][-1].pop('story_id')
            data['vendors'][-1]['lat'] = vendor.location.y
            data['vendors'][-1]['long'] = vendor.location.x

            vendor_products = vendor.vendorproduct_set.all()
            data['vendors'][-1]['products'] = []
            for vendor_product in vendor_products:
                product_data = {
<<<<<<< HEAD
                    'product_id': vendor_product.product_preparation.product.id,
                    'preparation_id': vendor_product.product_preparation.preparation.id,
                    'preparation': vendor_product.product_preparation.preparation.name,
                    'name': vendor_product.product_preparation.product.name
=======
                    'id': vendor_product.product_preparation.product.id,
                    'name': vendor_product.product_preparation.product.name,
                    'preparation':
                        vendor_product.product_preparation.preparation.name
>>>>>>> Rebase/merge off of develop
                }
                data['vendors'][-1]['products'].append(product_data)

        data['error'] = {
            'debug': None,
            'status': False,
            'level': None,
            'text': None,
            'name': None
        }
        return HttpResponse(json.dumps(data), content_type="application/json")

    except Exception as e:
        data['error'] = {
<<<<<<< HEAD
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Error',
            'text': 'Error {0} occurred processing the vendors for product {1}'.format(e, id),
            'name': 'Unknown'
=======
            'error_status': True,
            'error_level': 'Severe',
            'error_text': """Error {0} occurred processing the \
vendors for product {1}""".format(e, id),
            'error_name': 'Unknown'
>>>>>>> Rebase/merge off of develop
        }
        return HttpResponseServerError(
            json.dumps(data),
            content_type="application/json"
        )

def vendor_details(request, id=None):
    data = {}

    try:
        vendor = Vendor.objects.get(id=id)
    except Exception as e:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Error',
            'text': 'Vendor id %s was not found.' % id,
            'name': 'Vendor Not Found'
        }
        return HttpResponseNotFound(
            json.dumps(data),
            content_type="application/json"
        )

    try:
<<<<<<< HEAD
        data = model_to_dict(vendor, fields=[], exclude=[])
<<<<<<< HEAD
=======
=======
        data = model_to_dict(
            vendor, fields=[], exclude=[
                'location', 'phone', 'products_preparations'])
>>>>>>> Rebase/merge off of develop

        data['story_id'] = vendor.story_id.id
        try:
            data['phone'] = vendor.phone.national_number
        except AttributeError:
            data['phone'] = None
>>>>>>> Write location-based tests

<<<<<<< HEAD
        data['story_id'] = vendor.story_id.id
        if data['phone']:
            data['phone'] = data['phone'].national_number
        else:
            data['phone'] = None
=======
        data['lat'] = vendor.location.y
        data['long'] = vendor.location.x

>>>>>>> Update Vendor to use GeoDjango point rather than lat/long
        data['created'] = str(vendor.created)
        data['updated'] = str(vendor.modified)
        data['ext'] = {}
        data['story'] = data.pop('story_id')
        data['id'] = vendor.id

        vendor_products = vendor.vendorproduct_set.all()
        data['products'] = []
        for vendor_product in vendor_products:
            product_data = {
<<<<<<< HEAD
                'product_id': vendor_product.product_preparation.product.id,
                'preparation_id': vendor_product.product_preparation.preparation.id,
                'preparation': vendor_product.product_preparation.preparation.name,
                'name': vendor_product.product_preparation.product.name
=======
                'id': vendor_product.product_preparation.product.id,
                'name': vendor_product.product_preparation.product.name,
                'preparation':
                    vendor_product.product_preparation.preparation.name
>>>>>>> Rebase/merge off of develop
            }
            data['products'].append(product_data)

        data['error'] = {
            'debug': None,
            'status': False,
            'level': None,
            'text': None,
            'name': None
        }
        return HttpResponse(json.dumps(data), content_type="application/json")

    except Exception as e:
        data['error'] = {
            'debug': "{0}: {1}".format(type(e).__name__, str(e)),
            'status': True,
            'level': 'Error',
            'text': 'An unknown error occurred processing vendor %s'
            % id,
            'name': e
        }

        return HttpResponseServerError(
            json.dumps(data),
            content_type="application/json"
        )
<|endoftext|>


<commit_before>from aiida import load_dbenv
load_dbenv()
from aiida.orm import Code, DataFactory

StructureData = DataFactory('structure')
ParameterData = DataFactory('parameter')

import numpy as np
import os
import pymatgen

##########################
structure_id = 'mp-12668'
##########################

def get_supercell_size(structure, max_atoms=100):

    cell = np.array(structure.cell)

    num_atoms = len(structure.sites)

    print cell
    print num_atoms
    print np.linalg.norm(cell, axis=1)

    supercell_size = [1, 1, 1]

    while True:

        norm = np.linalg.norm(cell, axis=1)
        index = np.argmin(norm)
        supercell_size_test = list(supercell_size)
        supercell_size_test[index] += 1

        anum_atoms_supercell = num_atoms * np.prod(supercell_size_test)
        if anum_atoms_supercell > max_atoms:
            atoms_minus =  num_atoms * np.prod(supercell_size)
            atoms_plus = num_atoms * np.prod(supercell_size_test)

            if max_atoms - atoms_minus < atoms_plus - max_atoms:
                return supercell_size
            else:
                return supercell_size_test
        else:
            supercell_size = supercell_size_test

        print supercell_size
        cell = np.dot(cell.T, np.diag(supercell_size)).T
        print cell

def get_potential_labels(functional, symbol_list, ftype=None):
    potential_labels =[]
    for symbol in np.unique(symbol_list):
        psp_dir = os.environ['VASP_PSP_DIR'] + '/POT_GGA_PAW_' + functional
        all_labels = [f[7:-3] for f in os.listdir(psp_dir)  if os.path.isfile(os.path.join(psp_dir, f))]
        candidates = [ s for s in all_labels if symbol == s.split('_')[0]]
        if ftype is not None:
            final = [s for s in candidates if '_{}'.format(ftype) == s[-(len(ftype) + 1):]]
            if len(final) > 0:
                potential_labels.append(final[0])
                continue
        potential_labels.append(candidates[0])
    return potential_labels

rester = pymatgen.MPRester(os.environ['PMG_MAPI_KEY'])

pmg_structure = rester.get_structure_by_material_id(structure_id)
spa = pymatgen.symmetry.analyzer.SpacegroupAnalyzer(pmg_structure)

conventional = spa.get_conventional_standard_structure()
primitive = spa.get_primitive_standard_structure()

primitive_matrix = np.dot(np.linalg.inv(conventional.lattice_vectors()), primitive.lattice_vectors())
primitive_matrix = np.round(primitive_matrix, decimals=6).tolist()

structure = StructureData(pymatgen=conventional).store()
print structure
# crystal_system = spa.get_crystal_system()
# if crystal_system == 'hexagonal':
#     supercell = [[3, 0, 0],
#                  [0, 3, 0],
#                  [0, 0, 3]]
# else:
#     supercell = [[2, 0, 0],
#                  [0, 2, 0],
#                  [0, 0, 2]]

supercell_size = get_supercell_size(structure)
supercell = np.diag(supercell_size)

print ('Supercell shape: {}'.format(supercell_size))
exit()

incar_dict = {
    'NELMIN' : 5,
    'NELM'   : 100,
    'ENCUT'  : 400,
    'ALGO'   : 38,
    'ISMEAR' : 0,
    'SIGMA'  : 0.01,
    'GGA'    : 'PS'
}

pseudo_dict = {'functional': 'PBE',
               'symbols': get_potential_labels('PBE', conventional.symbol_set)}

# Monkhorst-pack
kpoints_dict = {'points': [2, 2, 2],
                'shift' : [0.0, 0.0, 0.0]}

machine_dict = {
    'num_machines': 1,
    'parallel_env':'mpi*',
    'tot_num_mpiprocs': 16}

phonopy_parameters = {'supercell': supercell,
                      'primitive': primitive_matrix,
                      'distance': 0.01,
                      'mesh': [20, 20, 20]}

wf_parameters = {
     'structure': structure,
     'phonopy_input': {'parameters': phonopy_parameters},
     'input_force': {'code': 'vasp541mpi@stern',
                    'parameters': incar_dict,
                    'resources': machine_dict,
                    'pseudo': pseudo_dict,
                    'kpoints': kpoints_dict},
     'input_optimize': {'code': 'vasp541mpi@stern',
                       'parameters': incar_dict,
                       'resources': machine_dict,
                       'pseudo': pseudo_dict,
                       'kpoints': kpoints_dict}
}

#Submit workflow
from aiida.workflows.wf_phonon import WorkflowPhonon
wf = WorkflowPhonon(params=wf_parameters, optimize=True)

wf.label = 'VASP {}'.format(structure.get_formula())
wf.start()
print ('pk: {}'.format(wf.pk))

<commit_msg>add qha plot tool (function version)<commit_after>from aiida import load_dbenv
load_dbenv()
from aiida.orm import Code, DataFactory

StructureData = DataFactory('structure')
ParameterData = DataFactory('parameter')

import numpy as np
import os
import pymatgen

##########################
structure_id = 'mp-12668'
##########################

def get_supercell_size(structure, max_atoms=100):

    cell = np.array(structure.cell)

    num_atoms = len(structure.sites)

    print cell
    print num_atoms
    print np.linalg.norm(cell, axis=1)

    supercell_size = [1, 1, 1]

    while True:

        test_cell = np.dot(cell.T, np.diag(supercell_size)).T
        norm = np.linalg.norm(test_cell, axis=1)
        index = np.argmin(norm)
        supercell_size_test = list(supercell_size)
        supercell_size_test[index] += 1

        anum_atoms_supercell = num_atoms * np.prod(supercell_size_test)
        if anum_atoms_supercell > max_atoms:
            atoms_minus =  num_atoms * np.prod(supercell_size)
            atoms_plus = num_atoms * np.prod(supercell_size_test)

            if max_atoms - atoms_minus < atoms_plus - max_atoms:
                return supercell_size
            else:
                return supercell_size_test
        else:
            supercell_size = supercell_size_test

        print supercell_size
        print test_cell

def get_potential_labels(functional, symbol_list, ftype=None):
    potential_labels =[]
    for symbol in np.unique(symbol_list):
        psp_dir = os.environ['VASP_PSP_DIR'] + '/POT_GGA_PAW_' + functional
        all_labels = [f[7:-3] for f in os.listdir(psp_dir)  if os.path.isfile(os.path.join(psp_dir, f))]
        candidates = [ s for s in all_labels if symbol == s.split('_')[0]]
        if ftype is not None:
            final = [s for s in candidates if '_{}'.format(ftype) == s[-(len(ftype) + 1):]]
            if len(final) > 0:
                potential_labels.append(final[0])
                continue
        potential_labels.append(candidates[0])
    return potential_labels

rester = pymatgen.MPRester(os.environ['PMG_MAPI_KEY'])

pmg_structure = rester.get_structure_by_material_id(structure_id)
spa = pymatgen.symmetry.analyzer.SpacegroupAnalyzer(pmg_structure)

conventional = spa.get_conventional_standard_structure()
primitive = spa.get_primitive_standard_structure()

primitive_matrix = np.dot(np.linalg.inv(conventional.lattice_vectors()), primitive.lattice_vectors())
primitive_matrix = np.round(primitive_matrix, decimals=6).tolist()

structure = StructureData(pymatgen=conventional).store()
print structure
# crystal_system = spa.get_crystal_system()
# if crystal_system == 'hexagonal':
#     supercell = [[3, 0, 0],
#                  [0, 3, 0],
#                  [0, 0, 3]]
# else:
#     supercell = [[2, 0, 0],
#                  [0, 2, 0],
#                  [0, 0, 2]]

supercell_size = get_supercell_size(structure)
supercell = np.diag(supercell_size)

print ('Supercell shape: {}'.format(supercell_size))
exit()

incar_dict = {
    'NELMIN' : 5,
    'NELM'   : 100,
    'ENCUT'  : 400,
    'ALGO'   : 38,
    'ISMEAR' : 0,
    'SIGMA'  : 0.01,
    'GGA'    : 'PS'
}

pseudo_dict = {'functional': 'PBE',
               'symbols': get_potential_labels('PBE', conventional.symbol_set)}

# Monkhorst-pack
kpoints_dict = {'points': [2, 2, 2],
                'shift' : [0.0, 0.0, 0.0]}

machine_dict = {
    'num_machines': 1,
    'parallel_env':'mpi*',
    'tot_num_mpiprocs': 16}

phonopy_parameters = {'supercell': supercell,
                      'primitive': primitive_matrix,
                      'distance': 0.01,
                      'mesh': [20, 20, 20]}

wf_parameters = {
     'structure': structure,
     'phonopy_input': {'parameters': phonopy_parameters},
     'input_force': {'code': 'vasp541mpi@stern',
                    'parameters': incar_dict,
                    'resources': machine_dict,
                    'pseudo': pseudo_dict,
                    'kpoints': kpoints_dict},
     'input_optimize': {'code': 'vasp541mpi@stern',
                       'parameters': incar_dict,
                       'resources': machine_dict,
                       'pseudo': pseudo_dict,
                       'kpoints': kpoints_dict}
}

#Submit workflow
from aiida.workflows.wf_phonon import WorkflowPhonon
wf = WorkflowPhonon(params=wf_parameters, optimize=True)

wf.label = 'VASP {}'.format(structure.get_formula())
wf.start()
print ('pk: {}'.format(wf.pk))

<|endoftext|>


<commit_before>import hashlib
import os
import socket
import sys
import base64
import time
from django.db.models import F, Q
from xos.config import Config
from observer.syncstep import SyncStep
from observer.ansible import run_template_ssh
from core.models import Service, Slice
from cord.models import VCPEService, VCPETenant, VOLTTenant
from hpc.models import HpcService, CDNPrefix
from util.logger import Logger, logging

# hpclibrary will be in steps/..
parentdir = os.path.join(os.path.dirname(__file__),"..")
sys.path.insert(0,parentdir)

from broadbandshield import BBS

logger = Logger(level=logging.INFO)

class SyncVCPETenant(SyncStep):
    provides=[VCPETenant]
    observes=VCPETenant
    requested_interval=0
    template_name = "sync_vcpetenant.yaml"
    service_key_name = "/opt/xos/observers/vcpe/vcpe_private_key"

    def __init__(self, **args):
        SyncStep.__init__(self, **args)

    def defer_sync(self, o, reason):
        logger.info("defer object %s due to %s" % (str(o), reason))
        raise Exception("defer object %s due to %s" % (str(o), reason))

    def fetch_pending(self, deleted):
        if (not deleted):
            objs = VCPETenant.get_tenant_objects().filter(Q(enacted__lt=F('updated')) | Q(enacted=None),Q(lazy_blocked=False))
        else:
            objs = VCPETenant.get_deleted_tenant_objects()

        return objs

    def get_extra_attributes(self, o):
        # This is a place to include extra attributes that aren't part of the
        # object itself. In our case, it's handy to know the VLAN IDs when
        # configuring the VCPE.

        dnsdemux_ip = "none"
        for service in HpcService.objects.all():
            for slice in service.slices.all():
                if "dnsdemux" in slice.name:
                    for sliver in slice.slivers.all():
                        # Connect to a dnsdemux that's on the hpc_client network
                        # if one is available.
                        for ns in sliver.networkslivers.all():
                            if ns.ip and ns.network.labels and ("hpc_client" in ns.network.labels):
                                dnsdemux_ip = ns.ip
                        if dnsdemux_ip=="none":
                            try:
                                dnsdemux_ip = socket.gethostbyname(sliver.node.name)
                            except:
                                pass

        cdn_prefixes = []
        for prefix in CDNPrefix.objects.all():
            cdn_prefixes.append(prefix.prefix)

        vlan_ids = []
        if o.volt:
            vlan_ids.append(o.volt.vlan_id)

        bbs_addrs = []
        bbs_slices = Slice.objects.filter(name="mysite_bbs")
        if bbs_slices:
            bbs_slice = bbs_slices[0]
            for bbs_sliver in bbs_slice.slivers.all():
                for ns in bbs_sliver.networkslivers.all():
                    if ns.ip and ns.network.labels and ("hpc_client" in ns.network.labels):
                        bbs_addrs.append(ns.ip)

        if not bbs_addrs:
            bbs_addrs = ["198.105.255.10",
                         "198.105.255.11",
                         "198.105.255.12",
                         "198.105.255.13"]

        return {"vlan_ids": vlan_ids,
                "dnsdemux_ip": dnsdemux_ip,
                "cdn_prefixes": cdn_prefixes,
                "bbs_addrs": bbs_addrs}

    def get_sliver(self, o):
        # We need to know what slivers is associated with the object.
        # For vCPE this is easy, as VCPETenant has a sliver field.

        return o.sliver

    def sync_record(self, o):
        logger.info("sync'ing VCPETenant %s" % str(o))

        sliver = self.get_sliver(o)
        if not sliver:
            self.defer_sync(o, "waiting on sliver")
            return

        service = o.sliver.slice.service
        if not service:
            # Ansible uses the service's keypair in order to SSH into the
            # instance. It would be bad if the slice had no service.

            raise Exception("Slice %s is not associated with a service" % sliver.slice.name)

        if not os.path.exists(self.service_key_name):
            raise Exception("Service key %s does not exist" % self.service_key_name)

        service_key = file(self.service_key_name).read()

        fields = { "sliver_name": sliver.name,
                   "hostname": sliver.node.name,
                   "instance_id": sliver.instance_id,
                   "private_key": service_key,
                   "ansible_tag": "vcpe_tenant_" + str(o.id)
                 }

        if hasattr(o, "sync_attributes"):
            for attribute_name in o.sync_attributes:
                 fields[attribute_name] = getattr(o, attribute_name)

        fields.update(self.get_extra_attributes(o))

        ansible_hash = hashlib.md5(repr(sorted(fields.items()))).hexdigest()
        quick_update = (o.last_ansible_hash == ansible_hash)

        if quick_update:
            logger.info("quick_update triggered; skipping ansible recipe")
        else:
            tStart = time.time()
            run_template_ssh(self.template_name, fields)
            logger.info("playbook execution time %d" % int(time.time()-tStart))

        if o.url_filter_enable:
            if False: # (str(o.service_specific_id) != "SYNCME"):
                # XXX FIXME
                # Also fix the spot in cord/models.py
                logger.info("skipping sync of URL filter for SSID %s" % str(o.service_specific_id))
            else:
                tStart = time.time()
                bbs = BBS(o.bbs_account, "123")
                bbs.sync(o.url_filter_level, o.users)

                if o.hpc_client_ip:
                    logger.info("associate account %s with ip %s" % (o.bbs_account, o.hpc_client_ip))
                    bbs.associate(o.hpc_client_ip)
                else:
                    logger.info("no hpc_client_ip to associate")

                logger.info("bbs update tiem %d" % int(time.time()-tStart))

        o.last_ansible_hash = ansible_hash
        o.save()

    def delete_record(self, m):
        pass

<commit_msg>eliinate the SYNCME code<commit_after>import hashlib
import os
import socket
import sys
import base64
import time
from django.db.models import F, Q
from xos.config import Config
from observer.syncstep import SyncStep
from observer.ansible import run_template_ssh
from core.models import Service, Slice
from cord.models import VCPEService, VCPETenant, VOLTTenant
from hpc.models import HpcService, CDNPrefix
from util.logger import Logger, logging

# hpclibrary will be in steps/..
parentdir = os.path.join(os.path.dirname(__file__),"..")
sys.path.insert(0,parentdir)

from broadbandshield import BBS

logger = Logger(level=logging.INFO)

class SyncVCPETenant(SyncStep):
    provides=[VCPETenant]
    observes=VCPETenant
    requested_interval=0
    template_name = "sync_vcpetenant.yaml"
    service_key_name = "/opt/xos/observers/vcpe/vcpe_private_key"

    def __init__(self, **args):
        SyncStep.__init__(self, **args)

    def defer_sync(self, o, reason):
        logger.info("defer object %s due to %s" % (str(o), reason))
        raise Exception("defer object %s due to %s" % (str(o), reason))

    def fetch_pending(self, deleted):
        if (not deleted):
            objs = VCPETenant.get_tenant_objects().filter(Q(enacted__lt=F('updated')) | Q(enacted=None),Q(lazy_blocked=False))
        else:
            objs = VCPETenant.get_deleted_tenant_objects()

        return objs

    def get_extra_attributes(self, o):
        # This is a place to include extra attributes that aren't part of the
        # object itself. In our case, it's handy to know the VLAN IDs when
        # configuring the VCPE.

        dnsdemux_ip = "none"
        for service in HpcService.objects.all():
            for slice in service.slices.all():
                if "dnsdemux" in slice.name:
                    for sliver in slice.slivers.all():
                        # Connect to a dnsdemux that's on the hpc_client network
                        # if one is available.
                        for ns in sliver.networkslivers.all():
                            if ns.ip and ns.network.labels and ("hpc_client" in ns.network.labels):
                                dnsdemux_ip = ns.ip
                        if dnsdemux_ip=="none":
                            try:
                                dnsdemux_ip = socket.gethostbyname(sliver.node.name)
                            except:
                                pass

        cdn_prefixes = []
        for prefix in CDNPrefix.objects.all():
            cdn_prefixes.append(prefix.prefix)

        vlan_ids = []
        if o.volt:
            vlan_ids.append(o.volt.vlan_id)

        bbs_addrs = []
        bbs_slices = Slice.objects.filter(name="mysite_bbs")
        if bbs_slices:
            bbs_slice = bbs_slices[0]
            for bbs_sliver in bbs_slice.slivers.all():
                for ns in bbs_sliver.networkslivers.all():
                    if ns.ip and ns.network.labels and ("hpc_client" in ns.network.labels):
                        bbs_addrs.append(ns.ip)

        if not bbs_addrs:
            bbs_addrs = ["198.105.255.10",
                         "198.105.255.11",
                         "198.105.255.12",
                         "198.105.255.13"]

        return {"vlan_ids": vlan_ids,
                "dnsdemux_ip": dnsdemux_ip,
                "cdn_prefixes": cdn_prefixes,
                "bbs_addrs": bbs_addrs}

    def get_sliver(self, o):
        # We need to know what slivers is associated with the object.
        # For vCPE this is easy, as VCPETenant has a sliver field.

        return o.sliver

    def sync_record(self, o):
        logger.info("sync'ing VCPETenant %s" % str(o))

        sliver = self.get_sliver(o)
        if not sliver:
            self.defer_sync(o, "waiting on sliver")
            return

        service = o.sliver.slice.service
        if not service:
            # Ansible uses the service's keypair in order to SSH into the
            # instance. It would be bad if the slice had no service.

            raise Exception("Slice %s is not associated with a service" % sliver.slice.name)

        if not os.path.exists(self.service_key_name):
            raise Exception("Service key %s does not exist" % self.service_key_name)

        service_key = file(self.service_key_name).read()

        fields = { "sliver_name": sliver.name,
                   "hostname": sliver.node.name,
                   "instance_id": sliver.instance_id,
                   "private_key": service_key,
                   "ansible_tag": "vcpe_tenant_" + str(o.id)
                 }

        if hasattr(o, "sync_attributes"):
            for attribute_name in o.sync_attributes:
                 fields[attribute_name] = getattr(o, attribute_name)

        fields.update(self.get_extra_attributes(o))

        ansible_hash = hashlib.md5(repr(sorted(fields.items()))).hexdigest()
        quick_update = (o.last_ansible_hash == ansible_hash)

        if quick_update:
            logger.info("quick_update triggered; skipping ansible recipe")
        else:
            tStart = time.time()
            run_template_ssh(self.template_name, fields)
            logger.info("playbook execution time %d" % int(time.time()-tStart))

        if o.url_filter_enable:
            tStart = time.time()
            bbs = BBS(o.bbs_account, "123")
            bbs.sync(o.url_filter_level, o.users)

            if o.hpc_client_ip:
                logger.info("associate account %s with ip %s" % (o.bbs_account, o.hpc_client_ip))
                bbs.associate(o.hpc_client_ip)
            else:
                logger.info("no hpc_client_ip to associate")

            logger.info("bbs update tiem %d" % int(time.time()-tStart))

        o.last_ansible_hash = ansible_hash
        o.save()

    def delete_record(self, m):
        pass

<|endoftext|>


<commit_before>#!/usr/bin/env python
"""This script downloads NCBI records based on user search query.

   Requires the biopython package.
"""

__author__ = 'DV Klopfenstein'
__copyright__ = "Copyright (C) 2014-2015 DV Klopfenstein. All rights reserved."
__license__ = "GPL"

import sys
import os
import shutil
import re

from Bio import Entrez
import requests

# -------------------------------------------------------
def find_IDs_with_ESearch(db, retmax, email, query, log=sys.stdout):
  """Searches an NCBI database for a user search term, returns NCBI records."""

  # Provide your email to the NCBI so they may contact you prior
  # to shutting down server access in the case of inadvertant over-use.
  Entrez.email = email

  # ESearch for NCBI IDs related to User Search...
  socket_handle = Entrez.esearch(
    db=db,
    retmax=retmax,
    term=query,
    usehistory="y") # NCBI prefers we use history(QueryKey, WebEnv) for next acess
  record = Entrez.read(socket_handle)
  socket_handle.close()

  if 'IdList' in record and record['IdList']:
    return record

# -------------------------------------------------------
def EPost(db, IDs, email, log=sys.stdout, step=10):
  """Posts to NCBI WebServer of any number of UIDs."""
  ## http://biopython.org/DIST/docs/tutorial/Tutorial.html#htoc112
  Entrez.email = email
  # Load the first 1...(step-1) UIDs to Entrez using epost. Get WebEnv to finish post."""
  strIDs = map(str, IDs)
  id_str = ','.join(strIDs[:step])
  # epost produces WebEnv value ($web1) and QueryKey value ($key1)
  socket_handle = Entrez.epost(db, id=id_str)
  record = Entrez.read(socket_handle)
  if log is not None:
    log.write('FIRST EPOST RESULT: {}\n'.format(record))
    log.write("QueryKey({:>6})  IDs={}\n".format(record['QueryKey'], id_str))
  socket_handle.close()
  if 'WebEnv' in record:
    WebEnv = record['WebEnv']
    num_IDs = len(strIDs)
    # Load the remainder of the UIDs using epost
    for idx in range(step, num_IDs, step):
      end_pt = idx+step
      if num_IDs < end_pt:
        end_pt = num_IDs
      #print '{:3} {:3} {:3}'.format(num_IDs, idx, end_pt)
      id_str = ','.join(strIDs[idx:end_pt])
      socket_handle = Entrez.epost(db, id=id_str, WebEnv=WebEnv)
      record = Entrez.read(socket_handle)
      WebEnv = record['WebEnv']
      if log is not None: log.write("QueryKey({:>6})  strIDs={}\n".format(record['QueryKey'], id_str))
      socket_handle.close()
  else:
    raise Exception("NO WebEnv RETURNED FROM FIRST EPOST")
  if log is not None: log.write('LAST  EPOST RESULT: {}\n'.format(record))
  return record

# -------------------------------------------------------
def ESummary_from_Post(db, QueryKey, WebEnv, email):
  """Give a post reoord. Get an esummary ."""
  Entrez.email = email
  socket_handle = Entrez.esummary(db=db, query_key=QueryKey, WebEnv=WebEnv)
  record = Entrez.read(socket_handle)
  socket_handle.close()
  return record

# -------------------------------------------------------
# einfo: http://www.ncbi.nlm.nih.gov/books/NBK25499/
def get_DbList():
  """Gets a list of valid Entrez databases."""
  socket_handle = Entrez.einfo()
  record = Entrez.read(socket_handle)
  socket_handle.close()
  return record['DbList']

def prt_DbList(PRT=sys.stdout):
  dblist = get_DbList()
  for db in dblist:
    PRT.write('{}\n'.format(db))

def get_DbStats(db, version='2.0', retmode='xml'):
  """Return stats about an Entrez database.

     db:      Use get_DbList to get valid entries
     retmode: 'xml' or 'json'
  """
  socket_handle = Entrez.einfo(
    db=db,
    version=version,
    retmode=retmode)
  record = Entrez.read(socket_handle)
  socket_handle.close()
  return record

# -------------------------------------------------------
def EFetch_and_write(desc, db, fout, typemode, record, batch_size=100, PRT=sys.stdout):
  """Fetches NCBI records returned from last search.

  For NCBI's online documentation of efetch:
    http://www.ncbi.nlm.nih.gov/books/NBK25499/#chapter4.EFetch
  """
  N  = len(record['IdList'])
  # Write the list of UIDs in the record
  tsv = get_tsv_filename(fout)
  wr_IDs(tsv, record)
  EFetch_and_write_N(desc, db, fout, typemode, record, N, batch_size, PRT)

def EFetch_and_write_N(desc, db, fout, typemode, record, N, batch_size=100, PRT=sys.stdout):
  FOUT = open(fout, 'w')
  downloaded_data = EFetch_and_write_WEQK_N(desc, db, FOUT, typemode, record,
    record['WebEnv'], record['QueryKey'], N, batch_size=100, PRT=sys.stdout)
  FOUT.close(); PRT.write("  WROTE: {}\n".format(fout))

  sys.stdout.write("""
    Need to revisit reading XML reading another way.
    Biopython XML no longer working for XML files from NCBI Gene.\n""")
  if N > batch_size and typemode[1] == "xml": # DVK Biopython XML not working
    Entrez_strip_extra_eSummaryResult(fout)   # DVK Biopython XML not working
  return downloaded_data

def EFetch_and_write_WEQK_N(desc, db, FOUT, typemode, WE, QK, N, batch_size=100, PRT=sys.stdout):
  downloaded_data = None
  # EFetch records found for IDs returned in ESearch...
  # Search for IDs returned using ID of the above Web Search
  #print "AAAA", N, batch_size
  for start in range(0, N, batch_size):
    #print "BBBB", start
    socket_handle = None
    msg = '  QueryKey({:>6}) EFetching(db={}) up to {:5} records, starting at {}; {}\n'.format(
      QK, db, batch_size, start, desc)
    #PRT.write(msg); PRT.flush()
    #sys.stdout.write(msg)
    try:
      socket_handle = Entrez.efetch(
        db        = db,
        retstart  = start,       # dflt: 1
        retmax    = batch_size,  # max: 10,000
        rettype   = typemode[0], # Ex: medline
        retmode   = typemode[1], # Ex: text
        webenv    = WE,
        query_key = QK)
    except IOError, e:
      msg = "\n*FATAL: EFetching FAILED: {}".format(e)
      PRT.write(msg)
      sys.stdout.write(msg)
      sys.stdout.write("  db: {}\n".format(db))
      sys.stdout.write("  retstart: {}\n".format(start))
      sys.stdout.write("  retmax: {}\n".format(batch_size))
      sys.stdout.write("  rettype: {}\n".format(typemode[0]))
      sys.stdout.write("  retmode: {}\n".format(typemode[1]))
      sys.stdout.write("  WebEnv: {}\n".format(WE))
      sys.stdout.write("  QueryKey: {}\n".format(QK))
      if socket_handle is not None:
        socket_handle.close()
        socket_handle = None

    if socket_handle is not None:
      try:
        # Read the downloaded data from the socket handle
        downloaded_data = socket_handle.read()
        socket_handle.close()
        M = re.search(r'(ERROR.*\n)', downloaded_data)
        if M:
          sys.stdout.write(M.group(1))
        FOUT.write(downloaded_data)
        FOUT.flush()
      except Exception, e:
        sys.stdout.write("*FATAL: PROBLEM READING FROM SOCKET HANDLE: {}\n".format(str(e)))
    else:
      sys.stdout.write("*FATAL: NO SOCKET HANDLE TO READ FROM\n")
  return downloaded_data

def ELink_webenvQK_num(desc, webenv, querykey, linkname, num, batch_size=100, PRT=sys.stdout):
  """https://www.ncbi.nlm.nih.gov/books/NBK25499/#chapter4.ELink"""
  ids = []
  dbfrom = 'pubmed'
  # EFetch records found for IDs returned in ESearch...
  # Search for IDs returned using ID of the above Web Search
  #print "AAAA", N, batch_size
  for start in range(0, num, batch_size):
    #print "BBBB", start
    socket_handle = None
    msg = '  QueryKey({:>6}) ELinking(db={}) up to {:5} records, starting at {}; {}\n'.format(
      querykey, linkname, batch_size, start, desc)
    #PRT.write(msg); PRT.flush()
    #sys.stdout.write(msg)
    try:
      socket_handle = Entrez.elink(
        dbfrom    = dbfrom,
        retstart  = start,       # dflt: 1
        retmax    = batch_size,  # max: 10,000
        linkname  = linkname,
        webenv    = webenv,
        query_key = querykey)
    except IOError, e:
      msg = "\n*FATAL: EFetching FAILED: {}".format(e)
      PRT.write(msg)
      sys.stdout.write(msg)
      sys.stdout.write("  dbfrom: {}\n".format(dbfrom))
      sys.stdout.write("  retstart: {}\n".format(start))
      sys.stdout.write("  retmax: {}\n".format(batch_size))
      sys.stdout.write("  linkname: {}\n".format(linkname))
      sys.stdout.write("  WebEnv: {}\n".format(webenv))
      sys.stdout.write("  QueryKey: {}\n".format(querykey))
      if socket_handle is not None:
        socket_handle.close()
        socket_handle = None

    if socket_handle is not None:
      try:
        # Read the downloaded data from the socket handle
        # [{u'LinkSetDb': [{u'DbTo': 'pubmed', u'Link': [{u'Id': '5668'}, {u'Id': '4442'}, ...
        #print "ddddd", socket_handle
        record = Entrez.read(socket_handle)
        #print "DDDDD", record
        socket_handle.close()
        if record:
          if not record[0]["ERROR"]:
            # keys: DbTo Link LinkName.    Link list:[{'Id':'NNNNN'}, {'Id':'NNNNN'}, ....
            # keys: LinkSetDb DbFrom IdList LinkSetDbHistory ERROR
            if record[0][u'LinkSetDb']:
              ids.extend([int(id2val[u'Id']) for id2val in record[0][u'LinkSetDb'][0][u'Link']])
            #print "EEEEE", " ".join(record[0].keys())
            #sys.stdout.write(record[0].keys())
      except Exception, e:
        sys.stdout.write("*FATAL: PROBLEM READING FROM SOCKET HANDLE: {}\n".format(str(e)))
    else:
      sys.stdout.write("*FATAL: NO SOCKET HANDLE TO READ FROM\n")
  return ids

# -------------------------------------------------------
def Entrez_strip_extra_eSummaryResult(Entrez_datafile):
  """For an XML file EFetched in multiple batches, strip extra eSummaryResult lines.

  The first </eSummaryResult> will cause Entrez.parse to prematurely stop parsing.

  """
  trash = '_'.join([Entrez_datafile, 'trash'])
  # Moved the poorly formatted XML file to <filename>_trash
  shutil.move(Entrez_datafile, trash)
  # Write newly formatted XML file into original filename
  FOUT = open(Entrez_datafile, 'w')
  xml_version  = None
  DOCTYPE      = None
  eResult      = None
  S = re.compile(r'DOCTYPE\s+(\S+)\s+')
  """Only write the xml_version|DOCTYPE|eSummaryResult once."""
  with open(trash) as FIN:
    for line in FIN:
      if   'xml version' in line:
        if xml_version is None:
          FOUT.write(line) # Write line once, the first time it is seen
          xml_version = True
      elif 'DOCTYPE' in line:
        if DOCTYPE is None:
          M = S.search(line)
          if M:
            FOUT.write(line) # Write line once, the first time it is seen
            DOCTYPE = M.group(1)
          else:
            print "*FATAL: ERROR STRIPPING XML({})".format(Entrez_datafile)
      elif DOCTYPE is not None and re.search(r'^</?{}>\s*$'.format(DOCTYPE), line):
        if eResult is None:
          FOUT.write(line) # Write line once, the first time it is seen
          eResult = True
      else:
        FOUT.write(line)
  FOUT.write('</{}>\n'.format(DOCTYPE))
  FOUT.close()
  FIN.close()

# -------------------------------------------------------
def wr_IDs(tsv, record):
  if 'IdList' in record:
    TSV = open(tsv, 'w')
    chk_num_IDs(sys.stdout, record)
    for ID in record['IdList']:
      TSV.write('{}\n'.format(ID))
    TSV.close();
    sys.stdout.write("  WROTE: {}  ESearch Returned # {} IDs\n".format(
      tsv, len(record['IdList'])))

# -------------------------------------------------------
def chk_num_IDs(PRT, record):
  """ALerts the User to increase retmax if not all records are returned."""
  if 'Count' in record and 'RetMax' in record and record['Count'] != record['RetMax']:
    txt = '**Note: {} PubMed IDs returned out of {} total found.\n'.format(
      record['RetMax'],
      record['Count'])
    PRT.write(txt)
    sys.stdout.write(txt)

# -------------------------------------------------------
def get_tsv_filename(fout):
  """NCBI IDs are stored in a ".tsv" file.  Runtime notes written into ".log" file"""
  basename = os.path.basename(os.path.splitext(fout)[0])
  tsv = ''.join([basename, '_ESearch_IDs.tsv'])
  return tsv

# -------------------------------------------------------
def get_fetch_fout(filename):
  """Create a output file name where EFetch data will be stored."""
  return ''.join([os.path.basename(os.path.splitext(filename)[0]), '.Entrez_downloads'])

# -------------------------------------------------------
def get_args():
  email = "music_pupil@yahoo.com"
  query = "asthma[mesh] AND leukotrienes[mesh] AND 2009[pdat]"
  max_records = 500
  return email, query, max_records

<commit_msg>separated downloading from writing to a markdown file.<commit_after>#!/usr/bin/env python
"""This script downloads NCBI records based on user search query.

   Requires the biopython package.
"""

__author__ = 'DV Klopfenstein'
__copyright__ = "Copyright (C) 2014-2015 DV Klopfenstein. All rights reserved."
__license__ = "GPL"

import sys
import os
import shutil
import re

from Bio import Entrez
import requests

# -------------------------------------------------------
def find_IDs_with_ESearch(db, retmax, email, query, log=sys.stdout):
  """Searches an NCBI database for a user search term, returns NCBI records."""

  # Provide your email to the NCBI so they may contact you prior
  # to shutting down server access in the case of inadvertant over-use.
  Entrez.email = email

  # ESearch for NCBI IDs related to User Search...
  socket_handle = Entrez.esearch(
    db=db,
    retmax=retmax,
    term=query,
    usehistory="y") # NCBI prefers we use history(QueryKey, WebEnv) for next acess
  record = Entrez.read(socket_handle)
  socket_handle.close()

  if 'IdList' in record and record['IdList']:
    return record

# -------------------------------------------------------
def EPost(db, IDs, email, log=sys.stdout, step=10):
  """Posts to NCBI WebServer of any number of UIDs."""
  ## http://biopython.org/DIST/docs/tutorial/Tutorial.html#htoc112
  Entrez.email = email
  # Load the first 1...(step-1) UIDs to Entrez using epost. Get WebEnv to finish post."""
  strIDs = map(str, IDs)
  id_str = ','.join(strIDs[:step])
  # epost produces WebEnv value ($web1) and QueryKey value ($key1)
  socket_handle = Entrez.epost(db, id=id_str)
  record = Entrez.read(socket_handle)
  if log is not None:
    log.write('FIRST EPOST RESULT: {}\n'.format(record))
    log.write("QueryKey({:>6})  IDs={}\n".format(record['QueryKey'], id_str))
  socket_handle.close()
  if 'WebEnv' in record:
    WebEnv = record['WebEnv']
    num_IDs = len(strIDs)
    # Load the remainder of the UIDs using epost
    for idx in range(step, num_IDs, step):
      end_pt = idx+step
      if num_IDs < end_pt:
        end_pt = num_IDs
      #print '{:3} {:3} {:3}'.format(num_IDs, idx, end_pt)
      id_str = ','.join(strIDs[idx:end_pt])
      socket_handle = Entrez.epost(db, id=id_str, WebEnv=WebEnv)
      record = Entrez.read(socket_handle)
      WebEnv = record['WebEnv']
      if log is not None: log.write("QueryKey({:>6})  strIDs={}\n".format(record['QueryKey'], id_str))
      socket_handle.close()
  else:
    raise Exception("NO WebEnv RETURNED FROM FIRST EPOST")
  if log is not None: log.write('LAST  EPOST RESULT: {}\n'.format(record))
  return record

# -------------------------------------------------------
def ESummary_from_Post(db, QueryKey, WebEnv, email):
  """Give a post reoord. Get an esummary ."""
  Entrez.email = email
  socket_handle = Entrez.esummary(db=db, query_key=QueryKey, WebEnv=WebEnv)
  record = Entrez.read(socket_handle)
  socket_handle.close()
  return record

# -------------------------------------------------------
# einfo: http://www.ncbi.nlm.nih.gov/books/NBK25499/
def get_DbList():
  """Gets a list of valid Entrez databases."""
  socket_handle = Entrez.einfo()
  record = Entrez.read(socket_handle)
  socket_handle.close()
  return record['DbList']

def prt_DbList(PRT=sys.stdout):
  dblist = get_DbList()
  for db in dblist:
    PRT.write('{}\n'.format(db))

def get_DbStats(db, version='2.0', retmode='xml'):
  """Return stats about an Entrez database.

     db:      Use get_DbList to get valid entries
     retmode: 'xml' or 'json'
  """
  socket_handle = Entrez.einfo(
    db=db,
    version=version,
    retmode=retmode)
  record = Entrez.read(socket_handle)
  socket_handle.close()
  return record

# -------------------------------------------------------
def EFetch_and_write(desc, db, fout, typemode, record, batch_size=100, PRT=sys.stdout):
  """Fetches NCBI records returned from last search.

  For NCBI's online documentation of efetch:
    http://www.ncbi.nlm.nih.gov/books/NBK25499/#chapter4.EFetch
  """
  N  = len(record['IdList'])
  # Write the list of UIDs in the record
  tsv = get_tsv_filename(fout)
  wr_IDs(tsv, record)
  EFetch_and_write_N(desc, db, fout, typemode, record, N, batch_size)

def download_records(desc, db, fout_dl, typemode, record, bnum, batch_size=100):
  """Download records and write to a text file."""
  with open(fout_dl, 'wb') as ostrm:
      downloaded_data = EFetch_and_write_WEQK_N(desc, db, ostrm, typemode,
        record['WebEnv'], record['QueryKey'], bnum, batch_size)
      sys.stdout.write("  WROTE: {}\n".format(fout_dl))
      return downloaded_data

def EFetch_and_write_N(desc, db, fout_dl, typemode, record, bnum, batch_size=100):
  downloaded_data = download_records(desc, db, fout_dl, typemode, record, bnum, batch_size)
  sys.stdout.write("""
    Need to revisit reading XML reading another way.
    Biopython XML no longer working for XML files from NCBI Gene.\n""")
  if bnum > batch_size and typemode[1] == "xml": # DVK Biopython XML not working
    Entrez_strip_extra_eSummaryResult(fout)   # DVK Biopython XML not working
  return downloaded_data

def EFetch_and_write_WEQK_N(desc, db, ostrm, typemode, WE, QK, N, batch_size=100):
  downloaded_data = None
  # EFetch records found for IDs returned in ESearch...
  # Search for IDs returned using ID of the above Web Search
  #print "AAAA", N, batch_size
  for start in range(0, N, batch_size):
    #print "BBBB", start
    socket_handle = None
    msg = '  QueryKey({:>6}) EFetching(db={}) up to {:5} records, starting at {}; {}\n'.format(
      QK, db, batch_size, start, desc)
    #sys.stdout.write(msg)
    try:
      socket_handle = Entrez.efetch(
        db        = db,
        retstart  = start,       # dflt: 1
        retmax    = batch_size,  # max: 10,000
        rettype   = typemode[0], # Ex: medline
        retmode   = typemode[1], # Ex: text
        webenv    = WE,
        query_key = QK)
    except IOError, e:
      msg = "\n*FATAL: EFetching FAILED: {}".format(e)
      sys.stdout.write(msg)
      sys.stdout.write("  db: {}\n".format(db))
      sys.stdout.write("  retstart: {}\n".format(start))
      sys.stdout.write("  retmax: {}\n".format(batch_size))
      sys.stdout.write("  rettype: {}\n".format(typemode[0]))
      sys.stdout.write("  retmode: {}\n".format(typemode[1]))
      sys.stdout.write("  WebEnv: {}\n".format(WE))
      sys.stdout.write("  QueryKey: {}\n".format(QK))
      if socket_handle is not None:
        socket_handle.close()
        socket_handle = None

    if socket_handle is not None:
      try:
        # Read the downloaded data from the socket handle
        downloaded_data = socket_handle.read()
        socket_handle.close()
        M = re.search(r'(ERROR.*\n)', downloaded_data)
        if M:
          sys.stdout.write(M.group(1))
        ostrm.write(downloaded_data)
        ostrm.flush()
      except Exception, e:
        sys.stdout.write("*FATAL: PROBLEM READING FROM SOCKET HANDLE: {}\n".format(str(e)))
    else:
      sys.stdout.write("*FATAL: NO SOCKET HANDLE TO READ FROM\n")
  return downloaded_data

def ELink_webenvQK_num(desc, webenv, querykey, linkname, num, batch_size=100, PRT=sys.stdout):
  """https://www.ncbi.nlm.nih.gov/books/NBK25499/#chapter4.ELink"""
  ids = []
  dbfrom = 'pubmed'
  # EFetch records found for IDs returned in ESearch...
  # Search for IDs returned using ID of the above Web Search
  #print "AAAA", N, batch_size
  for start in range(0, num, batch_size):
    #print "BBBB", start
    socket_handle = None
    msg = '  QueryKey({:>6}) ELinking(db={}) up to {:5} records, starting at {}; {}\n'.format(
      querykey, linkname, batch_size, start, desc)
    #PRT.write(msg); PRT.flush()
    #sys.stdout.write(msg)
    try:
      socket_handle = Entrez.elink(
        dbfrom    = dbfrom,
        retstart  = start,       # dflt: 1
        retmax    = batch_size,  # max: 10,000
        linkname  = linkname,
        webenv    = webenv,
        query_key = querykey)
    except IOError, e:
      msg = "\n*FATAL: EFetching FAILED: {}".format(e)
      PRT.write(msg)
      sys.stdout.write(msg)
      sys.stdout.write("  dbfrom: {}\n".format(dbfrom))
      sys.stdout.write("  retstart: {}\n".format(start))
      sys.stdout.write("  retmax: {}\n".format(batch_size))
      sys.stdout.write("  linkname: {}\n".format(linkname))
      sys.stdout.write("  WebEnv: {}\n".format(webenv))
      sys.stdout.write("  QueryKey: {}\n".format(querykey))
      if socket_handle is not None:
        socket_handle.close()
        socket_handle = None

    if socket_handle is not None:
      try:
        # Read the downloaded data from the socket handle
        # [{u'LinkSetDb': [{u'DbTo': 'pubmed', u'Link': [{u'Id': '5668'}, {u'Id': '4442'}, ...
        #print "ddddd", socket_handle
        record = Entrez.read(socket_handle)
        #print "DDDDD", record
        socket_handle.close()
        if record:
          if not record[0]["ERROR"]:
            # keys: DbTo Link LinkName.    Link list:[{'Id':'NNNNN'}, {'Id':'NNNNN'}, ....
            # keys: LinkSetDb DbFrom IdList LinkSetDbHistory ERROR
            if record[0][u'LinkSetDb']:
              ids.extend([int(id2val[u'Id']) for id2val in record[0][u'LinkSetDb'][0][u'Link']])
            #print "EEEEE", " ".join(record[0].keys())
            #sys.stdout.write(record[0].keys())
      except Exception, e:
        sys.stdout.write("*FATAL: PROBLEM READING FROM SOCKET HANDLE: {}\n".format(str(e)))
    else:
      sys.stdout.write("*FATAL: NO SOCKET HANDLE TO READ FROM\n")
  return ids

# -------------------------------------------------------
def Entrez_strip_extra_eSummaryResult(Entrez_datafile):
  """For an XML file EFetched in multiple batches, strip extra eSummaryResult lines.

  The first </eSummaryResult> will cause Entrez.parse to prematurely stop parsing.

  """
  trash = '_'.join([Entrez_datafile, 'trash'])
  # Moved the poorly formatted XML file to <filename>_trash
  shutil.move(Entrez_datafile, trash)
  # Write newly formatted XML file into original filename
  FOUT = open(Entrez_datafile, 'w')
  xml_version  = None
  DOCTYPE      = None
  eResult      = None
  S = re.compile(r'DOCTYPE\s+(\S+)\s+')
  """Only write the xml_version|DOCTYPE|eSummaryResult once."""
  with open(trash) as FIN:
    for line in FIN:
      if   'xml version' in line:
        if xml_version is None:
          FOUT.write(line) # Write line once, the first time it is seen
          xml_version = True
      elif 'DOCTYPE' in line:
        if DOCTYPE is None:
          M = S.search(line)
          if M:
            FOUT.write(line) # Write line once, the first time it is seen
            DOCTYPE = M.group(1)
          else:
            print "*FATAL: ERROR STRIPPING XML({})".format(Entrez_datafile)
      elif DOCTYPE is not None and re.search(r'^</?{}>\s*$'.format(DOCTYPE), line):
        if eResult is None:
          FOUT.write(line) # Write line once, the first time it is seen
          eResult = True
      else:
        FOUT.write(line)
  FOUT.write('</{}>\n'.format(DOCTYPE))
  FOUT.close()
  FIN.close()

# -------------------------------------------------------
def wr_IDs(tsv, record):
  if 'IdList' in record:
    TSV = open(tsv, 'w')
    chk_num_IDs(sys.stdout, record)
    for ID in record['IdList']:
      TSV.write('{}\n'.format(ID))
    TSV.close();
    sys.stdout.write("  WROTE: {}  ESearch Returned # {} IDs\n".format(
      tsv, len(record['IdList'])))

# -------------------------------------------------------
def chk_num_IDs(PRT, record):
  """ALerts the User to increase retmax if not all records are returned."""
  if 'Count' in record and 'RetMax' in record and record['Count'] != record['RetMax']:
    txt = '**Note: {} PubMed IDs returned out of {} total found.\n'.format(
      record['RetMax'],
      record['Count'])
    PRT.write(txt)
    sys.stdout.write(txt)

# -------------------------------------------------------
def get_tsv_filename(fout):
  """NCBI IDs are stored in a ".tsv" file.  Runtime notes written into ".log" file"""
  basename = os.path.basename(os.path.splitext(fout)[0])
  tsv = ''.join([basename, '_ESearch_IDs.tsv'])
  return tsv

# -------------------------------------------------------
def get_fetch_fout(filename):
  """Create a output file name where EFetch data will be stored."""
  return ''.join([os.path.basename(os.path.splitext(filename)[0]), '.Entrez_downloads'])

# -------------------------------------------------------
def get_args():
  email = "music_pupil@yahoo.com"
  query = "asthma[mesh] AND leukotrienes[mesh] AND 2009[pdat]"
  max_records = 500
  return email, query, max_records

<|endoftext|>