bug fix: mongodb inserts of structured objects with lists as elements

This commit is contained in:
Steve Nyemba 2024-04-02 12:31:41 -05:00
parent 4b97994ec1
commit eb81f5a4d2
2 changed files with 19 additions and 4 deletions

View File

@ -1,5 +1,6 @@
import json import json
import numpy as np
from datetime import datetime
class IEncoder (json.JSONEncoder): class IEncoder (json.JSONEncoder):
def default (self,object): def default (self,object):

View File

@ -213,10 +213,24 @@ class Writer(Mongo):
_uid = self.collection if 'doc' not in _args else _args['doc'] _uid = self.collection if 'doc' not in _args else _args['doc']
if self._lock : if self._lock :
Mongo.lock.acquire() Mongo.lock.acquire()
if type(info) == list or type(info) == pd.DataFrame : if type(info) == list or type(info) == pd.DataFrame :
info if type(info) == list else info.to_dict(orient='records') if type(info) == pd.DataFrame :
info = json.loads(json.dumps(info,cls=IEncoder)) info = info.to_dict(orient='records')
# info if type(info) == list else info.to_dict(orient='records')
info = json.loads(json.dumps(info))
self.db[_uid].insert_many(info) self.db[_uid].insert_many(info)
else:
#
# sometimes a dictionary can have keys with arrays (odd shaped)
#
_keycount = len(info.keys())
_arraycount = [len(info[key]) for key in info if type(info[key]) in (list,np.array,np.ndarray)]
if _arraycount and len(_arraycount) == _keycount and np.max(_arraycount) == np.min(_arraycount) :
#
# In case an object with consistent structure is passed, we store it accordingly
#
self.write(pd.DataFrame(info),**_args)
else: else:
self.db[_uid].insert_one(json.loads(json.dumps(info,cls=IEncoder))) self.db[_uid].insert_one(json.loads(json.dumps(info,cls=IEncoder)))
finally: finally: