Add GitHub support to check_branch.py script

scripts/check_branch.py

@@ -1,3 +1,6 @@
+from dataclasses import (
+    dataclass,
+)
 import os
 import sys
 from typing import (
@@ -37,29 +40,57 @@ def __init__(self,
                          f'only {allowed}personal deployments.')
 
 
-def check_branch(branch: Optional[str], deployment: str) -> None:
-    deployment = config.Deployment(deployment)
-    if deployment.is_shared:
-        deployments = config.shared_deployments_for_branch(branch)
-        if deployments is None or deployment not in deployments:
-            raise BranchDeploymentMismatch(branch, deployment, deployments)
+@dataclass(kw_only=True)
+class BranchContext:
+    name: Optional[str]
+    base_ref: Optional[str]
+    head_ref: Optional[str]
+    platform: str
 
 
-def gitlab_branch() -> Optional[str]:
+def branch_context() -> BranchContext:
     """
-    Return the current branch if we're on GitLab, else `None`
+    Determines the current branch context by checking multiple environment
+    variables across different CI platforms and local environments.
     """
-    # Gitlab checks out a specific commit which results in a detached HEAD
-    # (no active branch). Extract the branch name from the runner environment.
-    return os.environ.get('CI_COMMIT_REF_NAME')
+    gh_base_branch = os.environ.get('GITHUB_BASE_REF')
+    gh_feature_branch = os.environ.get('GITHUB_HEAD_REF')
+    gitlab_branch = os.environ.get('CI_COMMIT_REF_NAME')
+    if gh_base_branch is not None or gh_feature_branch is not None:
+        context = BranchContext(
+            name=gh_feature_branch,
+            base_ref=gh_base_branch,
+            head_ref=gh_feature_branch,
+            platform='github'
+        )
+    elif gitlab_branch is not None:
+        # Gitlab checks out a specific commit which results in a detached HEAD
+        # (no active branch). So the branch name from the runner environment is
+        # used.
+        context = BranchContext(
+            name=gitlab_branch,
+            base_ref=None,
+            head_ref=None,
+            platform='gitlab'
+        )
+    else:
+        repo = git.Repo(config.project_root)
+        local_branch = None if repo.head.is_detached else repo.active_branch.name
+        context = BranchContext(
+            name=local_branch,
+            base_ref=None,
+            head_ref=None,
+            platform='local'
+        )
+    return context
 
 
-def local_branch() -> Optional[str]:
-    """
-    Return `None` if detached head, else the current branch
-    """
-    repo = git.Repo(config.project_root)
-    return None if repo.head.is_detached else repo.active_branch.name
+def check_branch(branch: Optional[str], deployment: str) -> None:
+    deployment = config.Deployment(deployment)
+    if deployment.is_shared:
+        deployments = config.shared_deployments_for_branch(branch)
+        if deployments is None or deployment not in deployments:
+            raise BranchDeploymentMismatch(branch, deployment, deployments)
 
 
 def main(argv):
@@ -71,15 +102,30 @@ def main(argv):
                         help='Print the deployment matching the current branch or exit '
                              'with non-zero status code if no such deployment exists.')
     args = parser.parse_args(argv)
-    branch = gitlab_branch() or local_branch()
+    branch_in_context = branch_context()
+    if branch_in_context.platform == 'github' and branch_in_context.base_ref:
+        branch = branch_in_context.base_ref
+    else:
+        branch = branch_in_context.name
+    deployment = default_deployment(branch)
     if args.print:
-        deployment = default_deployment(branch)
         if deployment is None:
             sys.exit(1)
         else:
             print(deployment)
     else:
         check_branch(branch, config.deployment_stage)
+        allowed_deployments = config.shared_deployments_for_branch(branch)
+        if deployment == 'prod':
+            if branch != 'prod' and branch_in_context.base_ref != 'prod':
+                raise BranchDeploymentMismatch(branch,
+                                               config.Deployment(deployment),
+                                               allowed_deployments)
+        elif deployment == 'dev':
+            if branch != 'develop' and branch_in_context.base_ref != 'develop':
+                raise BranchDeploymentMismatch(branch,
+                                               config.Deployment(deployment),
+                                               allowed_deployments)
 
 
 if __name__ == '__main__':