Introduce compile_as_first and compile_as_array helpers

Pipfile

@@ -6,6 +6,7 @@ name = "pypi"
 [packages]
 antlr4-python3-runtime = "~=4.10"
 python-dateutil = "~=2.8"
+fhirpy-types-r4b = "==0.1.1"
 
 [dev-packages]
 pytest = "==7.1.1"

fhirpathpy/__init__.py

@@ -1,8 +1,10 @@
-from fhirpathpy.engine.invocations.constants import constants
-from fhirpathpy.parser import parse
+from typing import Any, Callable, Optional, Type
+
 from fhirpathpy.engine import do_eval
-from fhirpathpy.engine.util import arraify, get_data, set_paths, process_user_invocation_table
+from fhirpathpy.engine.invocations.constants import constants
 from fhirpathpy.engine.nodes import FP_Type, ResourceNode
+from fhirpathpy.engine.util import arraify, get_data, process_user_invocation_table, set_paths
+from fhirpathpy.parser import parse
 
 __title__ = "fhirpathpy"
 __version__ = "2.1.0"
@@ -124,3 +126,61 @@ def compile(path, model=None, options=None):
     For example, you could pass in the result of require("fhirpath/fhir-context/r4")
     """
     return set_paths(apply_parsed_path, parsedPath=parse(path), model=model, options=options)
+
+
+type ResourceType = Any
+type ContextType = Optional[dict]
+
+
+def compile_as_array(
+    expression: str, input_type: Type, output_type: Type
+) -> Callable[[Any], Any]:
+    path_fn = compile(expression)
+
+    def fn(resource: input_type, context: ContextType = None) -> list[output_type]:
+        return _format_result(
+            path_fn(_validate_and_convert_resource(resource, input_type), context), output_type, False
+        )
+
+    return fn
+
+
+def compile_as_first(
+    expression: str, input_type: Type, output_type: Type
+) -> Callable[[Any], Any]:
+    path_fn = compile(expression)
+
+    def fn(resource: input_type, context: ContextType = None) -> output_type:
+        return _format_result(
+            path_fn(_validate_and_convert_resource(resource, input_type), context), output_type, True
+        )
+
+    return fn
+
+
+def _validate_and_convert_resource(resource: ResourceType, input_type: Type) -> dict:
+    if isinstance(resource, input_type):
+        if isinstance(resource, dict):
+            return resource
+        elif hasattr(resource, "model_dump"):
+            return resource.model_dump()
+        else:
+            raise Exception(f"Don't know how to work with type {type(resource)}")
+    else:
+        raise Exception(f"Resource type is {type(resource)}, expected {input_type}")
+
+
+def _format_result(result: list, output_type: Type, is_first=False) -> Any:
+    if isinstance(result, list):
+        if is_first:
+            if len(result) > 0:
+                if isinstance(result[0], output_type):
+                    return result[0]
+                else:
+                    raise Exception(f"Unexpected result type {type(result)}, expected {output_type}")
+            else:
+                return None
+        else:
+            return result
+    else:
+        raise Exception(f"Unexpected result type {type(result)}")

tests/test_helpers.py

@@ -0,0 +1,41 @@
+from fhirpathpy import compile_as_first, compile_as_array
+import fhirpy_types_r4b as r4b
+import pytest
+
+PATIENT_DATA = {
+    "resourceType": "Patient",
+    "name": [{"given": ["First", "Middle"], "family": "Last"}],
+}
+
+PATIENT_RESOURCE = r4b.Patient(**PATIENT_DATA)
+
+EXPRESSION = "Patient.name.given"
+
+@pytest.mark.parametrize(
+    ("fn", "resource", "path", "input_type", "output_type", "expected"),
+    [
+        (compile_as_first, PATIENT_DATA, EXPRESSION, dict, str, "First"),
+        (compile_as_first, PATIENT_RESOURCE, EXPRESSION, r4b.Patient, str, "First"),
+        (compile_as_array, PATIENT_DATA, EXPRESSION, dict, list[str], ["First", "Middle"]),
+        (compile_as_array, PATIENT_RESOURCE, EXPRESSION, r4b.Patient, list[str], ["First", "Middle"]),
+    ],
+)
+def compile_as_test(fn, resource, path, input_type, output_type, expected):
+    wrapper_fn = fn(path, input_type, output_type)
+    assert wrapper_fn(resource) == expected
+
+
+@pytest.mark.parametrize(
+    ("fn", "resource", "path", "input_type", "output_type", "expected"),
+    [
+        (compile_as_first, PATIENT_RESOURCE, EXPRESSION, r4b.Observation, str, "Resource type is <class 'fhirpy_types_r4b.Patient'>, expected <class 'fhirpy_types_r4b.Observation'>"),
+        (compile_as_array, PATIENT_RESOURCE, EXPRESSION, r4b.Observation, list[str], "Resource type is <class 'fhirpy_types_r4b.Patient'>, expected <class 'fhirpy_types_r4b.Observation'>"),
+    ],
+)
+def exception_compile_as_test(fn, resource, path, input_type, output_type, expected):
+    wrapper_fn = fn(path, input_type, output_type)
+    try:
+        wrapper_fn(resource)
+        assert False, "Expected exception not raised"
+    except Exception as e:
+        assert str(e) == expected