<commit_before>'''
This file contains the functions to read data from mongodb
'''

from pymongo import MongoClient
import datetime

def get_db_client():
    ''' get the db client or None if error '''
    client = MongoClient('db', 27017)
    try: # this is because of change in pymongo 3.0, we have to do this
        client.admin.command('ismaster') # The ismaster command is cheap and does not require auth.
        return client
    except errors.ConnectionFailure:
        print("Server not available")
        return None

def get_stock_data(stock_dict):
    '''Get the data of the stock specified by the input

    This function retreives data of ONLY one stock.

    Args:
        stock_dict: A dict mapping information about the stock 
            to be queried. For example:
            {
                'code': 'XXXXXX',
                'start_date': datatime.datatime(YYYY,MM,DD,tzinfo=timezone.utc)
                'end_date': datatime.datatime(YYYY,MM,DD,tzinfo=timezone.utc)
            }
    Returns:
        A list of dicts each contains the data of this stock of one day.
        And it is sorted descending by date.
        For example:
        [
            {
                'date': datatime.datatime(YYYY,MM,DD,tzinfo=timezone.utc)
                'open_price': xx.xx,
                'close_price': xx.xx
            },
            ...
        ]
    '''
    client = get_db_client()
    if not client:
        return False
    db = client.fin
    filter_dict = {
            'code':stock_dict['code'],
            'date':{
                '$gte':stock_dict['start_date'],
                '$lte':stock_dict['end_date']
                }
            }
    proj_dict = {
            'date': True,
            'open_price': True,
            'close_price': True
            }
    sort_list = [('date':pymongo.DESCENDING)]
    cursor = db.find(filter=filter_dict, projection=proj_dict, sort=sort_list)
    return list(cursor)<commit_msg>daily update<commit_after>'''
This file contains the functions to read data from mongodb
'''

from pymongo import MongoClient
import datetime

def get_db_client():
    ''' get the db client or None if error '''
    client = MongoClient('db', 27017)
    try: # this is because of change in pymongo 3.0, we have to do this
        client.admin.command('ismaster') # The ismaster command is cheap and does not require auth.
        return client
    except errors.ConnectionFailure:
        print("Server not available")
        return None

def get_recent_data(days, code=None):
    '''Get the data of the stock specified by the input

    This function retreives data of ONLY one stock.

    Args:
        days:   The length of the days
        code:   (Optional) the code of the stock. If omited,
                it will retreive data of all stock codes
    Returns:
        A list of dicts each contains the data of this stock of one day.
        And it is sorted descending by date.
        For example:
        [
            {
                'code': XXXX,
                'date': datatime.datatime(YYYY,MM,DD,tzinfo=timezone.utc)
                'open_price': xx.xx,
                'close_price': xx.xx
            },
            ...
        ]
    '''
    client = get_db_client()
    if not client:
        return False
    db = client.fin
    if code != None:
        filter_dict = { 'code': code }
        proj_dict = {
                'code': True,
                'date': True,
                'open_price': True,
                'close_price': True
                }
        sort_list = [('date':pymongo.DESCENDING)]
        cursor = db.stocks.find(filter=filter_dict, projection=proj_dict, sort=sort_list, limit=days)
        return list(cursor)
    else:
        codes = get_code_list()
        return [doc for x in codes for doc in get_recent_data(days,x)]

def get_code_list():
    client = get_db_client()
    if not client:
        return False
    db = client.fin
    pipeline = [ { '$group': {'_id':'code'} } ]
    cursor = db.stocks.aggregate(pipeline)
    return [doc['_id'] for doc in list(cursor)]
{'fraction_non_alphanumeric': 0.09595092024539878, 'fraction_numerical': 0.0034355828220858898, 'mean_word_length': 1.9472161966738972, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5169919', 'n_tokens_mistral': 1216, 'n_tokens_neox': 1161, 'n_words': 372}