131 lines
3.7 KiB
Python
131 lines
3.7 KiB
Python
import pymongo
|
|
import xarray_mongodb
|
|
import bson
|
|
import builtins
|
|
|
|
from ToolFunction.ToolFunction import get_date
|
|
|
|
|
|
npTypeDict = {v: getattr(builtins, k) for k, v in np.sctypeDict.items() if k in vars(builtins)}
|
|
npArrayType = type(np.array([0]))
|
|
|
|
|
|
class MongoDB:
|
|
|
|
def __init__(self, mongoClient, mongoDB, date=None) -> None:
|
|
self.mongoClient = mongoClient
|
|
self.mongoDB = mongoDB
|
|
self.xdb = xarray_mongodb.XarrayMongoDB(mongoDB)
|
|
|
|
if date is None:
|
|
date= get_date()
|
|
self.set_date(date)
|
|
|
|
def _convert_numpy_type(self, data):
|
|
for key in data:
|
|
typeKey = type(data[key])
|
|
if typeKey in npTypeDict:
|
|
data[key] = data[key].item()
|
|
elif typeKey == npArrayType:
|
|
data[key] = data[key].tolist()
|
|
else:
|
|
try:
|
|
data[key] = data[key].item()
|
|
except:
|
|
pass
|
|
return data
|
|
|
|
def set_date(self, date):
|
|
date = date.split("/")
|
|
self.year = int(date[0])
|
|
self.month = int(date[1])
|
|
self.day = int(date[2])
|
|
|
|
def create_global(self, shotNum, dataSet=None, date=None):
|
|
if not date is None:
|
|
self.set_date(date)
|
|
|
|
data = {
|
|
'year': self.year,
|
|
'month': self.month,
|
|
'day': self.day,
|
|
'shotNum': shotNum,
|
|
'global_parameters' : {},
|
|
}
|
|
|
|
global_parameters = self._convert_numpy_type(dataSet.attrs)
|
|
|
|
if not dataSet is None:
|
|
data['global_parameters'].update(global_parameters)
|
|
|
|
data = self._convert_numpy_type(data)
|
|
|
|
self.mongoDB['global'].insert_one(data)
|
|
|
|
def _add_data_normal(self, shotNum, data):
|
|
filter = {
|
|
'year': self.year,
|
|
'month': self.month,
|
|
'day': self.day,
|
|
'shotNum': shotNum,
|
|
}
|
|
|
|
self.mongoDB['global'].update_one(filter, {"$set": data}, upsert=False)
|
|
|
|
def _add_data_xarray_dataArray(self, shotNum, dataArray):
|
|
filter = {
|
|
'year': self.year,
|
|
'month': self.month,
|
|
'day': self.day,
|
|
'shotNum': shotNum,
|
|
}
|
|
|
|
dataArray.attrs = self._convert_numpy_type(dataArray.attrs)
|
|
|
|
mongoID, _ = self.xdb.put(dataArray)
|
|
|
|
data_label = {
|
|
dataArray.name:
|
|
{
|
|
'name': dataArray.name,
|
|
'mongoID': mongoID,
|
|
'engine': 'xarray',
|
|
'dtype': 'dataArray',
|
|
}
|
|
}
|
|
|
|
self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False)
|
|
|
|
def _add_data_xarray_dataSet(self, shotNum, dataSet, name):
|
|
filter = {
|
|
'year': self.year,
|
|
'month': self.month,
|
|
'day': self.day,
|
|
'shotNum': shotNum,
|
|
}
|
|
|
|
dataSet.attrs = self._convert_numpy_type(dataSet.attrs)
|
|
|
|
for key in list(dataSet.data_vars):
|
|
dataSet[key].attrs = self._convert_numpy_type(dataSet[key].attrs)
|
|
print(key)
|
|
|
|
mongoID, _ = self.xdb.put(dataSet)
|
|
|
|
data_label = {
|
|
name:
|
|
{
|
|
'name': name,
|
|
'mongoID': mongoID,
|
|
'engine': 'xarray',
|
|
'dtype': 'dataSet',
|
|
}
|
|
}
|
|
|
|
self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False)
|
|
|
|
def add_data(self, shotNum, data, date=None, name=None, engine='normal'):
|
|
pass
|
|
|
|
def read_data(self, shotNum, data=None):
|
|
pass |