PY: type-check fixes for python-us-core example (mypy clean)

Makefile

@@ -171,7 +171,7 @@ test-python-fhirpy-sdk: typecheck prepare-aidbox-runme generate-python-sdk-fhirp
 test-python-us-core-example: typecheck generate-python-us-core-sdk python-us-core-test-setup
 	cd $(PYTHON_US_CORE_EXAMPLE) && \
 	     . venv/bin/activate && \
-	     mypy --config-file mypy.ini fhir_types/
+	     mypy --config-file mypy.ini .
 
 	cd $(PYTHON_US_CORE_EXAMPLE) && \
 	     . venv/bin/activate && \

examples/python-us-core/fhir_types/hl7_fhir_r4_core/profiles/observation_observation_vitalsigns.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import Literal
+from typing import Literal, overload
 
 from fhir_types.hl7_fhir_r4_core.observation import Observation
 from fhir_types.hl7_fhir_r4_core.base import CodeableConcept, Period, Reference
@@ -110,7 +110,11 @@ def set_effective_period(self, value: Period) -> "ObservationVitalsignsProfile":
         setattr(self._resource, "effective_period", value)
         return self
 
-    def get_vscat(self, mode: str | None = None) -> dict | None:
+    @overload
+    def get_vscat(self) -> dict | None: ...
+    @overload
+    def get_vscat(self, mode: Literal["raw"]) -> CodeableConcept | None: ...
+    def get_vscat(self, mode: Literal["raw"] | None = None) -> dict | CodeableConcept | None:
         match = self.__class__._vscat_slice_match
         item = get_array_slice(getattr(self._resource, "category", None), match)
         if mode == "raw":

examples/python-us-core/fhir_types/hl7_fhir_us_core/profiles/extension_uscore_ethnicity_extension.py

@@ -142,23 +142,35 @@ def set_text(self, value: "Extension | dict") -> "UscoreEthnicityExtension":
             push_extension(self._resource, {"url": "text", **value})
         return self
 
-    def get_extension_omb_category(self, mode: str | None = None) -> Coding | None:
+    @overload
+    def get_extension_omb_category(self) -> Coding | None: ...
+    @overload
+    def get_extension_omb_category(self, mode: Literal["raw"]) -> Extension | None: ...
+    def get_extension_omb_category(self, mode: Literal["raw"] | None = None) -> Coding | Extension | None:
         match = self.__class__._omb_category_slice_match
         item = get_array_slice(getattr(self._resource, "extension", None), match)
         if mode == "raw":
             return item
         item_dict = item if isinstance(item, dict) else item.model_dump(by_alias=True, exclude_none=True)
         return unwrap_slice_choice(item_dict, ["url"], "valueCoding")
 
-    def get_extension_detailed(self, mode: str | None = None) -> Coding | None:
+    @overload
+    def get_extension_detailed(self) -> Coding | None: ...
+    @overload
+    def get_extension_detailed(self, mode: Literal["raw"]) -> Extension | None: ...
+    def get_extension_detailed(self, mode: Literal["raw"] | None = None) -> Coding | Extension | None:
         match = self.__class__._detailed_slice_match
         item = get_array_slice(getattr(self._resource, "extension", None), match)
         if mode == "raw":
             return item
         item_dict = item if isinstance(item, dict) else item.model_dump(by_alias=True, exclude_none=True)
         return unwrap_slice_choice(item_dict, ["url"], "valueCoding")
 
-    def get_extension_text(self, mode: str | None = None) -> dict | None:
+    @overload
+    def get_extension_text(self) -> dict | None: ...
+    @overload
+    def get_extension_text(self, mode: Literal["raw"]) -> Extension | None: ...
+    def get_extension_text(self, mode: Literal["raw"] | None = None) -> dict | Extension | None:
         match = self.__class__._text_slice_match
         item = get_array_slice(getattr(self._resource, "extension", None), match)
         if mode == "raw":

examples/python-us-core/fhir_types/hl7_fhir_us_core/profiles/extension_uscore_race_extension.py

@@ -145,23 +145,35 @@ def set_text(self, value: "Extension | dict") -> "UscoreRaceExtension":
             push_extension(self._resource, {"url": "text", **value})
         return self
 
-    def get_extension_omb_category(self, mode: str | None = None) -> Coding | None:
+    @overload
+    def get_extension_omb_category(self) -> Coding | None: ...
+    @overload
+    def get_extension_omb_category(self, mode: Literal["raw"]) -> Extension | None: ...
+    def get_extension_omb_category(self, mode: Literal["raw"] | None = None) -> Coding | Extension | None:
         match = self.__class__._omb_category_slice_match
         item = get_array_slice(getattr(self._resource, "extension", None), match)
         if mode == "raw":
             return item
         item_dict = item if isinstance(item, dict) else item.model_dump(by_alias=True, exclude_none=True)
         return unwrap_slice_choice(item_dict, ["url"], "valueCoding")
 
-    def get_extension_detailed(self, mode: str | None = None) -> Coding | None:
+    @overload
+    def get_extension_detailed(self) -> Coding | None: ...
+    @overload
+    def get_extension_detailed(self, mode: Literal["raw"]) -> Extension | None: ...
+    def get_extension_detailed(self, mode: Literal["raw"] | None = None) -> Coding | Extension | None:
         match = self.__class__._detailed_slice_match
         item = get_array_slice(getattr(self._resource, "extension", None), match)
         if mode == "raw":
             return item
         item_dict = item if isinstance(item, dict) else item.model_dump(by_alias=True, exclude_none=True)
         return unwrap_slice_choice(item_dict, ["url"], "valueCoding")
 
-    def get_extension_text(self, mode: str | None = None) -> dict | None:
+    @overload
+    def get_extension_text(self) -> dict | None: ...
+    @overload
+    def get_extension_text(self, mode: Literal["raw"]) -> Extension | None: ...
+    def get_extension_text(self, mode: Literal["raw"] | None = None) -> dict | Extension | None:
         match = self.__class__._text_slice_match
         item = get_array_slice(getattr(self._resource, "extension", None), match)
         if mode == "raw":

examples/python-us-core/fhir_types/hl7_fhir_us_core/profiles/extension_uscore_tribal_affiliation_extension.py

@@ -115,15 +115,23 @@ def set_is_enrolled(self, value: "Extension | dict") -> "UscoreTribalAffiliation
             push_extension(self._resource, {"url": "isEnrolled", **value})
         return self
 
-    def get_extension_tribal_affiliation(self, mode: str | None = None) -> CodeableConcept | None:
+    @overload
+    def get_extension_tribal_affiliation(self) -> CodeableConcept | None: ...
+    @overload
+    def get_extension_tribal_affiliation(self, mode: Literal["raw"]) -> Extension | None: ...
+    def get_extension_tribal_affiliation(self, mode: Literal["raw"] | None = None) -> CodeableConcept | Extension | None:
         match = self.__class__._tribal_affiliation_slice_match
         item = get_array_slice(getattr(self._resource, "extension", None), match)
         if mode == "raw":
             return item
         item_dict = item if isinstance(item, dict) else item.model_dump(by_alias=True, exclude_none=True)
         return unwrap_slice_choice(item_dict, ["url"], "valueCodeableConcept")
 
-    def get_extension_is_enrolled(self, mode: str | None = None) -> dict | None:
+    @overload
+    def get_extension_is_enrolled(self) -> dict | None: ...
+    @overload
+    def get_extension_is_enrolled(self, mode: Literal["raw"]) -> Extension | None: ...
+    def get_extension_is_enrolled(self, mode: Literal["raw"] | None = None) -> dict | Extension | None:
         match = self.__class__._is_enrolled_slice_match
         item = get_array_slice(getattr(self._resource, "extension", None), match)
         if mode == "raw":

examples/python-us-core/fhir_types/hl7_fhir_us_core/profiles/observation_uscore_blood_pressure_profile.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import Literal
+from typing import Literal, overload
 
 from fhir_types.hl7_fhir_r4_core.observation import Observation
 from fhir_types.hl7_fhir_r4_core.base import (
@@ -315,23 +315,35 @@ def set_value_period(self, value: Period) -> "UscoreBloodPressureProfile":
         setattr(self._resource, "value_period", value)
         return self
 
-    def get_vscat(self, mode: str | None = None) -> dict | None:
+    @overload
+    def get_vscat(self) -> dict | None: ...
+    @overload
+    def get_vscat(self, mode: Literal["raw"]) -> CodeableConcept | None: ...
+    def get_vscat(self, mode: Literal["raw"] | None = None) -> dict | CodeableConcept | None:
         match = self.__class__._vscat_slice_match
         item = get_array_slice(getattr(self._resource, "category", None), match)
         if mode == "raw":
             return item
         item_dict = item if isinstance(item, dict) else item.model_dump(by_alias=True, exclude_none=True)
         return strip_match_keys(item_dict, ["coding"])
 
-    def get_systolic(self, mode: str | None = None) -> Quantity | None:
+    @overload
+    def get_systolic(self) -> Quantity | None: ...
+    @overload
+    def get_systolic(self, mode: Literal["raw"]) -> ObservationComponent | None: ...
+    def get_systolic(self, mode: Literal["raw"] | None = None) -> Quantity | ObservationComponent | None:
         match = self.__class__._systolic_slice_match
         item = get_array_slice(getattr(self._resource, "component", None), match)
         if mode == "raw":
             return item
         item_dict = item if isinstance(item, dict) else item.model_dump(by_alias=True, exclude_none=True)
         return unwrap_slice_choice(item_dict, ["code"], "valueQuantity")
 
-    def get_diastolic(self, mode: str | None = None) -> Quantity | None:
+    @overload
+    def get_diastolic(self) -> Quantity | None: ...
+    @overload
+    def get_diastolic(self, mode: Literal["raw"]) -> ObservationComponent | None: ...
+    def get_diastolic(self, mode: Literal["raw"] | None = None) -> Quantity | ObservationComponent | None:
         match = self.__class__._diastolic_slice_match
         item = get_array_slice(getattr(self._resource, "component", None), match)
         if mode == "raw":

examples/python-us-core/fhir_types/hl7_fhir_us_core/profiles/observation_uscore_body_weight_profile.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import Literal
+from typing import Literal, overload
 
 from fhir_types.hl7_fhir_r4_core.observation import Observation
 from fhir_types.hl7_fhir_r4_core.base import (
@@ -189,7 +189,11 @@ def set_value_period(self, value: Period) -> "UscoreBodyWeightProfile":
         setattr(self._resource, "value_period", value)
         return self
 
-    def get_vscat(self, mode: str | None = None) -> dict | None:
+    @overload
+    def get_vscat(self) -> dict | None: ...
+    @overload
+    def get_vscat(self, mode: Literal["raw"]) -> CodeableConcept | None: ...
+    def get_vscat(self, mode: Literal["raw"] | None = None) -> dict | CodeableConcept | None:
         match = self.__class__._vscat_slice_match
         item = get_array_slice(getattr(self._resource, "category", None), match)
         if mode == "raw":

examples/python-us-core/fhir_types/hl7_fhir_us_core/profiles/observation_uscore_vital_signs_profile.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import Literal
+from typing import Literal, overload
 
 from fhir_types.hl7_fhir_r4_core.observation import Observation
 from fhir_types.hl7_fhir_r4_core.base import (
@@ -299,7 +299,11 @@ def set_value_period(self, value: Period) -> "UscoreVitalSignsProfile":
         setattr(self._resource, "value_period", value)
         return self
 
-    def get_vscat(self, mode: str | None = None) -> dict | None:
+    @overload
+    def get_vscat(self) -> dict | None: ...
+    @overload
+    def get_vscat(self, mode: Literal["raw"]) -> CodeableConcept | None: ...
+    def get_vscat(self, mode: Literal["raw"] | None = None) -> dict | CodeableConcept | None:
         match = self.__class__._vscat_slice_match
         item = get_array_slice(getattr(self._resource, "category", None), match)
         if mode == "raw":

examples/python-us-core/test_profile_bodyweight.py

@@ -113,6 +113,6 @@ def test_get_vscat_raw_includes_discriminator():
         subject=Reference(reference="Patient/example"),
     )
 
-    raw = profile.get_vscat("raw")  # type: ignore[call-arg]
+    raw = profile.get_vscat("raw")
     assert raw is not None
     assert raw.coding is not None

examples/python-us-core/test_profile_bp.py

@@ -129,7 +129,7 @@ def test_set_systolic_get_systolic_get_systolic_raw():
         "code": "mm[Hg]",
     }
 
-    raw = profile.get_systolic("raw")  # type: ignore[call-arg]
+    raw = profile.get_systolic("raw")
     assert raw.value_quantity.value == 120
     assert raw.code.coding[0].code == "8480-6"
 
@@ -145,7 +145,7 @@ def test_set_diastolic_get_diastolic_get_diastolic_raw():
         "code": "mm[Hg]",
     }
 
-    raw = profile.get_diastolic("raw")  # type: ignore[call-arg]
+    raw = profile.get_diastolic("raw")
     assert raw.value_quantity.value == 80
     assert raw.code.coding[0].code == "8462-4"
 
@@ -161,7 +161,7 @@ def test_set_systolic_replaces_an_existing_systolic_component():
     profile.set_systolic({"value": 130, "unit": "mmHg"})
     obs = profile.to_resource()
     assert len(obs.component) == 2
-    assert profile.get_systolic("raw").value_quantity.value == 130  # type: ignore[call-arg]
+    assert profile.get_systolic("raw").value_quantity.value == 130
 
 
 def test_set_vscat_adds_category_with_discriminator_values():

examples/python-us-core/test_profile_patient.py

@@ -188,16 +188,14 @@ def test_apply_wraps_an_existing_patient():
 
 
 def test_all_three_methods_produce_equivalent_resources():
-    args = dict(
-        identifier=[Identifier(system="http://hospital.example.org", value="12345")],
-        name=[HumanName(family="Smith", given=["John"])],
-    )
-    from_create = UscorePatientProfile.create(**args).to_resource()
-    from_create_resource = UscorePatientProfile.create_resource(**args)
+    identifier = [Identifier(system="http://hospital.example.org", value="12345")]
+    name = [HumanName(family="Smith", given=["John"])]
+    from_create = UscorePatientProfile.create(identifier=identifier, name=name).to_resource()
+    from_create_resource = UscorePatientProfile.create_resource(identifier=identifier, name=name)
 
     bare = Patient(resource_type="Patient")
     profile = UscorePatientProfile.apply(bare)
-    profile.set_identifier(args["identifier"]).set_name(args["name"])
+    profile.set_identifier(identifier).set_name(name)
     from_apply = profile.to_resource()
 
     for res in (from_create, from_create_resource, from_apply):

examples/python-us-core/test_profile_typed_bundle.py

@@ -53,18 +53,18 @@ def test_get_patient_entry_raw_mode_returns_bundle_entry_instance():
 def test_get_patient_entry_returns_stored_entry_including_resource():
     """Getter returns the stored entry as-is — resource data is preserved."""
     bundle = ExampleTypedBundleProfile.create(type="collection")
-    bundle.set_patient_entry({"resource": smith_patient.model_dump(by_alias=True, exclude_none=True)})
+    bundle.set_patient_entry(BundleEntry(resource=smith_patient))
 
     entry = bundle.get_patient_entry()
     assert entry is not None
-    resource = entry.resource if hasattr(entry, "resource") else (entry or {}).get("resource")
+    resource = entry.resource
     assert resource is not None
 
 
 def test_set_patient_entry_replaces_existing():
     bundle = ExampleTypedBundleProfile.create(type="collection")
-    bundle.set_patient_entry({"resource": smith_patient.model_dump(by_alias=True, exclude_none=True)})
-    bundle.set_patient_entry({"resource": active_patient.model_dump(by_alias=True, exclude_none=True)})
+    bundle.set_patient_entry(BundleEntry(resource=smith_patient))
+    bundle.set_patient_entry(BundleEntry(resource=active_patient))
 
     # Only one patient entry — the second call replaced the first.
     assert len(bundle.to_resource().entry) == 1
@@ -79,7 +79,7 @@ def test_set_organization_entry_accepts_a_list():
     from fhir_types.hl7_fhir_r4_core.organization import Organization
 
     bundle = ExampleTypedBundleProfile.create(type="collection")
-    bundle.set_patient_entry({"resource": smith_patient.model_dump(by_alias=True, exclude_none=True)})
+    bundle.set_patient_entry(BundleEntry(resource=smith_patient))
 
     clinic = Organization(resource_type="Organization", name="Clinic")
     acme = Organization(resource_type="Organization", name="Acme")
@@ -95,7 +95,7 @@ def test_get_organization_entry_returns_list():
     from fhir_types.hl7_fhir_r4_core.organization import Organization
 
     bundle = ExampleTypedBundleProfile.create(type="collection")
-    bundle.set_patient_entry({"resource": smith_patient.model_dump(by_alias=True, exclude_none=True)})
+    bundle.set_patient_entry(BundleEntry(resource=smith_patient))
     org = Organization(resource_type="Organization", name="Clinic")
     bundle.set_organization_entry([BundleEntry(resource=org)])
 

src/api/writer-generator/python/profile-slices.ts

@@ -193,7 +193,14 @@ export const generateSliceGetters = (
             const flatRetType = sliceDef.constrainedChoice
                 ? pyTypeFromIdentifier(sliceDef.constrainedChoice.variantType)
                 : "dict";
-            w.line(`def get_${baseName}(self, mode: str | None = None) -> ${flatRetType} | None:`);
+            const rawRetType = sliceDef.elementTypeName ?? "Any";
+            w.line("@overload");
+            w.line(`def get_${baseName}(self) -> ${flatRetType} | None: ...`);
+            w.line("@overload");
+            w.line(`def get_${baseName}(self, mode: Literal["raw"]) -> ${rawRetType} | None: ...`);
+            w.line(
+                `def get_${baseName}(self, mode: Literal["raw"] | None = None) -> ${flatRetType} | ${rawRetType} | None:`,
+            );
             w.indentBlock(() => {
                 w.line(`match = self.__class__.${staticName}`);
                 if (sliceDef.array) {

src/api/writer-generator/python/profile.ts

@@ -160,8 +160,10 @@ const collectTypeImports = (
         // Constrained-choice slice getters return the variant type (e.g. `valueCoding` -> `Coding`).
         if (s.constrainedChoice)
             addExactTypeImport(typeImports, rootPackageName, baseTypeName, s.constrainedChoice.variantType);
-        if (!s.isTypeDiscriminated) continue;
+        // The element type is referenced by both the setter (`ElementType(**merged)`) and the
+        // raw getter overload, for type-discriminated and plain slices alike.
         if (s.elementTypeId) addExactTypeImport(typeImports, rootPackageName, baseTypeName, s.elementTypeId);
+        if (!s.isTypeDiscriminated) continue;
         if (!s.typeDiscriminatorResource) continue;
         const resourceId = schemas.find(
             (schema) => schema.identifier.kind === "resource" && schema.identifier.name === s.typeDiscriminatorResource,
@@ -204,17 +206,24 @@ const emitModuleImports = (
     typeImports: Map<string, Set<string>>,
     extProfileImports: Map<string, ExtensionProfileInfo>,
     helperImports: string[],
+    sliceDefs: SliceDef[],
 ): void => {
     w.line("from __future__ import annotations");
     w.line();
 
     const usesLiteral = [...factoryInfo.params, ...factoryInfo.accessors, ...factoryInfo.sliceAutoFields].some((f) =>
         f.pyType.includes("Literal["),
     );
+    // Non-type-discriminated slice getters emit `@overload`s with a `Literal["raw"]` mode.
+    const hasNonTypedSlice = sliceDefs.some((s) => !s.isTypeDiscriminated);
+    const needsOverload = extensions.length > 0 || hasNonTypedSlice;
+    // `Any` is needed for extension setter inputs (`X | Extension | Any`) and for raw
+    // slice getters whose element type is unknown (`rawRetType = "Any"`).
+    const needsAny = extensions.length > 0 || sliceDefs.some((s) => !s.isTypeDiscriminated && !s.elementTypeName);
     const typingNames: string[] = [];
-    // `Any` is only needed for extension setter inputs (`X | Extension | Any`).
-    if (extensions.length > 0) typingNames.push("Any", "Literal", "overload");
-    else if (usesLiteral) typingNames.push("Literal");
+    if (needsAny) typingNames.push("Any");
+    if (needsOverload || usesLiteral) typingNames.push("Literal");
+    if (needsOverload) typingNames.push("overload");
     if (typingNames.length > 0) {
         w.pyImportFrom("typing", ...[...typingNames].sort());
         w.line();
@@ -422,6 +431,7 @@ const generateProfileModule = (w: Python, tsIndex: TypeSchemaIndex, flatProfile:
         typeImports,
         extProfileImports,
         helperImports,
+        sliceDefs,
     );
 
     w.line(`class ${className}:`);