Spaces:
Configuration error
Configuration error
| import datetime as dt | |
| import pandas as pd | |
| import itertools | |
| import pymongo # requires dnspython | |
| import configparser | |
| from .mongo_credentials import credentials | |
| CREDENTIALS = credentials | |
| class _MongoAbstract(object): | |
| _instance = None | |
| _credentials = None | |
| _name = None | |
| _key_ls = None | |
| _del_many_limit = 50 | |
| _override_boolean_key = None | |
| def __init__(self, instance='remote', **query): | |
| self._set_instance(instance) | |
| self._set_credentials() | |
| if not isinstance(self._name, str) or not isinstance(self._instance, str) or not isinstance(self._key_ls, list): | |
| raise Exception('Instance (str), collection _name (str) and _key_ls (list) need to be properly declared!') | |
| db_obj = self._get_db_object() | |
| collection_names_ls = db_obj.list_collection_names() | |
| idx_tup_ls = [(i, pymongo.ASCENDING) for i in self._key_ls] | |
| if self._name not in collection_names_ls: | |
| print(f"Creating new {self._instance} collection: {self._name}\n") | |
| collection = db_obj[self._name] # this will create the collection | |
| collection.create_index(idx_tup_ls, unique=True, name=self._name) # this will index it | |
| else: | |
| collection = db_obj[self._name] # collection already exists | |
| self._db_obj = db_obj | |
| self._collection = collection | |
| query = self._override_boolean_keys(**query) | |
| self._is_adv_query = sum([isinstance(v, list) for v in query.values()]) | |
| self._set_query(query) | |
| self._set_advanced_query() | |
| self._out = None | |
| self._cursor = self._get_cursor() | |
| self._full_key_ls = self._key_ls + ['values'] | |
| def _set_credentials(self): | |
| if self._credentials is None: | |
| self._credentials = CREDENTIALS | |
| def _get_credentials(self): | |
| if self._credentials is None: | |
| self._set_credentials() | |
| return self._credentials | |
| def _set_instance(self, instance): | |
| self._instance = instance | |
| def _override_boolean_keys(self, **query): | |
| if self._get_boolean_keys() and query: | |
| for key in self._get_boolean_keys(): | |
| if key in query.keys(): | |
| if isinstance(query[key], (tuple, list)): | |
| query[key] = [str(int(i)) if isinstance(i, bool) else i for i in query[key]] | |
| else: | |
| if isinstance(query[key], bool): | |
| query[key] = str(int(query[key])) | |
| return query | |
| def _get_boolean_keys(self): | |
| if self._override_boolean_key is not None: | |
| if isinstance(self._override_boolean_key, str): | |
| return [self._override_boolean_key] | |
| elif isinstance(self._override_boolean_key, list): | |
| return self._override_boolean_key | |
| else: | |
| raise Exception(f'_override_boolean_key not recognized {self._override_boolean_key}') | |
| else: | |
| return None | |
| def __repr__(self): | |
| print_str = f'database: {self._get_db_name()}' | |
| print_str += f'\ninstance: {self._get_instance()}' | |
| print_str += f'\ncollection: {self._get_name()}' | |
| key_str = ', '.join([i for i in self.get_keys() if i != 'values']) | |
| print_str += f'\nkeys: [{key_str}]' | |
| storage_size = self._db_obj.command('collstats', self._get_name())['size'] / 1e6 | |
| print_str += f'\nstorage size: {storage_size:.1f}MB' | |
| num_docs = self._get_cursor().count() | |
| print_str += f'\ndocuments: {num_docs}' | |
| return print_str | |
| def _get_collection(self): | |
| return self._collection | |
| def _get_instance(self): | |
| return self._instance | |
| def _get_name(self): | |
| return self._name | |
| def _get_host_and_port(self): | |
| user = self._get_credentials()['user'] | |
| password = self._get_credentials()['password'] | |
| host = self._get_credentials()['host'] | |
| port = self._get_credentials()['port'] | |
| full_hostname = f"mongodb+srv://{user}:{password}@{host}/{port}" | |
| return full_hostname, port | |
| def _get_db_object(self): | |
| client = pymongo.MongoClient(host=self._get_host_and_port()[0], port=self._get_host_and_port()[1]) | |
| return client[self._get_db_name()] | |
| def _get_db_name(self): | |
| return self._get_credentials()['dbname'] | |
| def _set_advanced_query(self): | |
| if self._is_adv_query: | |
| and_ls = [] | |
| ls_dd = self._product_dict(**self._query) | |
| for dd in ls_dd: | |
| and_ls.extend([{'$and': [dd]}]) | |
| or_dd = {'$or': and_ls} | |
| self._query = or_dd | |
| def _get_query(self): | |
| return self._query | |
| def _get_cursor(self): | |
| if self._out is None: | |
| self._cursor = self._collection.find(self._query) | |
| return self._cursor | |
| def _product_dict(**kwargs): | |
| for key, val in kwargs.items(): # make sure adv query values are all lists | |
| if not isinstance(val, list): | |
| kwargs.update({key: [val]}) | |
| keys = kwargs.keys() | |
| vals = kwargs.values() | |
| for instance in itertools.product(*vals): | |
| yield dict(zip(keys, instance)) | |
| def _set_query(self, query): | |
| query_dd = dict() | |
| for k, v in query.items(): | |
| if v is not None: # drop None from queries, but allow it to pass | |
| if isinstance(v, str) and '~' in v: | |
| query_dd[k] = {'$ne': v.replace('~', '')} | |
| elif isinstance(v, list): | |
| new_ls = list() | |
| neg_ls = list() | |
| for i in v: | |
| if isinstance(i, str) and '~' in i: | |
| neg_ls.append(i.replace('~', '')) | |
| else: | |
| new_ls.append(i) | |
| if len(neg_ls) > 0: | |
| new_ls.append({'$nin': neg_ls}) | |
| query_dd[k] = new_ls | |
| else: | |
| query_dd[k] = v | |
| self._query = query_dd | |
| def drop_collection(self, confirm=''): | |
| if confirm != 'confirm': | |
| print('drop_collection not confirmed. Make sure you know what you are doing. Pass "confirm" to drop.') | |
| else: | |
| db_obj = self._get_db_object() | |
| db_obj.drop_collection(self._name) | |
| print('Database: %s' % self._get_db_name()) | |
| print('Collection dropped: %s' % self._name) | |
| def _reindex(self, confirm=''): | |
| if confirm != 'confirm': | |
| print('reindex not confirmed. Nothing done. Pass "confirm" to reindex.') | |
| else: | |
| collection = self._get_collection() | |
| collection.drop_indexes() | |
| idx_tup_ls = [(i, pymongo.ASCENDING) for i in self._key_ls] | |
| collection.create_index(idx_tup_ls, unique=True, name=self._name) # this will index it | |
| print('collection %s successfully re-indexed' % self._name) | |
| def count(self): | |
| print(f'Database: {self._get_db_name()}') | |
| print(f'Collection: {self._name}') | |
| print(f'Query: {self._query}') | |
| print(f'Documents: {self._collection.count_documents(self._query)}') | |
| def replace_field_value(self, existing_field=True, **kwargs): | |
| """ | |
| make sure you declare the appropriate key=value arguments to be modified in your query! | |
| """ | |
| if len(kwargs) != 1: | |
| raise Exception('please rename only one field at a time') | |
| field = list(kwargs.keys())[0] | |
| value = kwargs[field] | |
| if field not in self.get_keys() and existing_field: # check init that everything is ok to reindex | |
| raise Exception('new field needs to be declared in collection keys before replace / re-index') | |
| if field == 'values': | |
| raise Exception('values cannot be replace with this method. Use .update_values() instead') | |
| if existing_field: | |
| if not self._query: | |
| raise Exception('Cannot replace key value without filters.') | |
| if self._get_boolean_keys() and field in self._get_boolean_keys(): | |
| value = str(int(value)) | |
| cursor = self._get_cursor() | |
| request_ls = [] | |
| for doc in cursor: | |
| if existing_field: | |
| doc.pop(field) | |
| update_dd = {'$set': {field: value}} | |
| request = pymongo.UpdateOne(doc, update_dd, upsert=False) | |
| request_ls.append(request) | |
| if len(request_ls) > 0: | |
| self._get_collection().bulk_write(request_ls) | |
| print(f'collection field {field} successfully updated to {value}') | |
| if not existing_field: | |
| self._reindex('confirm') | |
| else: | |
| print('no operations to handle') | |
| def get_keys(self): | |
| return [i for i in self._key_ls] # pass a copy | |
| def get_key_values(self): | |
| match_dd = {'$match': self._query} | |
| key_dd = {k: f'${k}' for k in self._key_ls} | |
| group_dd = {'$group': {'_id': key_dd}} | |
| cursor = self._collection.aggregate([match_dd, group_dd]) | |
| ls = [i['_id'] for i in list(cursor)] | |
| keys_df = pd.DataFrame(ls) | |
| if not keys_df.empty: | |
| keys_df = keys_df.sort_values(by=self._key_ls) | |
| keys_df = keys_df.reset_index(drop=True) | |
| return keys_df | |
| def _is_local_instance(self): | |
| return self._get_instance() == 'local' | |
| def delete_many(self, confirm=''): | |
| if not self._query: | |
| raise Exception('Cannot delete_many without filters. Use drop_collection() instead for a full drop.') | |
| lmt = self._del_many_limit | |
| cnt = self._collection.count_documents(self._query) | |
| if (cnt > lmt) and confirm != 'confirm': | |
| raise Exception(f'confirm delete_many() if more than {lmt} documents. Potentially deleted: {cnt}') | |
| before = self._collection.count_documents({}) | |
| self._get_collection().delete_many(self._query) | |
| after = self._collection.count_documents({}) | |
| self._out = None | |
| print(f'{before - after} document(s) deleted') | |
| def distinct(self, field): | |
| if isinstance(field, str): | |
| field = field | |
| cursor = self._get_cursor() | |
| return cursor.distinct(field) # faster | |
| elif isinstance(field, list): | |
| group_dd = {"_id": {i: f'${i}' for i in field}} | |
| if len(self._query) > 0: | |
| cursor = self._collection.aggregate([{'$match': self._query}, | |
| {"$group": group_dd}]) | |
| else: | |
| cursor = self._collection.aggregate([{"$group": group_dd}]) | |
| dd = [doc['_id'] for doc in cursor] | |
| return pd.DataFrame().from_dict(dd) | |
| class _MongoSeries(_MongoAbstract): | |
| _drop_weekends = None | |
| def __init__(self, instance='remote', **query): | |
| _MongoAbstract.__init__(self, instance=instance, **query) | |
| def _get_drop_weekends(self): | |
| return self._drop_weekends | |
| def _update_query_date_range(self, start, end): | |
| if start is not None and end is None: | |
| self._query['values.date'] = start | |
| # self._query['values'] = {'$elemMatch': {'date': {'$gte': start}}} | |
| elif start is None and end is not None: | |
| self._query['values.date'] = {'$lte': end} | |
| elif start is not None and end is not None: | |
| self._query['values.date'] = {'$gte': start, '$lte': end} | |
| else: # both none | |
| pass | |
| def _get_query_cursor(self, start, end): | |
| if start is None and end is None: | |
| return self._get_cursor() # simple case | |
| else: | |
| co = self._get_collection() | |
| key_ls = self.get_keys() | |
| match_dd = {'$match': self._query} | |
| if start is not None and end is None: | |
| if isinstance(start, (dt.datetime, dt.date)): | |
| start = start.strftime('%Y%m%d') | |
| start = start.replace('-', '') | |
| filter_dd = { | |
| '$filter': | |
| {'input': '$values', | |
| 'as': 'vals', | |
| 'cond': {'$gte': ['$$vals.date', start]} | |
| } | |
| } | |
| elif start is None and end is not None: | |
| if isinstance(end, dt.datetime): | |
| end = end.strftime('%Y%m%d') | |
| end = end.replace('-', '') | |
| filter_dd = { | |
| '$filter': | |
| {'input': '$values', | |
| 'as': 'vals', | |
| 'cond': {'$lte': ['$$vals.date', end]} | |
| } | |
| } | |
| else: # both defined | |
| if isinstance(start, dt.datetime): | |
| start = start.strftime('%Y%m%d') | |
| start = start.replace('-', '') | |
| if isinstance(end, dt.datetime): | |
| end = end.strftime('%Y%m%d') | |
| end = end.replace('-', '') | |
| filter_dd = { | |
| '$filter': | |
| {'input': '$values', | |
| 'as': 'vals', | |
| 'cond': {'$and': [{'$gte': ['$$vals.date', start]}, | |
| {'$lte': ['$$vals.date', end]}]} | |
| } | |
| } | |
| key_dd = {i: 1 for i in key_ls} | |
| key_dd['values'] = filter_dd | |
| project_dd = {'$project': key_dd} | |
| cursor = co.aggregate([match_dd, project_dd]) | |
| return cursor | |
| def insert(self, pd_obj, col_key): | |
| if pd_obj.empty: | |
| print('pd_obj empty, no data to insert') | |
| return | |
| if isinstance(pd_obj, pd.Series): | |
| pd_obj = pd.DataFrame(pd_obj) | |
| # check no dups | |
| if pd_obj.columns.has_duplicates: | |
| raise Exception('inserted object cannot have duplicated keys!') | |
| # initial checks | |
| if isinstance(col_key, str): | |
| col_key = [col_key] | |
| if len(self._key_ls) != len(col_key): | |
| name = self._name | |
| full = len(self._key_ls) | |
| given = len(col_key) | |
| raise Exception(f'insert error: {name}. col_key must have {full} elements. provided was {given}') | |
| df = pd_obj.copy() | |
| df.index = df.index.map(lambda i: i.strftime('%Y%m%d')) | |
| # integrity check | |
| for col in pd_obj.columns: | |
| col_val_ls = col.split('.') | |
| if len(self._full_key_ls) - 1 != len(col_val_ls): | |
| raise Exception(f'all columns must have {len(col_key)} elements. provided was {col_val_ls}') | |
| # do not store full NaN rows or columns | |
| df = df.dropna(how='all', axis=0) | |
| df = df.dropna(how='all', axis=1) | |
| request_ls = [] | |
| for col in df.columns: | |
| col_val_ls = col.split('.') | |
| tup = zip(col_key, col_val_ls) | |
| query_dd = {k: v for k, v in tup} | |
| values_dd = df[col].dropna().to_dict() | |
| values_ls = [{'date': k, 'value': v} for k, v in values_dd.items()] | |
| update_dd = {'$set': {'values': values_ls}} | |
| request = pymongo.UpdateOne(query_dd, update_dd, upsert=True) | |
| request_ls.append(request) | |
| self._get_collection().bulk_write(request_ls) | |
| def last_update(self, how='each', order='min'): | |
| match_dd = {'$match': self._query} | |
| key_dd = {k: f'${k}' for k in self._key_ls} | |
| group_dd = { | |
| '$group': | |
| {'_id': | |
| {'_id': key_dd}, | |
| 'max': {'$max': "$values.date"} | |
| } | |
| } | |
| sort_dd = {'$sort': {'values.date': 1}} | |
| cursor = self._get_collection().aggregate([match_dd, group_dd, sort_dd]) | |
| last_df = pd.DataFrame() | |
| for num, doc in enumerate(cursor): | |
| dd = doc['_id']['_id'] | |
| dd['last_update'] = doc['max'][-1] | |
| tmp_df = pd.DataFrame.from_dict(dd, orient='index').T | |
| last_df = pd.concat([last_df, tmp_df], axis=0, sort=True) | |
| last_df = last_df[self._key_ls + ['last_update']] | |
| last_df['last_update'] = last_df['last_update'].apply(pd.to_datetime) | |
| days_f = lambda x: (dt.datetime.today() - x).days | |
| last_df['last_update_days'] = last_df['last_update'].apply(days_f) | |
| if how == 'all': | |
| if order == 'max': | |
| return last_df['last_update'].max() | |
| elif order == 'min': | |
| return last_df['last_update'].min() | |
| else: | |
| raise Exception(f'argument order not recognized {order}') | |
| elif how == 'each': | |
| return last_df | |
| else: | |
| raise Exception(f'argument how not recognized: {how}') | |
| def query(self, start=None, end=None, rename=None, expand=False): | |
| """ | |
| override is for self.last_update() | |
| """ | |
| if self._out is None: | |
| cursor = self._get_query_cursor(start, end) | |
| if rename is not None: | |
| if isinstance(rename, str): | |
| rename = [rename] | |
| exclude_ls = ['_id', 'values'] | |
| df = pd.DataFrame() | |
| for doc in cursor: | |
| if len(doc['values']) > 0: | |
| if rename is not None: | |
| name = '.'.join([doc[i] for i in rename]) | |
| else: | |
| name = '.'.join([doc[i] for i in sorted(doc.keys()) if i not in exclude_ls]) | |
| doc_df = pd.DataFrame().from_dict(doc['values']).set_index('date') | |
| doc_df.columns = [name] | |
| df = pd.concat([df, doc_df], axis=1, sort=True) | |
| if df.empty: | |
| return df | |
| df.index = df.index.map(pd.to_datetime) | |
| if self._get_drop_weekends(): | |
| week_days_ls = df.index.weekday < 5 | |
| df = df.loc[week_days_ls].copy() | |
| if expand: | |
| df.columns = df.columns.str.split('.', expand=True) | |
| # return series if only one column or row | |
| if len(df.columns) == 1: | |
| df = df[df.columns[0]].copy() | |
| if len(df.index) == 1: | |
| df = df.iloc[-1].copy() | |
| self._out = df.copy() | |
| return self._out | |
| def drop_datapoint(self, date_str): | |
| if self._collection.count_documents(self._query) == 0: | |
| raise Exception('no documents found with given filters') | |
| if self._collection.count_documents(self._query) > 1: | |
| raise Exception('you can only drop datapoints one series at a time') | |
| if isinstance(date_str, dt.datetime): | |
| date_str = date_str.strftime('%Y%m%d') | |
| cursor = self._get_cursor() | |
| request_ls = [] | |
| for doc in cursor: | |
| doc.pop('_id') | |
| doc.pop('values') | |
| doc['values.date'] = date_str | |
| update_dd = {'$unset': {'values.$.value': ""}} | |
| request = pymongo.UpdateOne(doc, update_dd, upsert=False) | |
| request_ls.append(request) | |
| self._get_collection().bulk_write(request_ls) | |
| def update_values(self, pd_obj, col_key): | |
| if pd_obj.empty: | |
| print('pd_obj empty, no data to update') | |
| return | |
| # initial checks | |
| if isinstance(col_key, str): | |
| col_key = [col_key] | |
| if len(self._key_ls) != len(col_key): | |
| name = self._name | |
| full = len(self._key_ls) | |
| given = len(col_key) | |
| raise Exception(f'insert error: {name}. col_key must have {full} elements. provided was {given}') | |
| if isinstance(pd_obj, pd.Series): | |
| df = pd.DataFrame(pd_obj) | |
| else: | |
| df = pd_obj.copy() | |
| # check no dups | |
| if df.columns.has_duplicates: | |
| raise Exception('inserted object cannot have duplicated keys!') | |
| # integrity check | |
| for col in df.columns: | |
| col_val_ls = col.split('.') | |
| if len(self._full_key_ls) - 1 != len(col_val_ls): | |
| raise Exception(f'all columns must have {len(col_key)} elements. provided was {col_val_ls}') | |
| df.index = df.index.map(lambda i: i.strftime('%Y%m%d')) | |
| df = df.dropna(axis=0, how='all') | |
| df = df.dropna(axis=1, how='all') | |
| request_ls = [] | |
| for col in df.columns: | |
| col_val_ls = col.split('.') | |
| tup = zip(col_key, col_val_ls) | |
| query_dd = {k: v for k, v in tup} | |
| values_dd = df[col].dropna().to_dict() | |
| # pull many | |
| date_ls = list(values_dd.keys()) | |
| update_dd = { | |
| '$pull': {'values': {'date': {'$in': date_ls}}} | |
| } | |
| request = pymongo.UpdateMany(query_dd, update_dd, upsert=True) | |
| request_ls.append(request) | |
| # push many | |
| new_ls = [{'date': k, 'value': v} for k, v in values_dd.items()] | |
| update_dd = {'$push': {'values': {'$each': new_ls}}} | |
| request = pymongo.UpdateMany(query_dd, update_dd, upsert=True) | |
| request_ls.append(request) | |
| self._get_collection().bulk_write(request_ls, ordered=True) | |
| class _MongoDoc(_MongoAbstract): | |
| def __init__(self, instance='remote', **query): | |
| _MongoAbstract.__init__(self, instance=instance, **query) | |
| def insert(self, dd): | |
| if not isinstance(dd, dict): | |
| print('insert dict empty, no data to insert') | |
| return | |
| # initial checks | |
| if sorted(list(self._get_query().keys())) != sorted(self.get_keys()): | |
| raise Exception(f'specify a full query to insert into a _MongoTable! Keys are: {str(self.get_keys())}') | |
| cast_key_ls = set([str(k) for k in dd.keys()]) | |
| if len(cast_key_ls) != len(dd.keys()): | |
| raise Exception('identical string representations of keys are not allowed') | |
| insert_dd = {str(k): v for k, v in dd.items()} | |
| filter_dd = self._get_query() | |
| update_dd = {'$set': {'values': insert_dd}} | |
| self._get_collection().update_one(filter_dd, update=update_dd, upsert=True) | |
| def query(self, expand=False, rename=None): | |
| """ | |
| override is for self.last_update() | |
| """ | |
| if rename: | |
| if isinstance(rename, str): | |
| rename = [rename] | |
| if self._out is None: | |
| cursor = self._get_cursor() # simple case | |
| df = pd.DataFrame() | |
| for num, doc in enumerate(cursor): | |
| if doc['values']: | |
| doc.pop('_id') | |
| val_dd = doc.pop('values') | |
| if rename: | |
| name = '.'.join([str(doc[k]) for k in rename]) | |
| else: | |
| name = '.'.join([str(doc[k]) for k in sorted(doc.keys())]) | |
| s = pd.Series(val_dd, name=name, dtype=object) | |
| df = pd.concat([df, s], axis=1, sort=True) | |
| if expand: | |
| df.columns = df.columns.str.split('.', expand=True) | |
| # return series if only one column or row | |
| if len(df.columns) == 1: | |
| df = df[df.columns[0]].copy() | |
| if len(df.index) == 1: | |
| df = df.iloc[-1].copy() | |
| self._out = df | |
| return self._out | |
| class _MongoTable(_MongoAbstract): | |
| def __init__(self, instance='remote', **query): | |
| _MongoAbstract.__init__(self, instance=instance, **query) | |
| def insert(self, pd_obj): | |
| if pd_obj.empty: | |
| print('pd_obj empty, no data to insert') | |
| return | |
| if isinstance(pd_obj, pd.Series): | |
| pd_obj = pd.DataFrame(pd_obj) | |
| # check no dups | |
| if pd_obj.columns.has_duplicates: | |
| raise Exception('inserted object cannot have duplicated columns!') | |
| # initial checks | |
| if sorted(list(self._get_query().keys())) != sorted(self.get_keys()): | |
| raise Exception(f'specify a full query to insert into a _MongoTable! Keys are: {str(self.get_keys())}') | |
| # drop nan | |
| df = pd_obj.copy() | |
| df = df.dropna(how='all', axis=0) | |
| df = df.dropna(how='all', axis=1) | |
| if isinstance(df.columns, pd.core.indexes.multi.MultiIndex): | |
| df.columns = ['.'.join([i for i in col]) for col in df.columns] | |
| if isinstance(df.index, pd.core.indexes.multi.MultiIndex): | |
| df.index = ['.'.join([i for i in col]) for col in df.columns] | |
| df.columns = df.columns.map(str) | |
| df.index = df.index.map(str) | |
| values_dd = df.to_dict(orient='records') | |
| filter_dd = self._get_query() | |
| update_dd = {'$set': {'values': values_dd}} | |
| self._get_collection().update_one(filter_dd, update=update_dd, upsert=True) | |
| def query(self, rename=None, expand=False): | |
| """ | |
| override is for self.last_update() | |
| """ | |
| if self._out is None: | |
| cursor = self._get_cursor() # simple case | |
| if rename is not None: | |
| if isinstance(rename, str): | |
| rename = [rename] | |
| exclude_ls = ['_id', 'values'] | |
| out_df = pd.DataFrame() | |
| cnt = 0 | |
| for num, doc in enumerate(cursor): | |
| if len(doc['values']) > 0: | |
| df = pd.DataFrame().from_records(doc['values']) | |
| if rename is not None: | |
| name = '.'.join([doc[i] for i in rename]) | |
| else: | |
| name = '.'.join([doc[i] for i in sorted(doc.keys()) if i not in exclude_ls]) | |
| id_ls = [name] * df.shape[0] | |
| df['_id'] = id_ls | |
| out_df = pd.concat([out_df, df], axis=0, sort=True) | |
| cnt += 1 | |
| if out_df.empty: | |
| return out_df | |
| if cnt > 1: | |
| out_df = out_df.reset_index(drop=False) | |
| out_df = out_df.set_index('_id', drop=True) | |
| out_df.index.name = None | |
| else: | |
| out_df = out_df.drop('_id', axis=1) | |
| if expand: | |
| out_df.index = out_df.index.str.split('.', expand=True) | |
| # return series if only one column or row | |
| if len(out_df.columns) == 1: | |
| out_df = out_df[out_df.columns[0]].copy() | |
| if len(out_df.index) == 1: | |
| out_df = out_df.iloc[-1].copy() | |
| self._out = out_df.copy() | |
| return self._out | |
| class _MongoLog(_MongoAbstract): | |
| def __init__(self, instance='remote', **query): | |
| _MongoAbstract.__init__(self, instance=instance, **query) | |
| def _log(self, df): | |
| request_ls = [] | |
| if not df.empty: | |
| for (section, iso_dt), row_s in df.iterrows(): | |
| values_dd = row_s.to_dict() | |
| doc_id = {'section': section, | |
| 'iso_date': iso_dt} | |
| update_dd = {'$push': {'values': values_dd}} | |
| request = pymongo.UpdateOne(doc_id, update_dd, upsert=True) | |
| request_ls.append(request) | |
| self._get_collection().bulk_write(request_ls) | |
| def query(self, rename=None, expand=False): | |
| """ | |
| override is for self.last_update() | |
| """ | |
| if self._out is None: | |
| cursor = self._get_cursor() # simple case | |
| if rename is not None: | |
| if isinstance(rename, str): | |
| rename = [rename] | |
| exclude_ls = ['_id', 'values'] | |
| df = pd.DataFrame() | |
| for doc in cursor: | |
| if len(doc['values']) > 0: | |
| if rename is not None: | |
| name = '.'.join([doc[i] for i in rename]) | |
| else: | |
| name = '.'.join([doc[i] for i in sorted(doc.keys()) if i not in exclude_ls]) | |
| doc_df = pd.DataFrame().from_records(doc['values']) | |
| doc_df.index = [name] * len(doc_df.index) | |
| df = pd.concat([df, doc_df], axis=0, sort=True) | |
| if df.empty: | |
| return df | |
| if expand: | |
| df.index = df.index.str.split('.', expand=True) | |
| # return series if only one column or row | |
| if len(df.columns) == 1: | |
| df = df[df.columns[0]].copy() | |
| if len(df.index) == 1: | |
| df = df.iloc[-1].copy() | |
| self._out = df.copy() | |
| return self._out | |