Merge pull request #120 from jjjermiah/119-fix-patientcount-return

fix: return type in patientCount and Refactor response parsing
jjjermiah · Feb 25, 2024 · f0a7e11 · f0a7e11
2 parents 214151e + b119d93
commit f0a7e11
Showing 1 changed file with 21 additions and 11 deletions.
diff --git a/src/nbiatoolkit/nbia.py b/src/nbiatoolkit/nbia.py
@@ -248,10 +248,10 @@ def getCollectionDescriptions(
         ), "The response from the API is empty. Please check the collection name."
 
         response[0] = {
-            "collectionName": response[0]["collectionName"],
-            "description": clean_html(response[0]["description"]),
-            "descriptionURI": response[0]["descriptionURI"],
-            "lastUpdated": convertMillis(
+            "Collection": response[0]["collectionName"],
+            "Description": clean_html(response[0]["description"]),
+            "DescriptionURI": response[0]["descriptionURI"],
+            "LastUpdated": convertMillis(
                 millis=int(response[0]["collectionDescTimestamp"])
             ),
         }
@@ -340,16 +340,22 @@ def getCollectionPatientCount(
     ) -> List[dict[Any, Any]] | pd.DataFrame:
         returnType: ReturnType = self._get_return(return_type)
 
+        response: List[dict[Any, Any]]
         response = self.query_api(NBIA_ENDPOINTS.GET_COLLECTION_PATIENT_COUNT)
 
-        if prefix:
-            response = [
-                response_dict
-                for response_dict in response
-                if response_dict["criteria"].lower().startswith(prefix.lower())
-            ]
+        parsed_response: List[dict[Any, Any]] = []
 
-        return conv_response_list(response, returnType)
+        for collection in response:
+            Collection = collection["criteria"]
+            if Collection.lower().startswith(prefix.lower()):
+                parsed_response.append(
+                    {
+                        "Collection": Collection,
+                        "PatientCount": collection["count"],
+                    }
+                )
+
+        return conv_response_list(parsed_response, returnType)
 
     def getBodyPartCounts(
         self,
@@ -361,6 +367,7 @@ def getBodyPartCounts(
 
         PARAMS = self.parsePARAMS(locals())
 
+        response: List[dict[Any, Any]]
         response = self.query_api(
             endpoint=NBIA_ENDPOINTS.GET_BODY_PART_PATIENT_COUNT, params=PARAMS
         )
@@ -378,6 +385,7 @@ def getStudies(
 
         PARAMS: dict = self.parsePARAMS(locals())
 
+        response: List[dict[Any, Any]]
         response = self.query_api(endpoint=NBIA_ENDPOINTS.GET_STUDIES, params=PARAMS)
 
         return conv_response_list(response, returnType)
@@ -398,6 +406,7 @@ def getSeries(
 
         PARAMS: dict = self.parsePARAMS(locals())
 
+        response: List[dict[Any, Any]]
         response = self.query_api(endpoint=NBIA_ENDPOINTS.GET_SERIES, params=PARAMS)
 
         return conv_response_list(response, returnType)
@@ -459,6 +468,7 @@ def getDICOMTags(
         returnType: ReturnType = self._get_return(return_type)
         PARAMS = self.parsePARAMS({"SeriesUID": SeriesInstanceUID})
 
+        response: List[dict[Any, Any]]
         response = self.query_api(endpoint=NBIA_ENDPOINTS.GET_DICOM_TAGS, params=PARAMS)
 
         return conv_response_list(response, returnType)