From 5eaba4eaea9670839158721d19abe454f0ed9000 Mon Sep 17 00:00:00 2001
From: Pip Liggins <philippa.liggins@dtc.ox.ac.uk>
Date: Mon, 5 Aug 2024 12:23:50 +0100
Subject: [PATCH] Add types to util functions

---
 fhirflat/flat2fhir.py | 11 ++++++++---
 fhirflat/util.py      | 17 +++++++++++++----
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/fhirflat/flat2fhir.py b/fhirflat/flat2fhir.py
index b201bf1..20ec377 100644
--- a/fhirflat/flat2fhir.py
+++ b/fhirflat/flat2fhir.py
@@ -8,7 +8,12 @@
 from fhir.resources.quantity import Quantity
 from pydantic.v1.error_wrappers import ValidationError
 
-from .util import find_data_class, get_fhirtype, get_local_extension_type, group_keys
+from .util import (
+    find_data_class,
+    get_fhirtype,
+    get_local_extension_type,
+    group_keys,
+)
 
 
 def create_codeable_concept(
@@ -16,7 +21,7 @@ def create_codeable_concept(
 ) -> dict[str, list[str]]:
     """Re-creates a codeableConcept structure from the FHIRflat representation."""
 
-    # for reading in from ingestion pipeline
+    # for creating backbone elements
     if name + ".code" in old_dict and name + ".system" in old_dict:
         raw_codes: str | float | list[str | None] = old_dict.get(name + ".code")
         if raw_codes is not None and not isinstance(raw_codes, list):
@@ -24,7 +29,7 @@ def create_codeable_concept(
                 raw_codes if isinstance(raw_codes, str) else str(int(raw_codes))
             )
             codes = [old_dict[name + ".system"] + "|" + formatted_code]
-        elif raw_codes is None:
+        elif not raw_codes:
             codes = raw_codes
         else:
             formatted_codes = [
diff --git a/fhirflat/util.py b/fhirflat/util.py
index f6d0be8..9e21e56 100644
--- a/fhirflat/util.py
+++ b/fhirflat/util.py
@@ -1,16 +1,23 @@
 # Utility functions for FHIRflat
+from __future__ import annotations
+
 import datetime
 import importlib
 import re
 from collections.abc import KeysView
 from itertools import groupby
+from typing import TYPE_CHECKING
 
 import fhir.resources
 import numpy as np
+import pandas as pd
 
 import fhirflat
 from fhirflat.resources import extensions
 
+if TYPE_CHECKING:
+    from .resources.base import FHIRFlatBase
+
 
 def group_keys(data_keys: list[str] | KeysView) -> dict[str, list[str]]:
     """
@@ -79,7 +86,9 @@ def get_local_resource(t: str, case_insensitive: bool = False):
                 return getattr(fhirflat, a)
 
 
-def find_data_class(data_class, k):
+def find_data_class(
+    data_class: FHIRFlatBase | list[FHIRFlatBase], k: str
+) -> FHIRFlatBase:
     """
     Finds the type class for item k within the data class.
 
@@ -116,7 +125,7 @@ def find_data_class(data_class, k):
         return get_fhirtype(base_class)
 
 
-def code_or_codeable_concept(col_name, resource):
+def code_or_codeable_concept(col_name: str, resource: FHIRFlatBase) -> bool:
     search_terms = col_name.split(".")
     fhir_type = find_data_class(resource, search_terms[0])
 
@@ -138,7 +147,7 @@ def code_or_codeable_concept(col_name, resource):
         return code_or_codeable_concept(".".join(search_terms[1:]), fhir_type)
 
 
-def format_flat(flat_df, resource):
+def format_flat(flat_df: pd.DataFrame, resource: FHIRFlatBase) -> pd.DataFrame:
     """
     Performs formatting on dates/lists in FHIRflat resources.
     """
@@ -177,7 +186,7 @@ def format_flat(flat_df, resource):
     return flat_df
 
 
-def condense_codes(row, code_col):
+def condense_codes(row: pd.Series, code_col: str) -> pd.Series:
     raw_codes = row[(code_col + ".code")]
     if isinstance(raw_codes, (str, int, float)) and raw_codes == raw_codes:
         formatted_code = (