Add to_json() and from_json() methods to Cohort class

raiwidgets/raiwidgets/cohort.py

@@ -3,6 +3,7 @@
 
 """Module for defining cohorts in raiwidgets package."""
 
+import json
 from typing import Any, List, Optional
 
 import numpy as np
@@ -397,6 +398,75 @@ def __init__(self, name: str):
         self.name = name
         self.cohort_filter_list = None
 
+    @staticmethod
+    def _cohort_serializer(obj):
+        """The function to serialize the Cohort class object.
+
+        :param obj: Any member of the Cohort class object.
+        :type: Any
+        :return: Python dictionary.
+        :rtype: dict[Any, Any]
+        """
+        return obj.__dict__
+
+    def to_json(self):
+        """Returns a serialized JSON string for the Cohort object.
+
+        :return: The JSON string for the cohort.
+        :rtype: str
+        """
+        return json.dumps(self, default=Cohort._cohort_serializer)
+
+    @staticmethod
+    def _get_cohort_object(json_dict):
+        """Method to read a JSON dictionary and return a Cohort object.
+
+        :param json_dict: JSON dictionary containing cohort data.
+        :type: dict[str, str]
+        :return: The Cohort object.
+        :rtype: Cohort
+        """
+        if "name" not in json_dict:
+            raise UserConfigValidationException(
+                "No name field found for cohort deserialization")
+        if "cohort_filter_list" not in json_dict:
+            raise UserConfigValidationException(
+                "No cohort_filter_list field found for cohort deserialization")
+        if not isinstance(json_dict['cohort_filter_list'], list):
+            raise UserConfigValidationException(
+                "Field cohort_filter_list not of type list for "
+                "cohort deserialization")
+
+        deserialized_cohort = Cohort(json_dict['name'])
+        for serialized_cohort_filter in json_dict['cohort_filter_list']:
+            if "method" not in serialized_cohort_filter:
+                raise UserConfigValidationException(
+                    "Field method not found for cohort filter deserialization")
+            if "arg" not in serialized_cohort_filter:
+                raise UserConfigValidationException(
+                    "Field arg not found for cohort filter deserialization")
+            if "column" not in serialized_cohort_filter:
+                raise UserConfigValidationException(
+                    "Field column not found for cohort filter deserialization")
+            cohort_filter = CohortFilter(
+                method=serialized_cohort_filter['method'],
+                arg=serialized_cohort_filter['arg'],
+                column=serialized_cohort_filter['column'])
+            deserialized_cohort.add_cohort_filter(cohort_filter=cohort_filter)
+        return deserialized_cohort
+
+    @staticmethod
+    def from_json(json_str):
+        """Method to read a json string and return a Cohort object.
+
+        :param json_str: Serialized JSON string.
+        :type: str
+        :return: The Cohort object.
+        :rtype: Cohort
+        """
+        json_dict = json.loads(json_str)
+        return Cohort._get_cohort_object(json_dict)
+
     def add_cohort_filter(self, cohort_filter: CohortFilter):
         """Add a cohort filter into the cohort.
         :param cohort_filter: Cohort filter defined by CohortFilter class.

raiwidgets/tests/test_cohort.py

@@ -497,40 +497,46 @@ def test_cohort_serialization_single_value_method(self, method):
                                        arg=[65], column='age')
         cohort_1 = Cohort(name="Cohort New")
         cohort_1.add_cohort_filter(cohort_filter_1)
-        json_str = json.dumps(cohort_1,
-                              default=cohort_filter_json_converter)
+        json_str = cohort_1.to_json()
 
         assert 'Cohort New' in json_str
         assert method in json_str
         assert '[65]' in json_str
         assert 'age' in json_str
 
-    def test_cohort_serialization_in_range_method(self):
+    def test_cohort_serialization_deserialization_in_range_method(self):
         cohort_filter_1 = CohortFilter(
             method=CohortFilterMethods.METHOD_RANGE,
             arg=[65.0, 70.0], column='age')
         cohort_1 = Cohort(name="Cohort New")
         cohort_1.add_cohort_filter(cohort_filter_1)
-        json_str = json.dumps(cohort_1,
-                              default=cohort_filter_json_converter)
 
+        json_str = cohort_1.to_json()
         assert 'Cohort New' in json_str
         assert CohortFilterMethods.METHOD_RANGE in json_str
         assert '65.0' in json_str
         assert '70.0' in json_str
         assert 'age' in json_str
 
+        cohort_1_new = Cohort.from_json(json_str)
+        assert cohort_1_new.name == cohort_1.name
+        assert len(cohort_1_new.cohort_filter_list) == \
+            len(cohort_1.cohort_filter_list)
+        assert cohort_1_new.cohort_filter_list[0].method == \
+            cohort_1.cohort_filter_list[0].method
+
     @pytest.mark.parametrize('method',
                              [CohortFilterMethods.METHOD_INCLUDES,
                               CohortFilterMethods.METHOD_EXCLUDES])
-    def test_cohort_serialization_include_exclude_methods(self, method):
+    def test_cohort_serialization_deserialization_include_exclude_methods(
+            self, method):
         cohort_filter_str = CohortFilter(method=method,
                                          arg=['val1', 'val2', 'val3'],
                                          column='age')
         cohort_str = Cohort(name="Cohort New Str")
         cohort_str.add_cohort_filter(cohort_filter_str)
-        json_str = json.dumps(cohort_str,
-                              default=cohort_filter_json_converter)
+
+        json_str = cohort_str.to_json()
         assert method in json_str
         assert 'val1' in json_str
         assert 'val2' in json_str
@@ -542,14 +548,21 @@ def test_cohort_serialization_include_exclude_methods(self, method):
                                          column='age')
         cohort_int = Cohort(name="Cohort New Int")
         cohort_int.add_cohort_filter(cohort_filter_int)
-        json_str = json.dumps(cohort_filter_int,
-                              default=cohort_filter_json_converter)
+
+        json_str = cohort_int.to_json()
         assert method in json_str
         assert '1' in json_str
         assert '2' in json_str
         assert '3' in json_str
         assert 'age' in json_str
 
+        cohort_int_new = Cohort.from_json(json_str)
+        assert cohort_int_new.name == cohort_int.name
+        assert len(cohort_int_new.cohort_filter_list) == \
+            len(cohort_int.cohort_filter_list)
+        assert cohort_int_new.cohort_filter_list[0].method == \
+            cohort_int.cohort_filter_list[0].method
+
 
 class TestCohortList:
     def test_cohort_list_serialization(self):