pre-commit tidying

dotimplement · May 12, 2024 · 88e02ff · 88e02ff
1 parent 27e9bff
commit 88e02ff
Show file tree

Hide file tree

Showing 18 changed files with 1,499 additions and 520 deletions.
diff --git a/README.md b/README.md
@@ -7,4 +7,3 @@ Simplify prototyping and testing LLM applications in healthcare context.
 ## Quickstart
 
 You can install this package by running:
-
diff --git a/healthchain/data_generator/base_generators.py b/healthchain/data_generator/base_generators.py
@@ -3,135 +3,159 @@
 import random
 import string
 
-from healthchain.fhir_resources.base_resources import booleanModel, canonicalModel, codeModel, dateModel, dateTimeModel, decimalModel, idModel, instantModel, integerModel, markdownModel, positiveIntModel, stringModel, timeModel, unsignedIntModel, uriModel, urlModel, uuidModel
+from healthchain.fhir_resources.base_resources import (
+    booleanModel,
+    canonicalModel,
+    codeModel,
+    dateModel,
+    dateTimeModel,
+    decimalModel,
+    idModel,
+    instantModel,
+    integerModel,
+    markdownModel,
+    positiveIntModel,
+    stringModel,
+    timeModel,
+    unsignedIntModel,
+    uriModel,
+    urlModel,
+    uuidModel,
+)
 
 from faker import Faker
+
 faker = Faker()
 
 generator_registry = {}
 
+
 def register_generator(cls):
     generator_registry[cls.__name__] = cls
     return cls
 
+
 @register_generator
 class BaseGenerator:
     @staticmethod
     def generate():
         raise NotImplementedError("Each generator must implement a 'generate' method.")
 
+
 @register_generator
 class booleanGenerator(BaseGenerator):
     @staticmethod
     def generate():
-        return booleanModel(random.choice(['true', 'false']))
-    
+        return booleanModel(random.choice(["true", "false"]))
+
 
 @register_generator
 class canonicalGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return canonicalModel(f"https://example/{faker.uri_path()}")
-    
+
 
 @register_generator
 class codeGenerator(BaseGenerator):
     # TODO: Codes can technically have whitespace but here I've left it out for simplicity
     @staticmethod
     def generate():
-        return codeModel(''.join(random.choices(string.ascii_uppercase + string.digits, k=6)))
-
+        return codeModel(
+            "".join(random.choices(string.ascii_uppercase + string.digits, k=6))
+        )
+
 
 @register_generator
 class dateGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return dateModel(faker.date())
-    
+
 
 @register_generator
 class dateTimeGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return dateTimeModel(faker.date_time().isoformat())
-    
+
 
 @register_generator
 class decimalGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return decimalModel(faker.random_number())
-    
+
 
 @register_generator
 class idGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return idModel(faker.uuid4())
-    
+
 
 @register_generator
 class instantGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return instantModel(faker.date_time().isoformat())
-    
+
 
 @register_generator
 class integerGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return integerModel(faker.random_int())
-    
+
 
 @register_generator
 class markdownGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return markdownModel(faker.text())
-    
+
 
 @register_generator
 class positiveIntGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return positiveIntModel(faker.random_int(min=1))
-    
+
 
 @register_generator
 class stringGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return stringModel(faker.word())
-    
+
 
 @register_generator
 class timeGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return timeModel(faker.time())
-    
+
 
 @register_generator
 class unsignedIntGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return unsignedIntModel(faker.random_int(min=0))
-    
+
 
 @register_generator
 class uriGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return uriModel(f"https://example/{faker.uri_path()}")
-    
+
 
 @register_generator
 class urlGenerator(BaseGenerator):
     @staticmethod
     def generate():
         return urlModel(f"https://example/{faker.uri_path()}")
-    
+
 
 @register_generator
 class uuidGenerator(BaseGenerator):

diff --git a/healthchain/data_generator/encounter_generator.py b/healthchain/data_generator/encounter_generator.py
@@ -1,43 +1,69 @@
 from healthchain.fhir_resources.encounter_resources import EncounterModel
 from healthchain.fhir_resources.base_resources import CodingModel, CodeableConceptModel
-from healthchain.data_generator.base_generators import BaseGenerator, generator_registry, register_generator
-
+from healthchain.data_generator.base_generators import (
+    BaseGenerator,
+    generator_registry,
+    register_generator,
+)
+from typing import Optional
 from faker import Faker
 
 faker = Faker()
-from typing import Optional
+
 
 @register_generator
 class ClassGenerator(BaseGenerator):
     @staticmethod
     def generate():
-        patient_class_mapping = {
-        'IMP': 'inpatient',
-        'AMB': 'ambulatory'
-    }
-        patient_class = faker.random_element(elements=('IMP', 'AMB'))
-        return CodeableConceptModel(coding=[CodingModel(
-            system="http://terminology.hl7.org/CodeSystem/v3-ActCode",
-            code=patient_class,
-            display=patient_class_mapping.get(patient_class)
-        )]
-    )
+        patient_class_mapping = {"IMP": "inpatient", "AMB": "ambulatory"}
+        patient_class = faker.random_element(elements=("IMP", "AMB"))
+        return CodeableConceptModel(
+            coding=[
+                CodingModel(
+                    system="http://terminology.hl7.org/CodeSystem/v3-ActCode",
+                    code=patient_class,
+                    display=patient_class_mapping.get(patient_class),
+                )
+            ]
+        )
 
 
 @register_generator
 class EncounterTypeGenerator(BaseGenerator):
-    encounter_type_mapping = {
-        'ADMS': 'admission',
-        'EMER': 'emergency'
-    }
-    encounter_type = faker.random_element(elements=('ADMS', 'EMER'))
     @staticmethod
     def generate():
-        return CodeableConceptModel(coding=[CodingModel(
-            system="http://terminology.hl7.org/CodeSystem/v3-ActCode",
-            code=EncounterTypeGenerator.encounter_type,
-            display=EncounterTypeGenerator.encounter_type_mapping.get(EncounterTypeGenerator.encounter_type)
-        )])
+        encounter_type_mapping = {"ADMS": "admission", "EMER": "emergency"}
+        encounter_type = faker.random_element(elements=("ADMS", "EMER"))
+        return CodeableConceptModel(
+            coding=[
+                CodingModel(
+                    system="http://terminology.hl7.org/CodeSystem/v3-ActCode",
+                    code=encounter_type,
+                    display=encounter_type_mapping.get(encounter_type),
+                )
+            ]
+        )
+
+
+@register_generator
+class EncounterPriorityGenerator(BaseGenerator):
+    @staticmethod
+    def generate():
+        encounter_priority_mapping = {
+            "ALRT": "alert",
+            "CRIT": "critical",
+            "NRM": "normal",
+        }
+        encounter_priority = faker.random_element(elements=("ALRT", "CRIT", "NRM"))
+        return CodeableConceptModel(
+            coding=[
+                CodingModel(
+                    system="http://terminology.hl7.org/CodeSystem/v3-ActCode",
+                    code=encounter_priority,
+                    display=encounter_priority_mapping.get(encounter_priority),
+                )
+            ]
+        )
 
 
 @register_generator
@@ -48,19 +74,24 @@ def generate(patient_reference: Optional[str]):
             patient_reference = "Patient/123"
         return EncounterModel(
             resourceType="Encounter",
-            id=generator_registry.get('idGenerator').generate(),
+            id=generator_registry.get("idGenerator").generate(),
             text={
                 "status": "generated",
-                "div": "<div xmlns=\"http://www.w3.org/1999/xhtml\">Encounter with patient @example</div>"
+                "div": '<div xmlns="http://www.w3.org/1999/xhtml">Encounter with patient @example</div>',
             },
             # TODO: Move the elements to live with the resources
-            status=faker.random_element(elements=('planned', 'in-progress', 'on-hold', 'discharged', 'cancelled')),
-            class_field=[generator_registry.get('ClassGenerator').generate()],
-            type_field=[generator_registry.get('EncounterTypeGenerator').generate()],
-            subject={
-                "reference": patient_reference,
-                "display": patient_reference
-            },
+            status=faker.random_element(
+                elements=(
+                    "planned",
+                    "in-progress",
+                    "on-hold",
+                    "discharged",
+                    "cancelled",
+                )
+            ),
+            class_field=[generator_registry.get("ClassGenerator").generate()],
+            type_field=[generator_registry.get("EncounterTypeGenerator").generate()],
+            subject={"reference": patient_reference, "display": patient_reference},
             participant=[],
             reason=[],
         )
diff --git a/healthchain/data_generator/generator_templates/templates.py b/healthchain/data_generator/generator_templates/templates.py
@@ -1,11 +1,6 @@
-from pydantic import Field, BaseModel
+from pydantic import BaseModel
+
 
 class patient_template_1(BaseModel):
-    name: list = [
-        {
-            "family": "Doe",
-            "given": ["John"],
-            "prefix": ["Mr."]
-        }
-    ]
-    birthDate: str = "1999-01-01"
+    name: list = [{"family": "Doe", "given": ["John"], "prefix": ["Mr."]}]
+    birthDate: str = "1999-01-01"