From 82263258e94de4eb3c0a767bdabe8ff63f63b206 Mon Sep 17 00:00:00 2001
From: Scott K Logan <logans@cottsay.net>
Date: Fri, 31 May 2024 15:29:06 -0500
Subject: [PATCH] Add get_package_condition_context function (#182)

This function captures package condition context variables for use when
resolving conditions in package manifests. This function is currently
implemented in ros_buildfarm, but is useful in other contexts as well.
---
 src/rosdistro/__init__.py | 23 +++++++++++++++++++++++
 test/files/index_v4.yaml  |  1 +
 test/test_index.py        | 13 +++++++++++++
 3 files changed, 37 insertions(+)

diff --git a/src/rosdistro/__init__.py b/src/rosdistro/__init__.py
index 927dc1542..a7044c8f0 100644
--- a/src/rosdistro/__init__.py
+++ b/src/rosdistro/__init__.py
@@ -179,6 +179,29 @@ def get_distribution_cache(index, dist_name):
     return DistributionCache(dist_name, data)
 
 
+def get_package_condition_context(index, dist_name):
+    if dist_name not in index.distributions.keys():
+        raise RuntimeError("Unknown distribution: '{0}'. Valid distribution names are: {1}".format(dist_name, ', '.join(sorted(index.distributions.keys()))))
+
+    condition_context = {
+        'ROS_DISTRO': dist_name,
+    }
+
+    dist = index.distributions[dist_name]
+    python_version = dist.get('python_version')
+    if python_version:
+        condition_context['ROS_PYTHON_VERSION'] = str(python_version)
+
+    ros_version = {
+        'ros1': '1',
+        'ros2': '2',
+    }.get(dist.get('distribution_type'))
+    if ros_version:
+        condition_context['ROS_VERSION'] = ros_version
+
+    return condition_context
+
+
 # internal
 
 def _get_dist_file_data(index, dist_name, type_):
diff --git a/test/files/index_v4.yaml b/test/files/index_v4.yaml
index 8d4a31bbe..8d412eeb7 100644
--- a/test/files/index_v4.yaml
+++ b/test/files/index_v4.yaml
@@ -8,5 +8,6 @@ distributions:
     distribution_cache: foo/release-cache.yaml
     distribution_status: active
     distribution_type: ros1
+    python_version: 3
 type: index
 version: 4
diff --git a/test/test_index.py b/test/test_index.py
index 7a20292cc..7d5bac459 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -4,6 +4,7 @@
 from rosdistro import get_distribution_files
 from rosdistro import get_index
 from rosdistro import get_index_url
+from rosdistro import get_package_condition_context
 
 from . import path_to_url
 
@@ -96,3 +97,15 @@ def test_get_index_from_http_with_query_parameters():
         get_distribution_file(i, 'foo')
     finally:
         proc.terminate()
+
+
+def test_get_condition_context():
+    url = path_to_url(os.path.join(FILES_DIR, 'index_v4.yaml'))
+    i = get_index(url)
+    condition_context = get_package_condition_context(i, 'foo')
+
+    assert condition_context == {
+        'ROS_DISTRO': 'foo',
+        'ROS_PYTHON_VERSION': '3',
+        'ROS_VERSION': '1',
+    }