def convert_dataframe(self, df: DataFrame) -> List[NormalizedData]:
name_mapping = {v.get("name"): k for k, v in self.mapping.items()}
self.fill_missing_columns(df, name_mapping)
df_ordernet = df[name_mapping.keys()]
df_ordernet.rename(columns=name_mapping, inplace=True)
self.run_pipeline(df_ordernet)
normalized_list = df_ordernet.to_dict(orient="records")
return [NormalizedData(**item) for item in normalized_list]
def get_user_normalized_data_in_range(
user_id: ObjectId,
start_date: dt = None,
end_date: dt = None) -> Generator[BaseData, None, None]:
query = {"user_id": user_id}
if start_date or end_date:
query["date"] = {}
if start_date:
query["date"]["$gte"] = start_date
if end_date:
query["date"]["$lt"] = end_date
raw_data_results = mongo_crud_wrapper.get_by_query(COLLECTION_NAME, query)
return (NormalizedData.from_doc(result) for result in raw_data_results)
def update(raw_data: NormalizedData) -> UpdateResult:
raw_data_doc = raw_data.to_doc()
return mongo_crud_wrapper.update(COLLECTION_NAME, raw_data_doc)
def get_bulk_by_user(user_id: ObjectId) -> Generator[BaseData, None, None]:
query = {"user_id": user_id}
raw_data_results = mongo_crud_wrapper.get_by_query(COLLECTION_NAME, query)
return (NormalizedData.from_doc(result) for result in raw_data_results)
def get(raw_data_id: ObjectId) -> BaseData:
raw_data_id_doc = mongo_crud_wrapper.get(COLLECTION_NAME, raw_data_id)
return NormalizedData.from_doc(raw_data_id_doc)
def create(raw_data: NormalizedData) -> BaseData:
raw_data_doc = raw_data.to_doc()
mongo_crud_wrapper.create(COLLECTION_NAME, raw_data_doc)
return NormalizedData.from_doc(raw_data_doc)
def test_to_json(normalized_data: NormalizedData, object_id):
normalized_data.uid = object_id
normalized_data_json = normalized_data.to_json()
loaded_json = NormalizedData.from_json(normalized_data_json)
assert loaded_json == normalized_data
def test_get_doc(normalized_data: NormalizedData, object_id):
normalized_data.uid = object_id
normalized_data_doc = normalized_data.to_doc()
data_obj = NormalizedData.from_doc(normalized_data_doc)
assert data_obj == normalized_data