From 8b460fc40637f9a4c0e09e3b6713aed71e6033c8 Mon Sep 17 00:00:00 2001
From: Vadim Laletin <rusmiligamer@gmail.com>
Date: Wed, 28 Aug 2024 19:06:01 +0200
Subject: [PATCH] Add extensions for primitive types #6

---
 fhir_py_types/__init__.py |  5 +++++
 fhir_py_types/ast.py      | 29 ++++++++++++++++++++++-------
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/fhir_py_types/__init__.py b/fhir_py_types/__init__.py
index 35d3492..98e292d 100644
--- a/fhir_py_types/__init__.py
+++ b/fhir_py_types/__init__.py
@@ -46,3 +46,8 @@ class StructureDefinition:
 
 def is_polymorphic(definition: StructureDefinition) -> bool:
     return len(definition.type) > 1
+
+
+def is_primitive_type(property_type: StructurePropertyType) -> bool:
+    # All primitive types starts with lowercased letters
+    return property_type.code[0].islower()
diff --git a/fhir_py_types/ast.py b/fhir_py_types/ast.py
index 52262ee..78b87df 100644
--- a/fhir_py_types/ast.py
+++ b/fhir_py_types/ast.py
@@ -13,6 +13,7 @@
     StructureDefinitionKind,
     StructurePropertyType,
     is_polymorphic,
+    is_primitive_type,
 )
 
 logger = logging.getLogger(__name__)
@@ -111,10 +112,24 @@ def remap_type(
 def zip_identifier_type(
     definition: StructureDefinition, identifier: str
 ) -> Iterable[tuple[str, StructurePropertyType]]:
-    return (
-        (format_identifier(definition, identifier, t), t)
-        for t in [remap_type(definition, t) for t in definition.type]
-    )
+    result = []
+
+    for t in [remap_type(definition, t) for t in definition.type]:
+        result.append((format_identifier(definition, identifier, t), t))
+        if is_primitive_type(t):
+            result.append(
+                (
+                    f"_{format_identifier(definition, identifier, t)}",
+                    StructurePropertyType(
+                        code="Element",
+                        target_profile=[],
+                        required=False,
+                        isarray=definition.type[0].isarray,
+                    ),
+                )
+            )
+
+    return result
 
 
 def make_assignment_statement(
@@ -133,7 +148,7 @@ def make_assignment_statement(
 
 
 def type_annotate(
-    defintion: StructureDefinition,
+    definition: StructureDefinition,
     identifier: str,
     form: Literal[AnnotationForm.Property, AnnotationForm.TypeAlias],
 ) -> Iterable[ast.stmt]:
@@ -145,9 +160,9 @@ def type_annotate(
                 form,
                 default=make_default_initializer(identifier_, type_),
             ),
-            ast.Expr(value=ast.Constant(defintion.docstring)),
+            ast.Expr(value=ast.Constant(definition.docstring)),
         ]
-        for (identifier_, type_) in zip_identifier_type(defintion, identifier)
+        for (identifier_, type_) in list(zip_identifier_type(definition, identifier))
     )