Merge pull request #5 from candleindark/union-translation

Augment implementation of `SlotGenerator._union_schema()`

src/pydantic2linkml/gen_linkml.py

@@ -1084,13 +1084,54 @@ def _union_schema(self, schema: core_schema.UnionSchema) -> None:
 
         :param schema: The schema representing the union restriction
         """
-        # TODO: the current implementation is just an annotation
-        #   A usable implementation is yet to be decided. Useful information
+        # TODO: the current implementation doesn't address all cases of `Union` partly
+        #   due to limitation of LinkML. Useful information
         #   can be found at, https://github.com/orgs/linkml/discussions/2154
-        self._attach_note(
-            "Warning: The translation is incomplete. Union types are yet to be "
-            "supported."
-        )
+
+        def get_model_slot_expression(
+            schema_: core_schema.CoreSchema,
+        ) -> AnonymousSlotExpression:
+            return AnonymousSlotExpression(
+                range=schema_["cls"].__name__,
+            )
+
+        # A map of supported type choices to the functions for generating the
+        # corresponding slot expression
+        supported_type_choices: dict[
+            str, Callable[[core_schema.CoreSchema], AnonymousSlotExpression]
+        ] = {"model": get_model_slot_expression}
+
+        choices = schema["choices"]
+
+        choice_slot_expressions = []
+        for c in choices:
+            # Exits early if a choice is a tuple
+            if isinstance(c, tuple):
+                self._attach_note(
+                    f"Warning: The translation is incomplete. The union core schema "
+                    f"contains a tuple as a choice. Tuples as choices are yet to be "
+                    f"supported. (core schema: {schema})."
+                )
+                return
+
+            # Exits early if a choice is of unsupported type
+            c_type = c["type"]
+            if c_type not in supported_type_choices:
+                self._attach_note(
+                    f"Warning: The translation is incomplete. The union core schema "
+                    f"contains a choice of type {c_type}. The choice type is yet to be "
+                    f"supported. (core schema: {schema})."
+                )
+                return
+
+            choice_slot_expressions.append(supported_type_choices[c_type](c))
+
+        self._slot.any_of = choice_slot_expressions
+
+        # This is needed because of the monotonicity nature of constraints
+        #   in LinkML. For more information,
+        #   see https://linkml.io/linkml/schemas/advanced.html#unions-as-ranges
+        self._slot.range = any_class_def.name
 
     def _tagged_union_schema(self, schema: core_schema.TaggedUnionSchema) -> None:
         """

tests/test_gen_linkml.py

@@ -7,6 +7,7 @@
 from uuid import UUID
 
 import pytest
+from linkml_runtime.linkml_model import SlotDefinition
 from linkml_runtime.linkml_model.meta import AnonymousSlotExpression
 from pydantic import (
     UUID3,
@@ -24,7 +25,7 @@
     conlist,
 )
 
-from pydantic2linkml.gen_linkml import LinkmlGenerator, SlotGenerator
+from pydantic2linkml.gen_linkml import LinkmlGenerator, SlotGenerator, any_class_def
 from pydantic2linkml.tools import (
     fetch_defs,
     get_all_modules,
@@ -86,6 +87,16 @@ def verify_str_lst(
         assert in_no_string(substr, str_lst)
 
 
+def translate_field_to_slot(model: type[BaseModel], fn: str) -> SlotDefinition:
+    """
+    Translate a field of a Pydantic model to a LinkML slot definition
+
+    :param model: The Pydantic model
+    :param fn: The field name of the field to be translated
+    """
+    return SlotGenerator(get_field_schema(model, fn)).generate()
+
+
 @pytest.fixture
 def models_and_enums(request) -> tuple[set[type[BaseModel]], set[type[Enum]]]:
     """
@@ -852,13 +863,55 @@ class Foo(BaseModel):
         assert slot.range == "integer"
 
     def test_union_schema(self):
-        class Foo(BaseModel):
+        class Bar1(BaseModel):
+            y: int
+
+        class Bar2(BaseModel):
+            z: str
+
+        class Foo0(BaseModel):
             x: Union[int, str]
 
-        field_schema = get_field_schema(Foo, "x")
+        # === A case, customized, of type choices expressed as tuples ===
+        field_schema = get_field_schema(Foo0, "x")
+        field_schema.schema["choices"] = [
+            (c, "label")
+            for c in field_schema.schema["choices"]
+            if not isinstance(c, tuple)
+        ]
         slot = SlotGenerator(field_schema).generate()
 
-        assert in_exactly_one_string("Union types are yet to be supported", slot.notes)
+        assert slot.range is None
+        assert in_exactly_one_string(
+            "The union core schema contains a tuple as a choice. "
+            "Tuples as choices are yet to be supported.",
+            slot.notes,
+        )
+
+        # === Union of base types and models ===
+        class Foo1(BaseModel):
+            x: Union[int, Bar1, str]
+
+        slot = translate_field_to_slot(Foo1, "x")
+
+        assert slot.range is None
+        assert in_exactly_one_string(
+            "The union core schema contains a choice of type int. "
+            "The choice type is yet to be supported.",
+            slot.notes,
+        )
+
+        # === Unions of two models ===
+        class Foo2(BaseModel):
+            x: Union[Bar1, Bar2]
+
+        slot = translate_field_to_slot(Foo2, "x")
+
+        assert slot.range == any_class_def.name
+        assert slot.any_of == [
+            AnonymousSlotExpression(range="Bar1"),
+            AnonymousSlotExpression(range="Bar2"),
+        ]
 
     def test_tagged_union_schema(self):
         class Cat(BaseModel):