From 8acdee15354a8e903ee85b48276310a400e52cae Mon Sep 17 00:00:00 2001
From: JB Lovland <jlov@equinor.com>
Date: Tue, 23 Jan 2024 16:57:34 +0100
Subject: [PATCH] Added callable discriminator

---
 boom.py                                       |  2 +
 ex.py                                         |  2 +
 src/fmu/dataio/models/meta/content.py         |  1 +
 src/fmu/dataio/models/meta/discriminator.py   | 45 +++++++++++++++++++
 src/fmu/dataio/models/meta/model.py           | 18 ++++++--
 .../test_schema/test_schema_logic_pydantic.py |  2 +-
 6 files changed, 66 insertions(+), 4 deletions(-)
 create mode 100644 src/fmu/dataio/models/meta/discriminator.py

diff --git a/boom.py b/boom.py
index e8f6acebb..0a5360e7d 100644
--- a/boom.py
+++ b/boom.py
@@ -1,3 +1,5 @@
+# type: ignore
+
 from __future__ import annotations
 
 from contextlib import suppress
diff --git a/ex.py b/ex.py
index b8954cba7..94e98c351 100644
--- a/ex.py
+++ b/ex.py
@@ -1,3 +1,5 @@
+# type: ignore
+
 import sys
 
 from fmu.dataio.models.meta.model import Root
diff --git a/src/fmu/dataio/models/meta/content.py b/src/fmu/dataio/models/meta/content.py
index a9897dc12..662f77576 100644
--- a/src/fmu/dataio/models/meta/content.py
+++ b/src/fmu/dataio/models/meta/content.py
@@ -302,6 +302,7 @@ class WellPicksContent(Content):
         FaultLinesContent,
         FieldOutlineContent,
         FieldRegionContent,
+        FluidContactContent,
         InplaceVolumesContent,
         KPProductContent,
         LiftCurvesContent,
diff --git a/src/fmu/dataio/models/meta/discriminator.py b/src/fmu/dataio/models/meta/discriminator.py
new file mode 100644
index 000000000..e43bb6371
--- /dev/null
+++ b/src/fmu/dataio/models/meta/discriminator.py
@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import ValidationError
+
+
+def fmu_discriminator(
+    value: dict,
+) -> Literal[
+    "FMUAggregation",
+    "FMURealization",
+    "FMUCase",
+]:
+    """
+        Discriminate the type of FMU based on the fields present in the value.
+
+        This function determines the type of an FMU object by checking the presence
+        of specific fields ('aggregation' and 'realization') in the given value.
+        It returns a string literal indicating the determined type of the FMU.
+
+        - If both 'aggregation' and 'realization' fields are present,
+            it raises a ValidationError.
+        - If 'aggregation' is present, it returns 'FMUAggregation'.
+        - If 'realization' is present, it returns 'FMURealization'.
+        - If neither is present, it defaults to 'FMUCase'.
+    """
+
+    if not isinstance(value, dict):
+        raise ValidationError("Input value must be a dictionary.")
+
+    if "aggregation" in value and "realization" in value:
+        raise ValidationError(
+            "Value cannot have both 'aggregation' and 'realization' "
+            "fields. It must exclusively be either an 'FMUAggregation' "
+            "or an 'FMURealization'."
+        )
+
+    if "aggregation" in value:
+        return "FMUAggregation"
+
+    if "realization" in value:
+        return "FMURealization"
+
+    return "FMUCase"
diff --git a/src/fmu/dataio/models/meta/model.py b/src/fmu/dataio/models/meta/model.py
index 0bd374a36..c4bb19b57 100644
--- a/src/fmu/dataio/models/meta/model.py
+++ b/src/fmu/dataio/models/meta/model.py
@@ -4,11 +4,11 @@
 from pathlib import Path
 from typing import Dict, Literal, Optional, Union
 
-from pydantic import BaseModel, Field, NaiveDatetime, RootModel
+from pydantic import BaseModel, Discriminator, Field, NaiveDatetime, RootModel, Tag
 from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import Annotated
 
-from . import content, enums
+from . import content, discriminator, enums
 
 
 class UUID(RootModel[str]):
@@ -396,7 +396,19 @@ class FMUDataClassMeta(ClassMeta):
         alias="class",
         title="Metadata class",
     )
-    fmu: Union[FMUAggregation, FMURealization, FMUCase]
+
+    # The presence of the a feild controlls what kind of
+    # FMUObj it is. The fmu_discriminator will inspects
+    # the obj. and returns a tag that tells pydantic
+    # what model to use.
+    fmu: Annotated[
+        Union[
+            Annotated[FMUAggregation, Tag("FMUAggregation")],
+            Annotated[FMURealization, Tag("FMURealization")],
+            Annotated[FMUCase, Tag("FMUCase")],
+        ],
+        Discriminator(discriminator.fmu_discriminator),
+    ]
     access: SsdlAccess
     data: content.AnyContent
     file: File
diff --git a/tests/test_schema/test_schema_logic_pydantic.py b/tests/test_schema/test_schema_logic_pydantic.py
index b7320a34f..1d628beb1 100644
--- a/tests/test_schema/test_schema_logic_pydantic.py
+++ b/tests/test_schema/test_schema_logic_pydantic.py
@@ -5,7 +5,7 @@
 import jsonschema
 import pytest
 from fmu.dataio._definitions import ALLOWED_CONTENTS
-from fmu.dataio.models.meta2 import dump
+from fmu.dataio.models.meta import dump
 
 # pylint: disable=no-member