From 1b81ea4bdbe96d244b36fd6a974cedc52b0666ed Mon Sep 17 00:00:00 2001
From: tmadlener <thomas.madlener@desy.de>
Date: Wed, 28 Aug 2024 16:41:33 +0200
Subject: [PATCH] Make datamodel versions accessible from python

---
 python/podio/base_reader.py | 20 ++++++++++++++++++++
 python/podio/test_Reader.py | 19 +++++++++++++++++++
 tools/podio-dump            | 12 ++++++++----
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/python/podio/base_reader.py b/python/podio/base_reader.py
index 4aa579159..20c8ad6f7 100644
--- a/python/podio/base_reader.py
+++ b/python/podio/base_reader.py
@@ -89,3 +89,23 @@ def get_podio_version(self):
                 write this file
         """
         return self._reader.currentFileVersion()
+
+    def get_datamodel_version(self, edm_name):
+        """Get the datamodel version
+
+        Args:
+            edm_name (str): The package name of the datamodel
+
+        Returns:
+            podio.version.Version: The version of the datamodel if available
+
+        Raises:
+            KeyError: If the datamodel does not have a version stored
+            RuntimeError: If the reader is a legacy reader
+        """
+        if self._is_legacy:
+            raise RuntimeError("Legacy readers do not store any version info")
+        maybe_version = self._reader.currentFileVersion(edm_name)
+        if maybe_version.has_value():
+            return maybe_version.value()
+        raise KeyError(f"No version information available for '{edm_name}'")
diff --git a/python/podio/test_Reader.py b/python/podio/test_Reader.py
index 1208feace..8324deab1 100644
--- a/python/podio/test_Reader.py
+++ b/python/podio/test_Reader.py
@@ -1,6 +1,8 @@
 #!/usr/bin/env python3
 """Unit tests for podio readers"""
 
+from podio.version import build_version
+
 
 class ReaderTestCaseMixin:
     """Common unittests for readers.
@@ -70,6 +72,23 @@ def test_frame_iterator_invalid_category(self):
             i += 1
         self.assertEqual(i, 0)
 
+    def test_available_datamodels(self):
+        """Make sure that the datamodel information can be retrieved"""
+        datamodels = self.reader.datamodel_definitions
+        self.assertEqual(len(datamodels), 2)
+        for model in datamodels:
+            self.assertTrue(model in ("datamodel", "extension_model"))
+
+        self.assertEqual(self.reader.get_datamodel_version("datamodel"), build_version)
+
+    def test_invalid_datamodel_version(self):
+        """Make sure that the necessary exceptions are raised"""
+        with self.assertRaises(KeyError):
+            self.reader.get_datamodel_version("extension_model")
+
+        with self.assertRaises(KeyError):
+            self.reader.get_datamodel_version("non-existant-model")
+
 
 class LegacyReaderTestCaseMixin:
     """Common test cases for the legacy readers python bindings.
diff --git a/tools/podio-dump b/tools/podio-dump
index 9ee9b02c6..42f21c3a6 100755
--- a/tools/podio-dump
+++ b/tools/podio-dump
@@ -25,10 +25,14 @@ def print_general_info(reader, filename):
         f"input file: {filename}{legacy_text}\n"
         f"            (written with podio version: {version_as_str(reader.get_podio_version())})\n"
     )
-    print(
-        "datamodel model definitions stored in this file: "
-        f'{", ".join(reader.datamodel_definitions)}'
-    )
+
+    print("datamodel model definitions stored in this file: ")
+    for edm_name in reader.datamodel_definitions:
+        try:
+            edm_version = reader.get_datamodel_version(edm_name)
+            print(f" - {edm_name} ({version_as_str(edm_version)})")
+        except KeyError:
+            print(f" - {edm_name}")
 
     print()
     print("Frame categories in this file:")