feat: add config-driven release_level utility

src/mozilla_taskgraph/__init__.py

@@ -26,6 +26,7 @@ def register(graph_config):
             "worker_types",
         ]
     )
+
     validate_graph_config(graph_config._config)
 
 

src/mozilla_taskgraph/config.py

@@ -1,5 +1,5 @@
 from textwrap import dedent
-from typing import Optional
+from typing import Optional, Union
 
 from taskgraph import config as tg
 from taskgraph.util.schema import Schema
@@ -26,6 +26,12 @@ class MozillaGraphConfigSchema(tg.graph_config_schema):
         # string to use for release tasks.
         # Defaults to ``mozilla_taskgraph.version:default_parser``.
         version_parser: Optional[str] = None
+        # Mapping of project to the branches that should be considered
+        # "production" releases. A value of ``True`` means all branches of the
+        # project are release branches, while a list restricts releases to the
+        # named branches. Consumed by
+        # ``mozilla_taskgraph.util.attributes:release_level``.
+        release_branches: Optional[dict[str, Union[bool, list[str]]]] = None
 
 else:
     # Legacy voluptuous-based graph_config_schema (e.g. gecko_taskgraph override).
@@ -48,6 +54,15 @@ class MozillaGraphConfigSchema(tg.graph_config_schema):
                     Defaults to ``mozilla_taskgraph.version:default_parser``.
                     """.lstrip()),
             ): str,
+            Vol_Optional(
+                "release-branches",
+                description=dedent("""
+                    Mapping of project to the branches that should be considered
+                    "production" releases. A value of ``True`` means all branches
+                    of the project are release branches, while a list restricts
+                    releases to the named branches.
+                    """.lstrip()),
+            ): {str: object},
         }
     )
 

src/mozilla_taskgraph/util/attributes.py

@@ -0,0 +1,29 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+import re
+
+
+def release_level(release_branches, params):
+    """Whether this is a production release or not.
+
+    ``release_branches`` is the graph config's ``release-branches`` mapping of
+    project to the branches considered "production" for it. A value of ``True``
+    for a project means every branch of that project is a release branch (the
+    model used by Mercurial based projects), while a list restricts releases to
+    the named branches.
+
+    :return str: One of "production" or "staging".
+    """
+
+    if params["level"] == "3":
+        branches = (release_branches or {}).get(params["project"])
+        if branches is True or (
+            branches
+            and (m := re.match(r"refs/heads/(\S+)$", params["head_ref"]))
+            and m.group(1) in branches
+        ):
+            return "production"
+
+    return "staging"

test/util/test_attributes.py

@@ -0,0 +1,56 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+import pytest
+
+from mozilla_taskgraph.util.attributes import release_level
+
+FIREFOX_BRANCHES = ["main", "beta", "release", "esr140"]
+RELEASE_BRANCHES = {
+    "firefox": FIREFOX_BRANCHES,
+    "mozilla-central": True,
+}
+
+
+@pytest.mark.parametrize(
+    "release_branches,params,expected",
+    (
+        # Not level 3 -> always staging, regardless of branch.
+        (RELEASE_BRANCHES, {"level": "1", "project": "mozilla-central"}, "staging"),
+        (
+            RELEASE_BRANCHES,
+            {"level": "1", "project": "firefox", "head_ref": "refs/heads/beta"},
+            "staging",
+        ),
+        # No `release-branches` mapping at all -> staging.
+        (
+            None,
+            {"level": "3", "project": "firefox", "head_ref": "refs/heads/beta"},
+            "staging",
+        ),
+        # Project not listed in the mapping -> staging (e.g. autoland).
+        (RELEASE_BRANCHES, {"level": "3", "project": "autoland"}, "staging"),
+        # Whole project is a release project (Mercurial model).
+        (RELEASE_BRANCHES, {"level": "3", "project": "mozilla-central"}, "production"),
+        # Git monorepo model: only listed branches are production.
+        (
+            RELEASE_BRANCHES,
+            {"level": "3", "project": "firefox", "head_ref": "refs/heads/beta"},
+            "production",
+        ),
+        (
+            RELEASE_BRANCHES,
+            {"level": "3", "project": "firefox", "head_ref": "refs/heads/test"},
+            "staging",
+        ),
+        # Only refs/heads/* match, not tags.
+        (
+            RELEASE_BRANCHES,
+            {"level": "3", "project": "firefox", "head_ref": "refs/tags/beta"},
+            "staging",
+        ),
+    ),
+)
+def test_release_level(release_branches, params, expected):
+    assert release_level(release_branches, params) == expected