Added the find section by code functionality

healthchain/cda_parser/cdaannotator.py

@@ -15,8 +15,8 @@
 )
 from healthchain.models.data.concept import Concept, Quantity, Range, TimeInterval
 
-from .model.cda import ClinicalDocument
-from .model.sections import (
+from healthchain.cda_parser.model.cda import ClinicalDocument
+from healthchain.cda_parser.model.sections import (
     Entry,
     Section,
     EntryRelationship,
@@ -138,6 +138,13 @@ class SectionId(Enum):
     NOTE = "1.2.840.114350.1.72.1.200001"
 
 
+class SectionCode(Enum):
+    PROBLEM = "11450-4"
+    MEDICATION = "10160-0"
+    ALLERGY = "48765-2"
+    NOTE = "51847-2"
+
+
 class ProblemCodes(Enum):
     CONDITION = "64572001"
     SYMPTOM = "418799008"
@@ -258,6 +265,31 @@ def _extract_data(self) -> None:
         self.allergy_list: List[AllergyConcept] = self._extract_allergies()
         self.note: str = self._extract_note()
 
+    def _find_section_by_code(self, section_code: str) -> Optional[Section]:
+        """
+        Finds a section in the clinical document by its code value.
+
+        Args:
+            section_code (str): The code of the section to find.
+
+        Returns:
+            Optional[Section]: The section with the specified code, or None if not found.
+        """
+        components = self.clinical_document.component.structuredBody.component
+
+        if not isinstance(components, list):
+            components = [components]
+
+        for component in components:
+            code = component.section.code.code
+
+            if code is None:
+                continue
+            if code == section_code:
+                return component.section
+        log.warning(f"unable to find section with code {section_code}")
+        return None
+
     def _find_section_by_template_id(self, section_id: str) -> Optional[Section]:
         """
         Finds a section in the clinical document by its template ID.
@@ -292,16 +324,24 @@ def _find_section_by_template_id(self, section_id: str) -> Optional[Section]:
         return None
 
     def _find_problems_section(self) -> Optional[Section]:
-        return self._find_section_by_template_id(SectionId.PROBLEM.value)
+        return self._find_section_by_template_id(
+            SectionId.PROBLEM.value
+        ) or self._find_section_by_code(SectionCode.PROBLEM.value)
 
     def _find_medications_section(self) -> Optional[Section]:
-        return self._find_section_by_template_id(SectionId.MEDICATION.value)
+        return self._find_section_by_template_id(
+            SectionId.MEDICATION.value
+        ) or self._find_section_by_code(SectionCode.MEDICATION.value)
 
     def _find_allergies_section(self) -> Optional[Section]:
-        return self._find_section_by_template_id(SectionId.ALLERGY.value)
+        return self._find_section_by_template_id(
+            SectionId.ALLERGY.value
+        ) or self._find_section_by_code(SectionCode.ALLERGY.value)
 
     def _find_notes_section(self) -> Optional[Section]:
-        return self._find_section_by_template_id(SectionId.NOTE.value)
+        return self._find_section_by_template_id(
+            SectionId.NOTE.value
+        ) or self._find_section_by_code(SectionCode.NOTE.value)
 
     def _extract_problems(self) -> List[ProblemConcept]:
         """

tests/conftest.py

@@ -356,6 +356,13 @@ def cda_annotator():
     return CdaAnnotator.from_xml(test_cda)
 
 
+@pytest.fixture
+def cda_annotator_code():
+    with open("./tests/data/test_cda_without_template_id.xml", "r") as file:
+        test_cda_without_template_id = file.read()
+    return CdaAnnotator.from_xml(test_cda_without_template_id)
+
+
 @pytest.fixture
 def cdsservices():
     return CDSServices()