From 97c693d2528f2db42c5ea8188ea3c546942646bb Mon Sep 17 00:00:00 2001
From: junos <junos.lukan@ijs.si>
Date: Thu, 19 Aug 2021 17:36:26 +0200
Subject: [PATCH] Add a getter for communication data.

---
 machine_learning/pipeline.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/machine_learning/pipeline.py b/machine_learning/pipeline.py
index f5c35ae..f08d2cf 100644
--- a/machine_learning/pipeline.py
+++ b/machine_learning/pipeline.py
@@ -50,6 +50,8 @@ class SensorFeatures:
     def get_sensor_data(self, data_type) -> pd.DataFrame:
         if data_type == "proximity":
             return self.df_proximity
+        elif data_type == "communication":
+            return self.df_calls_sms
         else:
             raise KeyError("This data type has not been implemented.")